From 0c1f26eec27e8ecbcbee46617fd6bb0fb0961a6b Mon Sep 17 00:00:00 2001 From: salt Date: Tue, 5 Mar 2024 05:56:29 -0800 Subject: [PATCH 01/22] use new docker --- .../config/inference-job.development.env | 1 + .../workflow/comfy_ui/comfy_ui_dependencies.rs | 18 +++++++++++++++++- .../comfy_ui/comfy_ui_inference_command.rs | 6 ++++++ 3 files changed, 24 insertions(+), 1 deletion(-) diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 9e1d5e5807..957a4f5132 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -152,6 +152,7 @@ COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRun # directory to where files should be downloaded to (default to just the comfy directory) COMFY_MOUNTS_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/" COMFY_VIDEO_PROCESSING_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/format_video.py" +COMFY_SETUP_SCRIPT="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/install.py" # ================================= # # # diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs index fee7faac54..4d1d3b16ef 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs @@ -1,4 +1,5 @@ use std::path::PathBuf; +use std::process::Command; use errors::AnyhowResult; use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand; @@ -10,8 +11,23 @@ pub struct ComfyDependencies { impl ComfyDependencies { pub fn setup() -> AnyhowResult { + // run the setup scripts for the comfy dependencies + let inference_command = ComfyInferenceCommand::from_env()?; + // check to see if the flag file exists + if PathBuf::from("comfy_setup_complete").exists() { + println!("Comfy setup already complete"); + } else { + // if the flag file does not exist, run the setup script + // run setup script + let comfy_setup_script = inference_command.clone().comfy_setup_script; + let output = Command::new("python3") + .arg(comfy_setup_script) + .output() + .expect("failed to run comfyui setup script"); + println!("stdout: {}", String::from_utf8_lossy(&output.stdout)); + } Ok(Self { - inference_command: ComfyInferenceCommand::from_env()?, + inference_command: inference_command, dependency_tokens: RequiredModels::init(), }) } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index f5999aef88..805494f956 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs @@ -44,6 +44,8 @@ pub struct ComfyInferenceCommand { // Video processing script pub(crate) processing_script: PathBuf, + pub(crate) comfy_setup_script: PathBuf, + /// Config file to use config_path: Option, @@ -95,6 +97,9 @@ impl ComfyInferenceCommand { } }; + let comfy_setup_script = easyenv::get_env_pathbuf_required( + "COMFY_SETUP_SCRIPT")?; + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( "COMFY_INFERENCE_MAYBE_VENV_COMMAND"); @@ -124,6 +129,7 @@ impl ComfyInferenceCommand { config_path, mounts_directory, processing_script, + comfy_setup_script, maybe_virtual_env_activation_command, maybe_docker_options, maybe_execution_timeout, From f49efe6c3959fb16ea02e5875c122a092a5b989e Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 19:34:29 -0500 Subject: [PATCH 02/22] logging for debug (why is r2d2 breaking?) --- crates/service/job/inference_job/src/main.rs | 10 ++++++++-- .../job/inference_job/src/util/instrumentation.rs | 4 +++- 2 files changed, 11 insertions(+), 3 deletions(-) diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 0b0abc000a..223030abc7 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -139,13 +139,15 @@ async fn main() -> AnyhowResult<()> { "MYSQL_URL", DEFAULT_MYSQL_CONNECTION_STRING); - info!("Connecting to database..."); + info!("Connecting to MySQL database..."); let mysql_pool = MySqlPoolOptions::new() .max_connections(2) .connect(&db_connection_string) .await?; + info!("Connected to MySQL."); + let common_env = CommonEnv::read_from_env()?; let sidecar_max_synthesizer_models = easyenv::get_env_num( @@ -191,8 +193,12 @@ async fn main() -> AnyhowResult<()> { let maybe_keepalive_redis_pool = match easyenv::get_env_string_optional("REDIS_FOR_KEEPALIVE_URL") { - None => None, + None => { + warn!("Redis for job keepalive is DISABLED! This might break some jobs."); + None + }, Some(redis_url) => { + info!("Connecting to Redis for keepalive signals... {}", redis_url); let redis_manager = RedisConnectionManager::new(redis_url)?; let redis_pool = r2d2::Pool::builder().build(redis_manager)?; Some(redis_pool) diff --git a/crates/service/job/inference_job/src/util/instrumentation.rs b/crates/service/job/inference_job/src/util/instrumentation.rs index a3c806ce06..d269a4d696 100644 --- a/crates/service/job/inference_job/src/util/instrumentation.rs +++ b/crates/service/job/inference_job/src/util/instrumentation.rs @@ -1,5 +1,5 @@ use std::time::Duration; -use log::warn; +use log::{info, warn}; use opentelemetry::KeyValue; use opentelemetry::metrics::{Counter, Histogram, Meter, Unit}; use opentelemetry_otlp::WithExportConfig; @@ -19,6 +19,8 @@ pub struct JobInstrumentLabels { pub fn init_otel_metrics_pipeline( job_instrument_labels: JobInstrumentLabels, ) -> Result<(), opentelemetry::metrics::MetricsError> { + info!("Setting up otel metrics pipeline..."); + let provider = opentelemetry_otlp::new_pipeline() .metrics(opentelemetry_sdk::runtime::Tokio) // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset? From bf8c189ab5432b36100de57c11f248636c167b3b Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 19:46:05 -0500 Subject: [PATCH 03/22] optional env vars --- crates/service/job/inference_job/src/main.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index 223030abc7..f418d7f387 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -177,8 +177,8 @@ async fn main() -> AnyhowResult<()> { // We want to turn this off in the on-premises workers since we're not tunneling to the production Redis. let job_progress_reporter : Box; - job_progress_reporter = match easyenv::get_env_string_optional("REDIS_FOR_JOB_PROGRESS") { - None => { + job_progress_reporter = match easyenv::get_env_string_optional("REDIS_FOR_JOB_PROGRESS").as_deref() { + None | Some("") => { warn!("Redis for job progress status reports is DISABLED! Users will not see in-flight details of inference progress."); Box::new(NoOpJobProgressReporterBuilder {}) }, @@ -192,8 +192,8 @@ async fn main() -> AnyhowResult<()> { }; let maybe_keepalive_redis_pool = - match easyenv::get_env_string_optional("REDIS_FOR_KEEPALIVE_URL") { - None => { + match easyenv::get_env_string_optional("REDIS_FOR_KEEPALIVE_URL").as_deref() { + None | Some("") => { warn!("Redis for job keepalive is DISABLED! This might break some jobs."); None }, From 046eddad99511827840647f7c7705993e8d139ea Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 23:06:56 -0500 Subject: [PATCH 04/22] name file based on type --- .../bvh_to_workflow/process_job.rs | 20 +++++++++++++++++-- 1 file changed, 18 insertions(+), 2 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs index 0296b55fa2..7a56aba4a8 100644 --- a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs @@ -69,8 +69,11 @@ pub async fn process_job(args: BvhToWorkflowJobArgs<'_>) -> Result) -> Result) -> Result Option<&'static str> { + match media_file.media_type { + MediaFileType::Bvh => Some("original.bvh"), + MediaFileType::Glb => Some("original.scn.gltf"), // TODO(bt): Is the extension "gltf" for "glb" ?? + MediaFileType::SceneRon => Some("original.scn.ron"), + _ => None, + } +} + #[derive(Deserialize, Default)] struct FileMetadata { pub duration_millis: Option, From 2ef8a4c5552528978f8296da5d478cf97d441103 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 23:19:55 -0500 Subject: [PATCH 05/22] fix logging --- .../job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs index 7a56aba4a8..204ce4de94 100644 --- a/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/bevy_to_workflow/bvh_to_workflow/process_job.rs @@ -123,7 +123,9 @@ pub async fn process_job(args: BvhToWorkflowJobArgs<'_>) -> Result Date: Fri, 1 Mar 2024 02:42:37 -0500 Subject: [PATCH 06/22] more logging for comfyui --- .../workflow_payload.rs | 11 +++ .../create/insert_media_file_from_comfy_ui.rs | 7 ++ .../comfy_ui/comfy_ui_inference_command.rs | 10 +-- .../workflow/comfy_ui/process_job.rs | 68 +++++++++++++++++-- 4 files changed, 85 insertions(+), 11 deletions(-) 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 b55f8b2994..9e8655ef02 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 @@ -12,6 +12,17 @@ pub enum NewValue { Bool(bool), } +impl NewValue { + pub fn to_string(&self) -> String { + match self { + NewValue::String(s) => s.to_string(), + NewValue::Int(s) => s.to_string(), + NewValue::Float(s) => s.to_string(), + NewValue::Bool(s) => s.to_string(), + } + } +} + #[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] pub struct WorkflowArgs { #[serde(rename = "sd")] diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs index 284cd6894e..cb6eb108ff 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_comfy_ui.rs @@ -9,6 +9,7 @@ use enums::by_table::media_files::media_file_origin_product_category::MediaFileO use enums::by_table::media_files::media_file_type::MediaFileType; use errors::AnyhowResult; use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::prompts::PromptToken; use tokens::tokens::users::UserToken; use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; @@ -24,6 +25,8 @@ pub struct InsertArgs<'a> { // TODO: Media duration. //pub duration_millis: u64, + pub maybe_prompt_token: Option<&'a PromptToken>, + pub public_bucket_directory_hash: &'a str, pub maybe_public_bucket_prefix: Option<&'a str>, pub maybe_public_bucket_extension: Option<&'a str>, @@ -85,6 +88,8 @@ SET checksum_sha2 = ?, + maybe_prompt_token = ?, + public_bucket_directory_hash = ?, maybe_public_bucket_prefix = ?, maybe_public_bucket_extension = ?, @@ -116,6 +121,8 @@ SET args.sha256_checksum, + args.maybe_prompt_token.map(|e| e.as_str()), + args.public_bucket_directory_hash, args.maybe_public_bucket_prefix, args.maybe_public_bucket_extension, diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index 805494f956..a4e66177c7 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs @@ -204,10 +204,12 @@ impl ComfyInferenceCommand { info!("stderr will be written to file: {:?}", args.stderr_output_file.as_os_str()); - let stderr_file = File::create(&args.stderr_output_file)?; - let stdout_file = File::create(&args.stdout_output_file)?; - config.stderr = Redirection::File(stderr_file); - config.stdout = Redirection::File(stdout_file); + let _stderr_file = File::create(&args.stderr_output_file)?; + let _stdout_file = File::create(&args.stdout_output_file)?; + + // NB(bt, 2024-03-01): Let's actually emit these lots to the rust logs so we can see what happens. + //config.stderr = Redirection::File(stderr_file); + //config.stdout = Redirection::File(stdout_file); if !env_vars.is_empty() { config.env = Some(env_vars); diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 8ebf49b884..233d01b063 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -14,6 +14,8 @@ use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloud use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::by_table::prompts::prompt_type::PromptType; +use errors::AnyhowResult; use filesys::check_file_exists::check_file_exists; use filesys::file_size::file_size; use filesys::path_to_string::path_to_string; @@ -27,7 +29,9 @@ use mysql_queries::queries::generic_inference::job::list_available_generic_infer use mysql_queries::queries::media_files::create::insert_media_file_from_comfy_ui::{insert_media_file_from_comfy_ui, InsertArgs}; use mysql_queries::queries::media_files::get_media_file::get_media_file; use mysql_queries::queries::model_weights::get::get_weight::get_weight_by_token; +use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; use thumbnail_generator::task_client::thumbnail_task::ThumbnailTaskBuilder; +use tokens::tokens::prompts::PromptToken; use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity}; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; @@ -200,15 +204,33 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result return Err(ProcessSingleJobError::Other(anyhow!("ComfyUi args not found"))), }; + let mut maybe_positive_prompt = None; + let mut maybe_negative_prompt = None; + // Apply modifications if they exist - if let Some(modifications) = comfy_args.maybe_json_modifications.clone() { + if let Some(modifications) = comfy_args.maybe_json_modifications.as_ref() { + info!("Prompt modifications: #{:?}", modifications); + + // Save prompts for later (if they exist) + maybe_positive_prompt = modifications.get("$.510.inputs.Text") + .map(|v| v.to_string()); + + maybe_negative_prompt = modifications.get("$.8.inputs.text") + .map(|v| v.to_string()); + // Load prompt.json let prompt_file = File::open(&workflow_path).unwrap(); let mut prompt_json: Value = serde_json::from_reader(prompt_file).unwrap(); + // Modify json - for (path, new_value) in modifications { - prompt_json = replace_json_value(prompt_json, &path, new_value).map_err(|e| ProcessSingleJobError::Other(e.into()))?; + for (path, new_value) in modifications + .iter() + .map(|(k, v)| (k.as_str(), v)) + { + prompt_json = replace_json_value(prompt_json, path, new_value) + .map_err(|e| ProcessSingleJobError::Other(anyhow!("error replacing prompt json: {:?}", e)))?; } + // Save prompt.json let prompt_file = File::create(&workflow_path).unwrap(); serde_json::to_writer(prompt_file, &prompt_json).unwrap(); @@ -409,6 +431,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result {} + Err(err) => { + error!("No prompt result token? something has failed: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("No prompt result token? something has failed."))); + } + } + } + info!("ComfyUI Done."); job_progress_reporter.log_status("done") @@ -581,7 +635,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result anyhow::Result { +fn replace_json_value(json: Value, path: &str, new_value: &NewValue) -> AnyhowResult { // First, attempt to read the value at the specified path let results = jsonpath_lib::select(&json, path).map_err(|err| { anyhow!("Invalid jsonpath '{}': {:?}", path, err) @@ -595,7 +649,7 @@ fn replace_json_value(json: Value, path: &str, new_value: NewValue) -> anyhow::R // If the path exists, proceed with the replacement // Assuming replace_with returns a Result, handle it appropriately jsonpath_lib::replace_with(json, path, &mut |_| { - match &new_value { + match new_value { NewValue::String(s) => Some(Value::String(s.clone())), NewValue::Float(f) => serde_json::Number::from_f64(*f as f64) .map(Value::Number) // Convert Option to Some(Value::Number) if Some, else None From 8f3453eb8df2750643a3646af381c767324e8061 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 02:43:37 -0500 Subject: [PATCH 07/22] codegen mysql --- ...1ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd.json | 12 ++++++++++++ ...6f881ebbfadc292939745e38754991650651edbe7acb.json | 12 ------------ 2 files changed, 12 insertions(+), 12 deletions(-) create mode 100644 .sqlx/query-ca7db50451b04da0b5521ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd.json delete mode 100644 .sqlx/query-f681f25833875063150b6f881ebbfadc292939745e38754991650651edbe7acb.json diff --git a/.sqlx/query-ca7db50451b04da0b5521ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd.json b/.sqlx/query-ca7db50451b04da0b5521ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd.json new file mode 100644 index 0000000000..13bf91b98e --- /dev/null +++ b/.sqlx/query-ca7db50451b04da0b5521ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nINSERT INTO media_files\nSET\n token = ?,\n\n origin_category = ?,\n origin_product_category = ?,\n maybe_origin_model_type = ?,\n\n media_type = ?,\n maybe_mime_type = ?,\n file_size_bytes = ?,\n\n checksum_sha2 = ?,\n\n maybe_prompt_token = ?,\n\n public_bucket_directory_hash = ?,\n maybe_public_bucket_prefix = ?,\n maybe_public_bucket_extension = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_anonymous_visitor_token = ?,\n creator_ip_address = ?,\n\n creator_set_visibility = ?,\n extra_file_modification_info = ?,\n\n maybe_creator_file_synthetic_id = ?,\n maybe_creator_category_synthetic_id = ?,\n\n is_generated_on_prem = ?,\n generated_by_worker = ?,\n generated_by_cluster = ?\n\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 22 + }, + "nullable": [] + }, + "hash": "ca7db50451b04da0b5521ebc12416280d4fe6cbc1da47b7fa314d810b5f3a6fd" +} diff --git a/.sqlx/query-f681f25833875063150b6f881ebbfadc292939745e38754991650651edbe7acb.json b/.sqlx/query-f681f25833875063150b6f881ebbfadc292939745e38754991650651edbe7acb.json deleted file mode 100644 index f6a1155e30..0000000000 --- a/.sqlx/query-f681f25833875063150b6f881ebbfadc292939745e38754991650651edbe7acb.json +++ /dev/null @@ -1,12 +0,0 @@ -{ - "db_name": "MySQL", - "query": "\nINSERT INTO media_files\nSET\n token = ?,\n\n origin_category = ?,\n origin_product_category = ?,\n maybe_origin_model_type = ?,\n\n media_type = ?,\n maybe_mime_type = ?,\n file_size_bytes = ?,\n\n checksum_sha2 = ?,\n\n public_bucket_directory_hash = ?,\n maybe_public_bucket_prefix = ?,\n maybe_public_bucket_extension = ?,\n\n maybe_creator_user_token = ?,\n maybe_creator_anonymous_visitor_token = ?,\n creator_ip_address = ?,\n\n creator_set_visibility = ?,\n extra_file_modification_info = ?,\n\n maybe_creator_file_synthetic_id = ?,\n maybe_creator_category_synthetic_id = ?,\n\n is_generated_on_prem = ?,\n generated_by_worker = ?,\n generated_by_cluster = ?\n\n ", - "describe": { - "columns": [], - "parameters": { - "Right": 21 - }, - "nullable": [] - }, - "hash": "f681f25833875063150b6f881ebbfadc292939745e38754991650651edbe7acb" -} From 5e5264c65cea10964e6cf375a236418fce309ac1 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:14:24 -0500 Subject: [PATCH 08/22] debugging errors --- .../workflow/comfy_ui/process_job.rs | 19 +++++++++++++++---- 1 file changed, 15 insertions(+), 4 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 233d01b063..7787d5b4e2 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -296,27 +296,38 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result {} } + // Download Input file if specified let mut maybe_input_path: Option = None; let mut maybe_original_input_path: Option = None; + match job_args.maybe_input_file { - Some(input_file) => { + Some(input_media_file_token) => { let input_dir = root_comfy_path.join("input"); // make if not exist if !input_dir.exists() { - std::fs::create_dir_all(&input_dir).unwrap(); + std::fs::create_dir_all(&input_dir) + .map_err(|err| ProcessSingleJobError::IoError(err))?; } + + info!("Querying input media file by token: {:?}", &input_media_file_token); + let retrieved_input_record = get_media_file( - input_file, + input_media_file_token, false, &deps.db.mysql_pool - ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("Input file not found")))?; + ).await?.ok_or_else(|| { + error!("input media_file not found: {:?}", &input_media_file_token); + ProcessSingleJobError::Other(anyhow!("input media_file not found: {:?}", input_media_file_token)) + })?; let media_file_bucket_path = MediaFileBucketPath::from_object_hash( &retrieved_input_record.public_bucket_directory_hash, retrieved_input_record.maybe_public_bucket_prefix.as_deref(), retrieved_input_record.maybe_public_bucket_extension.as_deref()); + info!("Input media file cloud bucket path: {:?}", media_file_bucket_path.get_full_object_path_str()); + //let input_filename = format!("input.{}", retrieved_input_record.maybe_public_bucket_extension.unwrap()); let input_filename = "video.mp4".to_string(); let input_path = path_to_string(input_dir.join(input_filename)); From 1fac0e634b0dc9374d6e0153695013dc65f38753 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:23:08 -0500 Subject: [PATCH 09/22] more error logging --- .../workflow/comfy_ui/process_job.rs | 31 ++++++++++++++----- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 7787d5b4e2..5af4d1d03c 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -240,15 +240,20 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result = None; + match job_args.maybe_sd_model { - Some(sd_model) => { + Some(sd_model_weight_token) => { let sd_dir = root_comfy_path.join("models").join("checkpoints"); // make if not exist if !sd_dir.exists() { - std::fs::create_dir_all(&sd_dir).unwrap(); + std::fs::create_dir_all(&sd_dir) + .map_err(|err| ProcessSingleJobError::IoError(err))?; } + + info!("Querying SD model by token: {:?} ...", &sd_model_weight_token); + let retrieved_sd_record = get_weight_by_token( - sd_model, + sd_model_weight_token, false, &deps.db.mysql_pool ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("SD model not found")))?; @@ -258,6 +263,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result {} } + // Download Lora model if specified let mut maybe_lora_path: Option = None; + match job_args.maybe_lora_model { - Some(lora_model) => { + Some(lora_model_weight_token) => { let lora_dir = root_comfy_path.join("models").join("loras"); // make if not exist if !lora_dir.exists() { - std::fs::create_dir_all(&lora_dir).unwrap(); + std::fs::create_dir_all(&lora_dir) + .map_err(|err| ProcessSingleJobError::IoError(err))?; } + + info!("Querying lora model by token: {:?} ...", &lora_model_weight_token); + let retrieved_lora_record = get_weight_by_token( - lora_model, + lora_model_weight_token, false, &deps.db.mysql_pool ).await?.ok_or_else(|| ProcessSingleJobError::Other(anyhow!("Lora model not found")))?; + let bucket_details = RemoteCloudBucketDetails { object_hash: retrieved_lora_record.public_bucket_hash, prefix: retrieved_lora_record.maybe_public_bucket_prefix.unwrap(), @@ -310,7 +323,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result "video", From cb68c1cb163ea0f239119ad4696916a287bbb520 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:31:21 -0500 Subject: [PATCH 10/22] more error logging --- .../src/job/job_types/workflow/comfy_ui/process_job.rs | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 5af4d1d03c..3068a3974b 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -256,7 +256,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result Date: Fri, 1 Mar 2024 03:32:04 -0500 Subject: [PATCH 11/22] allow deleted SD + lora in comfy --- .../src/job/job_types/workflow/comfy_ui/process_job.rs | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 3068a3974b..febb155337 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -254,7 +254,7 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result Date: Fri, 1 Mar 2024 03:41:45 -0500 Subject: [PATCH 12/22] more logging --- .../workflow/comfy_ui/process_job.rs | 41 ++++++++++--------- 1 file changed, 22 insertions(+), 19 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index febb155337..0030acfad2 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -384,7 +384,10 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result) -> Result {} - Err(err) => { - error!("No prompt result token? something has failed: {:?}", err); - return Err(ProcessSingleJobError::from_anyhow_error( - anyhow!("No prompt result token? something has failed."))); - } + if let Err(err) = prompt_result { + error!("No prompt result token? something has failed: {:?} (we'll ignore this error)", err); } } From 3acbf33c20abc9241f1c7ed6466473433a84595f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 05:25:29 -0500 Subject: [PATCH 13/22] rebase --- .../comfy_ui/comfy_ui_dependencies.rs | 3 + .../workflow/comfy_ui/process_job.rs | 70 +++++++++++++++---- 2 files changed, 61 insertions(+), 12 deletions(-) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs index 4d1d3b16ef..906cdbf82d 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs @@ -3,9 +3,11 @@ use std::process::Command; use errors::AnyhowResult; use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand; +use crate::util::common_commands::ffmpeg_logo_watermark_command::FfmpegLogoWatermarkCommand; pub struct ComfyDependencies { pub inference_command: ComfyInferenceCommand, + pub ffmpeg_watermark_command: FfmpegLogoWatermarkCommand, pub dependency_tokens: RequiredModels, } @@ -29,6 +31,7 @@ impl ComfyDependencies { Ok(Self { inference_command: inference_command, dependency_tokens: RequiredModels::init(), + ffmpeg_watermark_command: FfmpegLogoWatermarkCommand::from_env()?, }) } } diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index 0030acfad2..d5dfd9eb2a 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -38,6 +38,7 @@ use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::InferenceArgs; use crate::job::job_types::workflow::comfy_ui::validate_job::validate_job; use crate::job_dependencies::JobDependencies; +use crate::util::common_commands::ffmpeg_logo_watermark_command; fn get_file_extension(mimetype: &str) -> Result<&'static str> { let ext = match mimetype { @@ -442,19 +443,18 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result) -> Result) -> Result) -> Result) -> Result Date: Sun, 3 Mar 2024 22:20:36 -0500 Subject: [PATCH 14/22] refactor prompt classification --- .../prompt_engineering/src/classify_prompt.rs | 67 ++++--------------- .../src/keywords/contains_children_keyword.rs | 16 +++++ .../contains_potential_minor_keyword.rs | 16 +++++ .../src/keywords/contains_racist_keyword.rs | 16 +++++ .../src/keywords/contains_sex_keyword.rs | 15 +++++ .../prompt_engineering/src/keywords/mod.rs | 4 ++ crates/lib/prompt_engineering/src/lib.rs | 4 +- .../src/potential_minor_regex.rs | 62 ----------------- .../src/{ => regexes}/minor_regex.rs | 7 +- .../lib/prompt_engineering/src/regexes/mod.rs | 3 + .../src/regexes/nsfw_regex.rs | 40 +++++++++++ .../src/regexes/potential_minor_regex.rs | 62 +++++++++++++++++ .../tests/borderline_tests.rs | 3 +- .../tests/csam_prompt_tests.rs | 3 +- .../tests/nsfw_prompt_tests.rs | 3 +- .../tests/racist_prompt_tests.rs | 2 +- .../tests/sfw_prompt_tests.rs | 3 +- .../tests/spicy_prompt_tests.rs | 3 +- .../dictionary_potential_minor_terms.txt | 6 ++ 19 files changed, 209 insertions(+), 126 deletions(-) create mode 100644 crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs create mode 100644 crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs create mode 100644 crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs create mode 100644 crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs create mode 100644 crates/lib/prompt_engineering/src/keywords/mod.rs delete mode 100644 crates/lib/prompt_engineering/src/potential_minor_regex.rs rename crates/lib/prompt_engineering/src/{ => regexes}/minor_regex.rs (98%) create mode 100644 crates/lib/prompt_engineering/src/regexes/mod.rs create mode 100644 crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs create mode 100644 crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs create mode 100644 includes/binary_includes/dictionary_potential_minor_terms.txt diff --git a/crates/lib/prompt_engineering/src/classify_prompt.rs b/crates/lib/prompt_engineering/src/classify_prompt.rs index 80172c42c2..7f55578d95 100644 --- a/crates/lib/prompt_engineering/src/classify_prompt.rs +++ b/crates/lib/prompt_engineering/src/classify_prompt.rs @@ -1,22 +1,21 @@ -use std::collections::HashSet; - use once_cell::sync::Lazy; use regex::Regex; -use crate::minor_regex::lowercase_mentions_underage; -use crate::potential_minor_regex::lowercase_mentions_potential_underage; - -const CHILD_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_children_terms.txt"); -const SEX_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_sex_terms.txt"); -const RACIST_TERMS : &str = include_str!("../../../../includes/binary_includes/dictionary_racist_terms.txt"); +use crate::keywords::contains_children_keyword::contains_children_keyword; +use crate::keywords::contains_potential_minor_keyword::contains_potential_minor_keyword; +use crate::keywords::contains_racist_keyword::contains_racist_keyword; +use crate::keywords::contains_sex_keyword::contains_sex_keyword; +use crate::regexes::minor_regex::lowercase_mentions_underage; +use crate::regexes::nsfw_regex::nsfw_regex; +use crate::regexes::potential_minor_regex::potential_minor_regex; /// Classification of a text prompt pub struct PromptClassification { pub prompt_references_children: bool, pub prompt_references_potential_minors: bool, pub prompt_references_sex: bool, - pub prompt_references_violence: bool, pub prompt_references_racism: bool, + pub prompt_references_violence: bool, } impl PromptClassification { @@ -48,59 +47,21 @@ pub fn classify_prompt(text_prompt: &str) -> PromptClassification { prompt_tokens.extend(prompt_tokens_alpha); prompt_tokens.extend(prompt_tokens_lower); - let prompt_references_children = references_children(&prompt_tokens) || lowercase_mentions_underage(&text_prompt); - let prompt_references_potential_minors = lowercase_mentions_potential_underage(&text_prompt); - let prompt_references_sex = references_sex(&prompt_tokens); - let prompt_references_violence = references_violence(&prompt_tokens); - let prompt_references_racism = references_racism(&prompt_tokens); + let prompt_references_children = contains_children_keyword(&prompt_tokens) || lowercase_mentions_underage(&text_prompt); + let prompt_references_potential_minors = contains_potential_minor_keyword(&prompt_tokens) || potential_minor_regex(&text_prompt); + let prompt_references_sex = contains_sex_keyword(&prompt_tokens) || nsfw_regex(&text_prompt); + let prompt_references_racism = contains_racist_keyword(&prompt_tokens); + let prompt_references_violence = false; // TODO PromptClassification { prompt_references_children, prompt_references_potential_minors, prompt_references_sex, - prompt_references_violence, prompt_references_racism, + prompt_references_violence, } } -fn references_children(prompt_tokens: &[String]) -> bool { - static CHILD_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - CHILD_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| CHILD_TERM_DICTIONARY.contains(term)) -} - -fn references_sex(prompt_tokens: &[String]) -> bool { - static SEX_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - SEX_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| SEX_TERM_DICTIONARY.contains(term)) -} - -fn references_violence(_prompt_tokens: &[String]) -> bool { - // TODO(bt,2024-02-19): Need to implement - false -} - -fn references_racism(prompt_tokens: &[String]) -> bool { - static RACIST_TERM_DICTIONARY : Lazy> = Lazy::new(|| { - RACIST_TERMS.lines() - .map(|line| line.trim().to_lowercase()) - .filter(|line| !line.is_empty() && !line.starts_with("#")) - .collect::>() - }); - - prompt_tokens.iter().any(|term| RACIST_TERM_DICTIONARY.contains(term)) -} - static SPACES_REGEX : Lazy = Lazy::new(|| { // NB: Includes non-ascii, such as "，". Regex::new(r"[\s,;，]+").expect("regex should be valid") diff --git a/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs new file mode 100644 index 0000000000..375b4a9436 --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_children_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_children_terms.txt"); + +pub (crate) fn contains_children_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs new file mode 100644 index 0000000000..b6cc17b8fb --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_potential_minor_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS: &str = include_str!("../../../../../includes/binary_includes/dictionary_potential_minor_terms.txt"); + +pub (crate) fn contains_potential_minor_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY: Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs new file mode 100644 index 0000000000..6c1609cc8d --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_racist_keyword.rs @@ -0,0 +1,16 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_racist_terms.txt"); + +pub (crate) fn contains_racist_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs b/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs new file mode 100644 index 0000000000..f83455c35c --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/contains_sex_keyword.rs @@ -0,0 +1,15 @@ +use std::collections::HashSet; + +use once_cell::sync::Lazy; + +const TERMS : &str = include_str!("../../../../../includes/binary_includes/dictionary_sex_terms.txt"); +pub (crate) fn contains_sex_keyword(prompt_tokens: &[String]) -> bool { + static TERM_DICTIONARY : Lazy> = Lazy::new(|| { + TERMS.lines() + .map(|line| line.trim().to_lowercase()) + .filter(|line| !line.is_empty() && !line.starts_with("#")) + .collect::>() + }); + + prompt_tokens.iter().any(|term| TERM_DICTIONARY.contains(term)) +} diff --git a/crates/lib/prompt_engineering/src/keywords/mod.rs b/crates/lib/prompt_engineering/src/keywords/mod.rs new file mode 100644 index 0000000000..a6d4629d2d --- /dev/null +++ b/crates/lib/prompt_engineering/src/keywords/mod.rs @@ -0,0 +1,4 @@ +pub (crate) mod contains_children_keyword; +pub (crate) mod contains_potential_minor_keyword; +pub (crate) mod contains_racist_keyword; +pub (crate) mod contains_sex_keyword; \ No newline at end of file diff --git a/crates/lib/prompt_engineering/src/lib.rs b/crates/lib/prompt_engineering/src/lib.rs index af1c9d67ef..fe2fe32b30 100644 --- a/crates/lib/prompt_engineering/src/lib.rs +++ b/crates/lib/prompt_engineering/src/lib.rs @@ -18,7 +18,7 @@ #![allow(dead_code)] #![allow(non_snake_case)] -pub (crate) mod minor_regex; -pub (crate) mod potential_minor_regex; +pub (crate) mod regexes; +pub (crate) mod keywords; pub mod classify_prompt; diff --git a/crates/lib/prompt_engineering/src/potential_minor_regex.rs b/crates/lib/prompt_engineering/src/potential_minor_regex.rs deleted file mode 100644 index 1e7311b611..0000000000 --- a/crates/lib/prompt_engineering/src/potential_minor_regex.rs +++ /dev/null @@ -1,62 +0,0 @@ -use once_cell::sync::Lazy; -use regex::RegexSet; - -pub (crate) fn lowercase_mentions_potential_underage(text: &str) -> bool { - let text = text.to_lowercase(); - - // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward - // Matching digits "0" and "1" are tricky because they would match "19" and "0100". - static AGE_REGEX: Lazy = Lazy::new(|| { - let patterns = [ - // English - r"\b(teen(age(r)?)?s?)\b", - r"\b(young|tiny|small|petite)\b(.{0,7})\b(girl|boy)s?\b", - - ]; - RegexSet::new(&patterns).expect("regex should be valid") - }); - - //println!("text: {:?}", text); - - AGE_REGEX.matches(&text).matched_any() -} - -#[cfg(test)] -mod tests { - - mod english { - - mod positive_match { - use crate::potential_minor_regex::lowercase_mentions_potential_underage; - - #[test] - fn spot_check_matches() { - assert!(lowercase_mentions_potential_underage("teen")); - assert!(lowercase_mentions_potential_underage("teens")); - assert!(lowercase_mentions_potential_underage("teenage")); - assert!(lowercase_mentions_potential_underage("teenager")); - assert!(lowercase_mentions_potential_underage("teenagers")); - - assert!(lowercase_mentions_potential_underage("tiny girl")); - assert!(lowercase_mentions_potential_underage("tiny boy")); - assert!(lowercase_mentions_potential_underage("small boy")); - assert!(lowercase_mentions_potential_underage("small girl")); - - assert!(lowercase_mentions_potential_underage("tiny girls")); - assert!(lowercase_mentions_potential_underage("tiny boys")); - assert!(lowercase_mentions_potential_underage("small boys")); - assert!(lowercase_mentions_potential_underage("small girls")); - - assert!(lowercase_mentions_potential_underage("young girl")); - assert!(lowercase_mentions_potential_underage("young boy")); - assert!(lowercase_mentions_potential_underage("young boy")); - assert!(lowercase_mentions_potential_underage("young girl")); - - assert!(lowercase_mentions_potential_underage("petite girl")); - assert!(lowercase_mentions_potential_underage("petite boy")); - assert!(lowercase_mentions_potential_underage("petite boy")); - assert!(lowercase_mentions_potential_underage("petite girl")); - } - } - } -} diff --git a/crates/lib/prompt_engineering/src/minor_regex.rs b/crates/lib/prompt_engineering/src/regexes/minor_regex.rs similarity index 98% rename from crates/lib/prompt_engineering/src/minor_regex.rs rename to crates/lib/prompt_engineering/src/regexes/minor_regex.rs index 202ee08e55..da4110e326 100644 --- a/crates/lib/prompt_engineering/src/minor_regex.rs +++ b/crates/lib/prompt_engineering/src/regexes/minor_regex.rs @@ -35,7 +35,7 @@ pub (crate) fn lowercase_mentions_underage(text: &str) -> bool { #[cfg(test)] mod tests { - use crate::minor_regex::lowercase_mentions_underage; + use crate::regexes::minor_regex::lowercase_mentions_underage; mod english { use super::*; @@ -117,7 +117,7 @@ mod tests { mod spanish { use super::*; - //mod non_match { +//mod non_match { // #[test] // fn prefix_and_suffix_digits() { // // TODO @@ -184,7 +184,7 @@ mod tests { mod user_inputs { use speculoos::asserting; - use crate::minor_regex::lowercase_mentions_underage; + use super::*; #[test] fn test_bad_user_csam_inputs() { @@ -199,6 +199,7 @@ mod tests { "a blonde girl of 13 years old with blue eyes.she's naked ,showing her private parts. i'm fucking her now", "a blonde girl of 13 years old with blue eyes.she's naked ,showing her tits and ass", "a little boy have brown eyes and hair，nude body", + "Cute hot 15 year old girl, no bra, no shirt.", "Fiveteen years old, naked, Brown hair, big ass, Brest d, open mouth, close eyes, seamen face", "naked 8yo blonde girl", "seven year old girl naked legs spread", diff --git a/crates/lib/prompt_engineering/src/regexes/mod.rs b/crates/lib/prompt_engineering/src/regexes/mod.rs new file mode 100644 index 0000000000..0fdb5b806e --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/mod.rs @@ -0,0 +1,3 @@ +pub (crate) mod minor_regex; +pub (crate) mod nsfw_regex; +pub (crate) mod potential_minor_regex; diff --git a/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs b/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs new file mode 100644 index 0000000000..91f2c309b1 --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/nsfw_regex.rs @@ -0,0 +1,40 @@ +use once_cell::sync::Lazy; +use regex::RegexSet; + +pub (crate) fn nsfw_regex(text: &str) -> bool { + let text = text.to_lowercase(); + + // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward + // Matching digits "0" and "1" are tricky because they would match "19" and "0100". + static AGE_REGEX: Lazy = Lazy::new(|| { + let patterns = [ + // English + r"\b(no|without|missing)\b(.{0,7})\b(clothes|clothing|bra|panties|panty|clothing)s?\b", + + ]; + RegexSet::new(&patterns).expect("regex should be valid") + }); + + //println!("text: {:?}", text); + + AGE_REGEX.matches(&text).matched_any() +} + +#[cfg(test)] +mod tests { + + mod english { + + mod positive_match { + use crate::regexes::nsfw_regex::nsfw_regex; + + #[test] + fn spot_check_matches() { + assert!(nsfw_regex("no panties")); + assert!(nsfw_regex("without panties")); + assert!(nsfw_regex("no clothing")); + assert!(nsfw_regex("missing clothing")); + } + } + } +} diff --git a/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs b/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs new file mode 100644 index 0000000000..2d4b5761ec --- /dev/null +++ b/crates/lib/prompt_engineering/src/regexes/potential_minor_regex.rs @@ -0,0 +1,62 @@ +use once_cell::sync::Lazy; +use regex::RegexSet; + +pub (crate) fn potential_minor_regex(text: &str) -> bool { + let text = text.to_lowercase(); + + // NB: Unfortunately Rust Regex does not support look-around, so our age matching is awkward + // Matching digits "0" and "1" are tricky because they would match "19" and "0100". + static AGE_REGEX: Lazy = Lazy::new(|| { + let patterns = [ + // English + r"\b(teen(age(r)?)?s?)\b", + r"\b(young|tiny|small|petite)\b(.{0,7})\b(girl|boy)s?\b", + + ]; + RegexSet::new(&patterns).expect("regex should be valid") + }); + + //println!("text: {:?}", text); + + AGE_REGEX.matches(&text).matched_any() +} + +#[cfg(test)] +mod tests { + + mod english { + + mod positive_match { + use crate::regexes::potential_minor_regex::potential_minor_regex; + + #[test] + fn spot_check_matches() { + assert!(potential_minor_regex("teen")); + assert!(potential_minor_regex("teens")); + assert!(potential_minor_regex("teenage")); + assert!(potential_minor_regex("teenager")); + assert!(potential_minor_regex("teenagers")); + + assert!(potential_minor_regex("tiny girl")); + assert!(potential_minor_regex("tiny boy")); + assert!(potential_minor_regex("small boy")); + assert!(potential_minor_regex("small girl")); + + assert!(potential_minor_regex("tiny girls")); + assert!(potential_minor_regex("tiny boys")); + assert!(potential_minor_regex("small boys")); + assert!(potential_minor_regex("small girls")); + + assert!(potential_minor_regex("young girl")); + assert!(potential_minor_regex("young boy")); + assert!(potential_minor_regex("young boy")); + assert!(potential_minor_regex("young girl")); + + assert!(potential_minor_regex("petite girl")); + assert!(potential_minor_regex("petite boy")); + assert!(potential_minor_regex("petite boy")); + assert!(potential_minor_regex("petite girl")); + } + } + } +} diff --git a/crates/lib/prompt_engineering/tests/borderline_tests.rs b/crates/lib/prompt_engineering/tests/borderline_tests.rs index 5828707ffe..422544948d 100644 --- a/crates/lib/prompt_engineering/tests/borderline_tests.rs +++ b/crates/lib/prompt_engineering/tests/borderline_tests.rs @@ -1,7 +1,8 @@ //! Tests that could be close to CSAM, but appear to be completely SFW use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn borderline_cases() { diff --git a/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs b/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs index aaa5bec405..40cf010328 100644 --- a/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/csam_prompt_tests.rs @@ -3,7 +3,7 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn sfw_prompts_including_children() { @@ -44,6 +44,7 @@ fn test_english_bad_user_csam_prompts() { "A young girl how eats her classmates with her butt-hole", "Beautiful young naked preteen girl hugging a young naked boy on a couch in front of a window. Boy Touch tits. Perfect teen’s breasts. Laneya Grace. superflat, a photorealistic 16k", "Boy jerking off", + "Cute hot 15 year old girl, no bra, no shirt.", "Fiveteen years old, naked, Brown hair, big ass, Brest d, open mouth, close eyes, seamen face", "Little girl pussy cum dripping out", "Twelve years old girl , sucking thirty years old cock men", diff --git a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs index 343bbac1a8..4e077d7550 100644 --- a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs @@ -1,5 +1,6 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn english_nsfw_prompts_test() { diff --git a/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs b/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs index 6e148d9ab6..38288ccd62 100644 --- a/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/racist_prompt_tests.rs @@ -1,6 +1,6 @@ use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn test_bad_user_racist_prompts() { diff --git a/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs b/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs index 9639a58e4f..f1dfc731b0 100644 --- a/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/sfw_prompt_tests.rs @@ -1,7 +1,8 @@ //! These prompts should be fully innocent and allowed use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn sfw_english_test() { diff --git a/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs b/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs index 85cb1b7f2c..7190369ab3 100644 --- a/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/spicy_prompt_tests.rs @@ -3,7 +3,8 @@ //! to steer the image generation away from NSFW territory. //! use speculoos::asserting; -use classify_text::classify_prompt::classify_prompt; + +use prompt_engineering::classify_prompt::classify_prompt; #[test] fn permissible_sexy_tests() { diff --git a/includes/binary_includes/dictionary_potential_minor_terms.txt b/includes/binary_includes/dictionary_potential_minor_terms.txt new file mode 100644 index 0000000000..99f52fae54 --- /dev/null +++ b/includes/binary_includes/dictionary_potential_minor_terms.txt @@ -0,0 +1,6 @@ +# These are words that could refer to adults (18+), but also might be interpreted as minors +# English +teen +teenage +young +younger From debce383769ebed143594bbf0ffa2b4e13daecab Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 3 Mar 2024 22:43:10 -0500 Subject: [PATCH 15/22] tests --- .../tests/borderline_tests.rs | 7 +++- .../tests/nsfw_prompt_tests.rs | 42 +++++++++---------- .../dictionary_children_terms.txt | 1 - .../dictionary_potential_minor_terms.txt | 6 ++- 4 files changed, 32 insertions(+), 24 deletions(-) diff --git a/crates/lib/prompt_engineering/tests/borderline_tests.rs b/crates/lib/prompt_engineering/tests/borderline_tests.rs index 422544948d..fe065ee080 100644 --- a/crates/lib/prompt_engineering/tests/borderline_tests.rs +++ b/crates/lib/prompt_engineering/tests/borderline_tests.rs @@ -16,6 +16,7 @@ fn borderline_cases() { "Create an Korean child playing with a polish child on play ground", "Demon boy teen that is cute hot in a black suit and tie", "girl in a strapless dress in a feald of flower", + "navel of a girl is eating banana", "two girls kissing in bed", ]); } @@ -31,6 +32,10 @@ fn assert_prompts_are_sfw(prompts: &[&str]) { asserting(&format!("is child abuse - prompt: {}", prompt)) .that(&classification.is_child_abuse()) .is_equal_to(false); + + asserting(&format!("references children or potential minors - prompt: {}", prompt)) + .that(&(classification.prompt_references_children || classification.prompt_references_potential_minors)) + + .is_equal_to(true); } } - diff --git a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs index 4e077d7550..84b9f3d7f5 100644 --- a/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs +++ b/crates/lib/prompt_engineering/tests/nsfw_prompt_tests.rs @@ -9,6 +9,7 @@ fn english_nsfw_prompts_test() { "1girl, nude, large breasts", "2B. white swimsuit, butt, black eyeliner, white droplets on ass, sunscreen bottle in hand", "3 naked Latina women in a swimming pool", + "A beautiful, young, curvy woman with ginger hair, huge breasts twice the size of her head, wearing (red glasses) and a (((sexy white nurse costume with a front red zip))), playfully flirting with the viewer; ((pulling on the front zip, showing off her cleavage)), (no background; solid colour background). Mind-broken; sexual ecstasy, ahegao, happy, (Horny, naughty face), ((cock hungry, cock worship, penis awe)), 3/4 head shot; head, hair and arms in frame", "a gay black guy with nipples and has undies on", "A giant woman living in a city she loves putting people in her butthole", "A girl who used her ass to eat vore", @@ -20,23 +21,6 @@ fn english_nsfw_prompts_test() { "a woman having sex with a man", "Adult film star Lisa Ann, naked, posing for viewer ", "ahh daddy ahh fuck me", - "hot and pretty faced naked girls", - "hot girl with a tight vagina", - "vagina", - "hot naked girl", - "hot anime girl gaping pussy and ass legs wide open", - "hot gaping pussy and ass close up", - "girls gaping pussy and ass close up in the camera", - "girl getting fucked in a pussy and another girl showing her ass in the image", - "hot girl gaping pussy and ass close up in the camera", - "hot cat girl pussy getting fucked by a dick", - "hot wet gaping pussy close up", - "hot naked girl full body", - "juicy wet pussy close up", - "Anime girls squiring from thier vaginas", - "Naked girl showing her cum covered vagina ", - "Wet vagina after cream pie\n", - "The wettest vagina", "Anal vore", "Anime Girl (NSFW)", "anime girl getting fucked in the pussy", @@ -68,6 +52,7 @@ fn english_nsfw_prompts_test() { "Futanari holding her penis in one hand", "Futanari", "Generate a beautiful picture of a naked woman.", + "girl getting fucked in a pussy and another girl showing her ass in the image", "girl in a tight thong. NSFW, k, realistic, photorealsitic.", "Girl naked, masturbating, ahegao, squatting", "Girl naked, masturbating, ahegao, squatting, ultra realistic", @@ -75,17 +60,23 @@ fn english_nsfw_prompts_test() { "Girl with big ass in white pants shake her ass", "Girl with big ass sits on chair on tight skirt", "Girl with big boobs and small bro touch her boobs and make a lusty eye contact", + "girls gaping pussy and ass close up in the camera", "guy shitting in the toilet", "hentai", "Hitomi Tanaka with nude big boobs and pussy open", "hot and pretty faced naked girls", + "hot anime girl gaping pussy and ass legs wide open", "hot cat girl pussy getting fucked by a dick", + "hot gaping pussy and ass close up", "hot gaping pussy and ass legs wide open", + "hot girl gaping pussy and ass close up in the camera", "hot girl with a tight vagina", "Hot girl with big boobs in white dress", "hot hentai fuck", "Hot naked anime maids", + "hot naked girl full body", "hot naked girl", + "hot wet gaping pussy close up", "hot white girl gaping pussy and ass close up in the camera", "Hottest anime girl spreading apart her vagina", "Huge breasts naked anime girl", @@ -99,6 +90,7 @@ fn english_nsfw_prompts_test() { "intercource between man and women", "Iran blowjob girl", "Iranian girl suck dick", + "juicy wet pussy close up", "Kyu Sugardust from the game Huniepop", "leaked sex tape", "Lesbians in a park nude", @@ -114,6 +106,7 @@ fn english_nsfw_prompts_test() { "naked anime girl", "Naked Gardevoir", "Naked girl showing her creampied vagina", + "Naked girl showing her cum covered vagina ", "Naked girl showing her cum covered vagina", "naked girl sitting in chair", "Naked girl spreading her pussy", @@ -144,6 +137,7 @@ fn english_nsfw_prompts_test() { "short blonde hair side shaved, black lingerie, black eyeliner, cum on face", "The Love Fairy Kyu Sugardust from the game Huniepop", "The wettest vagina ", + "The wettest vagina", "titfuck, huge boobs", "Vagina cuming", "Vagina hole", @@ -151,6 +145,7 @@ fn english_nsfw_prompts_test() { "vagina", "Vegeta female version with big boobs from Dragonball", "Wet vagina after cream pie", + "Wet vagina after cream pie\n", "woman pleasuring herself stylized in photorealist, 4k, high-res, and NSFW (Nude)", "woman sucking dick.", "Woman with huge tits desnude", @@ -173,6 +168,15 @@ fn spanish_nsfw_prompts_test() { fn assert_prompts_are_nsfw(prompts: &[&str]) { for prompt in prompts { let classification = classify_prompt(prompt); + + asserting(&format!("is nsfw - prompt: {}", prompt)) + .that(&classification.prompt_references_sex) + .is_equal_to(true); + + asserting(&format!("references children - prompt: {}", prompt)) + .that(&classification.prompt_references_children) + .is_equal_to(false); + asserting(&format!("is abusive - prompt: {}", prompt)) .that(&classification.is_abusive()) .is_equal_to(false); @@ -180,9 +184,5 @@ fn assert_prompts_are_nsfw(prompts: &[&str]) { asserting(&format!("is child abuse - prompt: {}", prompt)) .that(&classification.is_child_abuse()) .is_equal_to(false); - - asserting(&format!("is nsfw - prompt: {}", prompt)) - .that(&classification.prompt_references_sex) - .is_equal_to(true); } } diff --git a/includes/binary_includes/dictionary_children_terms.txt b/includes/binary_includes/dictionary_children_terms.txt index 6c872dda5b..2773aa28c5 100644 --- a/includes/binary_includes/dictionary_children_terms.txt +++ b/includes/binary_includes/dictionary_children_terms.txt @@ -45,7 +45,6 @@ toddlers tot tyke underage -young youngster youngsters youth diff --git a/includes/binary_includes/dictionary_potential_minor_terms.txt b/includes/binary_includes/dictionary_potential_minor_terms.txt index 99f52fae54..526c097e38 100644 --- a/includes/binary_includes/dictionary_potential_minor_terms.txt +++ b/includes/binary_includes/dictionary_potential_minor_terms.txt @@ -1,6 +1,10 @@ # These are words that could refer to adults (18+), but also might be interpreted as minors # English +boy +boys +girl +girls teen teenage young -younger +younger \ No newline at end of file From cd6201867ad1a40f2957591ae665724ef7c83993 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:32:00 -0500 Subject: [PATCH 16/22] media_files - maybe_media_subtype column --- .../down.sql | 5 +++++ .../up.sql | 7 +++++++ _sql/migrations_squashed/media_files.sql | 6 ++++++ .../src/job/job_types/workflow/comfy_ui/process_job.rs | 3 ++- 4 files changed, 20 insertions(+), 1 deletion(-) create mode 100644 _sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/down.sql create mode 100644 _sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/up.sql diff --git a/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/down.sql b/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/down.sql new file mode 100644 index 0000000000..43187f2514 --- /dev/null +++ b/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files DROP COLUMN maybe_media_subtype; diff --git a/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/up.sql b/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/up.sql new file mode 100644 index 0000000000..99eb343ae1 --- /dev/null +++ b/_sql/migrations/2024-02-29-234250_alter_media_files_add_media_subtype/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files + ADD COLUMN maybe_media_subtype VARCHAR(32) DEFAULT NULL + AFTER media_type; diff --git a/_sql/migrations_squashed/media_files.sql b/_sql/migrations_squashed/media_files.sql index 79fa38ee21..9083e974a8 100644 --- a/_sql/migrations_squashed/media_files.sql +++ b/_sql/migrations_squashed/media_files.sql @@ -81,6 +81,12 @@ CREATE TABLE media_files ( -- * 'video' for a variety of video types. media_type VARCHAR(16) NOT NULL, + -- A media file's possible subtype. Typically used for Storyteller Studio. + -- * 'mixamo' for mixamo animations (eg. for BVH, GLB, FBX, etc. files) + -- * 'mocap_net' for mocapnet animations (eg. for BVH files) + -- * 'scene' for generic scenes (eg. for BVH, GLB, FBX, etc. files) + maybe_media_subtype VARCHAR(32) DEFAULT NULL, + -- The file's mime type. maybe_mime_type VARCHAR(32) DEFAULT NULL, diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs index d5dfd9eb2a..4f231fe6e4 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/process_job.rs @@ -523,7 +523,8 @@ pub async fn process_job(args: ComfyProcessJobArgs<'_>) -> Result Date: Mon, 4 Mar 2024 03:42:44 -0500 Subject: [PATCH 17/22] enum for MediaFileSubtype --- .../media_files/media_file_origin_category.rs | 1 + .../media_files/media_file_subtype.rs | 135 ++++++++++++++++++ .../enums/src/by_table/media_files/mod.rs | 1 + 3 files changed, 137 insertions(+) create mode 100644 crates/_types/enums/src/by_table/media_files/media_file_subtype.rs diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_category.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_category.rs index 7cbab0645c..743fcd2510 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_category.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_category.rs @@ -94,6 +94,7 @@ mod tests { assert_eq!(MediaFileOriginCategory::Inference.to_str(), "inference"); assert_eq!(MediaFileOriginCategory::Processed.to_str(), "processed"); assert_eq!(MediaFileOriginCategory::Upload.to_str(), "upload"); + assert_eq!(MediaFileOriginCategory::DeviceApi.to_str(), "device_api"); assert_eq!(MediaFileOriginCategory::StoryEngine.to_str(), "story_engine"); } diff --git a/crates/_types/enums/src/by_table/media_files/media_file_subtype.rs b/crates/_types/enums/src/by_table/media_files/media_file_subtype.rs new file mode 100644 index 0000000000..f332e96099 --- /dev/null +++ b/crates/_types/enums/src/by_table/media_files/media_file_subtype.rs @@ -0,0 +1,135 @@ +use std::collections::BTreeSet; + +use serde::Deserialize; +use serde::Serialize; +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// Used in the `media_files` table in a `VARCHAR` field. +/// +/// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum MediaFileSubtype { + /// Animation file from Mixamo + /// Primarily used for FBX and GLB. + Mixamo, + + /// Animation file from MocapNet + /// Primarily used for BVH. + MocapNet, + + /// Generic 3D scene file. + /// Can pertain to BVH, GLB, FBX, etc. + Scene, +} + +// TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(MediaFileSubtype); +impl_mysql_enum_coders!(MediaFileSubtype); +impl_mysql_from_row!(MediaFileSubtype); + +/// NB: Legacy API for older code. +impl MediaFileSubtype { + pub fn to_str(&self) -> &'static str { + match self { + Self::Mixamo => "mixamo", + Self::Scene => "scene", + Self::MocapNet => "mocap_net", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "mixamo" => Ok(Self::Mixamo), + "mocap_net" => Ok(Self::MocapNet), + "scene" => Ok(Self::Scene), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::Mixamo, + Self::MocapNet, + Self::Scene, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::by_table::media_files::media_file_subtype::MediaFileSubtype; + use crate::test_helpers::assert_serialization; + + mod explicit_checks { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(MediaFileSubtype::Mixamo, "mixamo"); + assert_serialization(MediaFileSubtype::MocapNet, "mocap_net"); + assert_serialization(MediaFileSubtype::Scene, "scene"); + } + + #[test] + fn test_to_str() { + assert_eq!(MediaFileSubtype::Mixamo.to_str(), "mixamo"); + assert_eq!(MediaFileSubtype::MocapNet.to_str(), "mocap_net"); + assert_eq!(MediaFileSubtype::Scene.to_str(), "scene"); + } + + #[test] + fn test_from_str() { + assert_eq!(MediaFileSubtype::from_str("mixamo").unwrap(), MediaFileSubtype::Mixamo); + assert_eq!(MediaFileSubtype::from_str("mocap_net").unwrap(), MediaFileSubtype::MocapNet); + assert_eq!(MediaFileSubtype::from_str("scene").unwrap(), MediaFileSubtype::Scene); + assert!(MediaFileSubtype::from_str("foo").is_err()); + } + + #[test] + fn all_variants() { + let mut variants = MediaFileSubtype::all_variants(); + assert_eq!(variants.len(), 3); + assert_eq!(variants.pop_first(), Some(MediaFileSubtype::Mixamo)); + assert_eq!(variants.pop_first(), Some(MediaFileSubtype::MocapNet)); + assert_eq!(variants.pop_first(), Some(MediaFileSubtype::Scene)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(MediaFileSubtype::all_variants().len(), MediaFileSubtype::iter().len()); + } + + #[test] + fn round_trip() { + for variant in MediaFileSubtype::all_variants() { + assert_eq!(variant, MediaFileSubtype::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, MediaFileSubtype::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, MediaFileSubtype::from_str(&format!("{:?}", variant)).unwrap()); + } + } + + #[test] + fn serialized_length_ok_for_database() { + const MAX_LENGTH : usize = 32; + for variant in MediaFileSubtype::all_variants() { + let serialized = variant.to_str(); + assert!(serialized.len() > 0, "variant {:?} is too short", variant); + assert!(serialized.len() <= MAX_LENGTH, "variant {:?} is too long", variant); + } + } + } +} diff --git a/crates/_types/enums/src/by_table/media_files/mod.rs b/crates/_types/enums/src/by_table/media_files/mod.rs index 009ae6ae72..d506650184 100644 --- a/crates/_types/enums/src/by_table/media_files/mod.rs +++ b/crates/_types/enums/src/by_table/media_files/mod.rs @@ -1,4 +1,5 @@ pub mod media_file_origin_category; pub mod media_file_origin_model_type; pub mod media_file_origin_product_category; +pub mod media_file_subtype; pub mod media_file_type; \ No newline at end of file From 0236d6251ab39e1af8838e2029fb4b18431f3015 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:51:38 -0500 Subject: [PATCH 18/22] add subtype to GET media file endpoint --- ...96864593642ef495335c135d8356dc9ac2c5.json} | 67 +++++++++++-------- ...9903da1fa4ca6c0c9925b2a4e2bc791d4083.json} | 67 +++++++++++-------- .../src/queries/media_files/get_media_file.rs | 6 ++ .../endpoints/media_files/get_media_file.rs | 9 ++- 4 files changed, 92 insertions(+), 57 deletions(-) rename .sqlx/{query-2c5083331bf226ef10c419f6a2b6dee6ee776a3d41eea06de1eafd1d1bd41791.json => query-09a1a654d1493a700e89bf7d828d96864593642ef495335c135d8356dc9ac2c5.json} (72%) rename .sqlx/{query-726a6c62dc20e8372e950a10ecd5a438769448c7ce315bc8743b728e80179420.json => query-b161ade93a0b6c8b7bd01ad627659903da1fa4ca6c0c9925b2a4e2bc791d4083.json} (72%) diff --git a/.sqlx/query-2c5083331bf226ef10c419f6a2b6dee6ee776a3d41eea06de1eafd1d1bd41791.json b/.sqlx/query-09a1a654d1493a700e89bf7d828d96864593642ef495335c135d8356dc9ac2c5.json similarity index 72% rename from .sqlx/query-2c5083331bf226ef10c419f6a2b6dee6ee776a3d41eea06de1eafd1d1bd41791.json rename to .sqlx/query-09a1a654d1493a700e89bf7d828d96864593642ef495335c135d8356dc9ac2c5.json index 77ba2f77ad..9185f9a075 100644 --- a/.sqlx/query-2c5083331bf226ef10c419f6a2b6dee6ee776a3d41eea06de1eafd1d1bd41791.json +++ b/.sqlx/query-09a1a654d1493a700e89bf7d828d96864593642ef495335c135d8356dc9ac2c5.json @@ -1,6 +1,6 @@ { "db_name": "MySQL", - "query": "\nSELECT\n m.token as `token: tokens::tokens::media_files::MediaFileToken`,\n\n m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`,\n\n users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n\n m.maybe_batch_token as `maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken`,\n\n m.maybe_text_transcript,\n\n m.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`,\n\n m.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n model_weights.token as `maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken`,\n model_weights.title as maybe_model_weights_title,\n model_weights.weights_type as `maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n model_weights.weights_category as `maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n cover_image.public_bucket_directory_hash as maybe_model_cover_image_public_bucket_hash,\n cover_image.maybe_public_bucket_prefix as maybe_model_cover_image_public_bucket_prefix,\n cover_image.maybe_public_bucket_extension as maybe_model_cover_image_public_bucket_extension,\n\n model_weight_creator.token as `maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken`,\n model_weight_creator.username as maybe_model_weight_creator_username,\n model_weight_creator.display_name as maybe_model_weight_creator_display_name,\n model_weight_creator.email_gravatar_hash as maybe_model_weight_creator_gravatar_hash,\n\n m.public_bucket_directory_hash,\n m.maybe_public_bucket_prefix,\n m.maybe_public_bucket_extension,\n\n entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n entity_stats.bookmark_count as maybe_bookmark_count,\n\n m.created_at,\n m.updated_at\n\nFROM media_files AS m\nLEFT OUTER JOIN users\n ON m.maybe_creator_user_token = users.token\nLEFT OUTER JOIN model_weights\n ON m.maybe_origin_model_token = model_weights.token\nLEFT OUTER JOIN media_files as cover_image\n ON cover_image.token = model_weights.maybe_cover_image_media_file_token\nLEFT OUTER JOIN users as model_weight_creator\n ON model_weight_creator.token = model_weights.creator_user_token\nLEFT OUTER JOIN entity_stats\n ON entity_stats.entity_type = \"media_file\"\n AND entity_stats.entity_token = m.token\nWHERE\n m.token = ?\n AND m.user_deleted_at IS NULL\n AND m.mod_deleted_at IS NULL\n ", + "query": "\nSELECT\n m.token as `token: tokens::tokens::media_files::MediaFileToken`,\n\n m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`,\n m.maybe_media_subtype as `maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype`,\n\n users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n\n m.maybe_batch_token as `maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken`,\n\n m.maybe_text_transcript,\n\n m.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`,\n\n m.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n model_weights.token as `maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken`,\n model_weights.title as maybe_model_weights_title,\n model_weights.weights_type as `maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n model_weights.weights_category as `maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n cover_image.public_bucket_directory_hash as maybe_model_cover_image_public_bucket_hash,\n cover_image.maybe_public_bucket_prefix as maybe_model_cover_image_public_bucket_prefix,\n cover_image.maybe_public_bucket_extension as maybe_model_cover_image_public_bucket_extension,\n\n model_weight_creator.token as `maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken`,\n model_weight_creator.username as maybe_model_weight_creator_username,\n model_weight_creator.display_name as maybe_model_weight_creator_display_name,\n model_weight_creator.email_gravatar_hash as maybe_model_weight_creator_gravatar_hash,\n\n m.public_bucket_directory_hash,\n m.maybe_public_bucket_prefix,\n m.maybe_public_bucket_extension,\n\n entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n entity_stats.bookmark_count as maybe_bookmark_count,\n\n m.created_at,\n m.updated_at\n\nFROM media_files AS m\nLEFT OUTER JOIN users\n ON m.maybe_creator_user_token = users.token\nLEFT OUTER JOIN model_weights\n ON m.maybe_origin_model_token = model_weights.token\nLEFT OUTER JOIN media_files as cover_image\n ON cover_image.token = model_weights.maybe_cover_image_media_file_token\nLEFT OUTER JOIN users as model_weight_creator\n ON model_weight_creator.token = model_weights.creator_user_token\nLEFT OUTER JOIN entity_stats\n ON entity_stats.entity_type = \"media_file\"\n AND entity_stats.entity_token = m.token\nWHERE\n m.token = ?\n ", "describe": { "columns": [ { @@ -25,6 +25,16 @@ }, { "ordinal": 2, + "name": "maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, "name": "maybe_creator_user_token: tokens::tokens::users::UserToken", "type_info": { "type": "VarString", @@ -34,7 +44,7 @@ } }, { - "ordinal": 3, + "ordinal": 4, "name": "maybe_creator_username", "type_info": { "type": "VarString", @@ -44,7 +54,7 @@ } }, { - "ordinal": 4, + "ordinal": 5, "name": "maybe_creator_display_name", "type_info": { "type": "VarString", @@ -54,7 +64,7 @@ } }, { - "ordinal": 5, + "ordinal": 6, "name": "maybe_creator_gravatar_hash", "type_info": { "type": "String", @@ -64,7 +74,7 @@ } }, { - "ordinal": 6, + "ordinal": 7, "name": "maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken", "type_info": { "type": "VarString", @@ -74,7 +84,7 @@ } }, { - "ordinal": 7, + "ordinal": 8, "name": "maybe_text_transcript", "type_info": { "type": "Blob", @@ -84,7 +94,7 @@ } }, { - "ordinal": 8, + "ordinal": 9, "name": "maybe_prompt_token: tokens::tokens::prompts::PromptToken", "type_info": { "type": "VarString", @@ -94,7 +104,7 @@ } }, { - "ordinal": 9, + "ordinal": 10, "name": "creator_set_visibility: enums::common::visibility::Visibility", "type_info": { "type": "String", @@ -104,7 +114,7 @@ } }, { - "ordinal": 10, + "ordinal": 11, "name": "maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken", "type_info": { "type": "VarString", @@ -114,7 +124,7 @@ } }, { - "ordinal": 11, + "ordinal": 12, "name": "maybe_model_weights_title", "type_info": { "type": "VarString", @@ -124,7 +134,7 @@ } }, { - "ordinal": 12, + "ordinal": 13, "name": "maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType", "type_info": { "type": "VarString", @@ -134,7 +144,7 @@ } }, { - "ordinal": 13, + "ordinal": 14, "name": "maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory", "type_info": { "type": "VarString", @@ -144,7 +154,7 @@ } }, { - "ordinal": 14, + "ordinal": 15, "name": "maybe_model_cover_image_public_bucket_hash", "type_info": { "type": "VarString", @@ -154,7 +164,7 @@ } }, { - "ordinal": 15, + "ordinal": 16, "name": "maybe_model_cover_image_public_bucket_prefix", "type_info": { "type": "VarString", @@ -164,7 +174,7 @@ } }, { - "ordinal": 16, + "ordinal": 17, "name": "maybe_model_cover_image_public_bucket_extension", "type_info": { "type": "VarString", @@ -174,7 +184,7 @@ } }, { - "ordinal": 17, + "ordinal": 18, "name": "maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken", "type_info": { "type": "VarString", @@ -184,7 +194,7 @@ } }, { - "ordinal": 18, + "ordinal": 19, "name": "maybe_model_weight_creator_username", "type_info": { "type": "VarString", @@ -194,7 +204,7 @@ } }, { - "ordinal": 19, + "ordinal": 20, "name": "maybe_model_weight_creator_display_name", "type_info": { "type": "VarString", @@ -204,7 +214,7 @@ } }, { - "ordinal": 20, + "ordinal": 21, "name": "maybe_model_weight_creator_gravatar_hash", "type_info": { "type": "String", @@ -214,7 +224,7 @@ } }, { - "ordinal": 21, + "ordinal": 22, "name": "public_bucket_directory_hash", "type_info": { "type": "VarString", @@ -224,7 +234,7 @@ } }, { - "ordinal": 22, + "ordinal": 23, "name": "maybe_public_bucket_prefix", "type_info": { "type": "VarString", @@ -234,7 +244,7 @@ } }, { - "ordinal": 23, + "ordinal": 24, "name": "maybe_public_bucket_extension", "type_info": { "type": "VarString", @@ -244,7 +254,7 @@ } }, { - "ordinal": 24, + "ordinal": 25, "name": "maybe_ratings_positive_count", "type_info": { "type": "Long", @@ -254,7 +264,7 @@ } }, { - "ordinal": 25, + "ordinal": 26, "name": "maybe_ratings_negative_count", "type_info": { "type": "Long", @@ -264,7 +274,7 @@ } }, { - "ordinal": 26, + "ordinal": 27, "name": "maybe_bookmark_count", "type_info": { "type": "Long", @@ -274,7 +284,7 @@ } }, { - "ordinal": 27, + "ordinal": 28, "name": "created_at", "type_info": { "type": "Timestamp", @@ -284,7 +294,7 @@ } }, { - "ordinal": 28, + "ordinal": 29, "name": "updated_at", "type_info": { "type": "Timestamp", @@ -307,6 +317,7 @@ true, true, true, + true, false, true, true, @@ -329,5 +340,5 @@ false ] }, - "hash": "2c5083331bf226ef10c419f6a2b6dee6ee776a3d41eea06de1eafd1d1bd41791" + "hash": "09a1a654d1493a700e89bf7d828d96864593642ef495335c135d8356dc9ac2c5" } diff --git a/.sqlx/query-726a6c62dc20e8372e950a10ecd5a438769448c7ce315bc8743b728e80179420.json b/.sqlx/query-b161ade93a0b6c8b7bd01ad627659903da1fa4ca6c0c9925b2a4e2bc791d4083.json similarity index 72% rename from .sqlx/query-726a6c62dc20e8372e950a10ecd5a438769448c7ce315bc8743b728e80179420.json rename to .sqlx/query-b161ade93a0b6c8b7bd01ad627659903da1fa4ca6c0c9925b2a4e2bc791d4083.json index ebad316c68..374d785e1f 100644 --- a/.sqlx/query-726a6c62dc20e8372e950a10ecd5a438769448c7ce315bc8743b728e80179420.json +++ b/.sqlx/query-b161ade93a0b6c8b7bd01ad627659903da1fa4ca6c0c9925b2a4e2bc791d4083.json @@ -1,6 +1,6 @@ { "db_name": "MySQL", - "query": "\nSELECT\n m.token as `token: tokens::tokens::media_files::MediaFileToken`,\n\n m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`,\n\n users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n\n m.maybe_batch_token as `maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken`,\n\n m.maybe_text_transcript,\n\n m.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`,\n\n m.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n model_weights.token as `maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken`,\n model_weights.title as maybe_model_weights_title,\n model_weights.weights_type as `maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n model_weights.weights_category as `maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n cover_image.public_bucket_directory_hash as maybe_model_cover_image_public_bucket_hash,\n cover_image.maybe_public_bucket_prefix as maybe_model_cover_image_public_bucket_prefix,\n cover_image.maybe_public_bucket_extension as maybe_model_cover_image_public_bucket_extension,\n\n model_weight_creator.token as `maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken`,\n model_weight_creator.username as maybe_model_weight_creator_username,\n model_weight_creator.display_name as maybe_model_weight_creator_display_name,\n model_weight_creator.email_gravatar_hash as maybe_model_weight_creator_gravatar_hash,\n\n m.public_bucket_directory_hash,\n m.maybe_public_bucket_prefix,\n m.maybe_public_bucket_extension,\n\n entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n entity_stats.bookmark_count as maybe_bookmark_count,\n\n m.created_at,\n m.updated_at\n\nFROM media_files AS m\nLEFT OUTER JOIN users\n ON m.maybe_creator_user_token = users.token\nLEFT OUTER JOIN model_weights\n ON m.maybe_origin_model_token = model_weights.token\nLEFT OUTER JOIN media_files as cover_image\n ON cover_image.token = model_weights.maybe_cover_image_media_file_token\nLEFT OUTER JOIN users as model_weight_creator\n ON model_weight_creator.token = model_weights.creator_user_token\nLEFT OUTER JOIN entity_stats\n ON entity_stats.entity_type = \"media_file\"\n AND entity_stats.entity_token = m.token\nWHERE\n m.token = ?\n ", + "query": "\nSELECT\n m.token as `token: tokens::tokens::media_files::MediaFileToken`,\n\n m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`,\n m.maybe_media_subtype as `maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype`,\n\n users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`,\n users.username as maybe_creator_username,\n users.display_name as maybe_creator_display_name,\n users.email_gravatar_hash as maybe_creator_gravatar_hash,\n\n m.maybe_batch_token as `maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken`,\n\n m.maybe_text_transcript,\n\n m.maybe_prompt_token as `maybe_prompt_token: tokens::tokens::prompts::PromptToken`,\n\n m.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n\n model_weights.token as `maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken`,\n model_weights.title as maybe_model_weights_title,\n model_weights.weights_type as `maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType`,\n model_weights.weights_category as `maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory`,\n\n cover_image.public_bucket_directory_hash as maybe_model_cover_image_public_bucket_hash,\n cover_image.maybe_public_bucket_prefix as maybe_model_cover_image_public_bucket_prefix,\n cover_image.maybe_public_bucket_extension as maybe_model_cover_image_public_bucket_extension,\n\n model_weight_creator.token as `maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken`,\n model_weight_creator.username as maybe_model_weight_creator_username,\n model_weight_creator.display_name as maybe_model_weight_creator_display_name,\n model_weight_creator.email_gravatar_hash as maybe_model_weight_creator_gravatar_hash,\n\n m.public_bucket_directory_hash,\n m.maybe_public_bucket_prefix,\n m.maybe_public_bucket_extension,\n\n entity_stats.ratings_positive_count as maybe_ratings_positive_count,\n entity_stats.ratings_negative_count as maybe_ratings_negative_count,\n entity_stats.bookmark_count as maybe_bookmark_count,\n\n m.created_at,\n m.updated_at\n\nFROM media_files AS m\nLEFT OUTER JOIN users\n ON m.maybe_creator_user_token = users.token\nLEFT OUTER JOIN model_weights\n ON m.maybe_origin_model_token = model_weights.token\nLEFT OUTER JOIN media_files as cover_image\n ON cover_image.token = model_weights.maybe_cover_image_media_file_token\nLEFT OUTER JOIN users as model_weight_creator\n ON model_weight_creator.token = model_weights.creator_user_token\nLEFT OUTER JOIN entity_stats\n ON entity_stats.entity_type = \"media_file\"\n AND entity_stats.entity_token = m.token\nWHERE\n m.token = ?\n AND m.user_deleted_at IS NULL\n AND m.mod_deleted_at IS NULL\n ", "describe": { "columns": [ { @@ -25,6 +25,16 @@ }, { "ordinal": 2, + "name": "maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 3, "name": "maybe_creator_user_token: tokens::tokens::users::UserToken", "type_info": { "type": "VarString", @@ -34,7 +44,7 @@ } }, { - "ordinal": 3, + "ordinal": 4, "name": "maybe_creator_username", "type_info": { "type": "VarString", @@ -44,7 +54,7 @@ } }, { - "ordinal": 4, + "ordinal": 5, "name": "maybe_creator_display_name", "type_info": { "type": "VarString", @@ -54,7 +64,7 @@ } }, { - "ordinal": 5, + "ordinal": 6, "name": "maybe_creator_gravatar_hash", "type_info": { "type": "String", @@ -64,7 +74,7 @@ } }, { - "ordinal": 6, + "ordinal": 7, "name": "maybe_batch_token: tokens::tokens::batch_generations::BatchGenerationToken", "type_info": { "type": "VarString", @@ -74,7 +84,7 @@ } }, { - "ordinal": 7, + "ordinal": 8, "name": "maybe_text_transcript", "type_info": { "type": "Blob", @@ -84,7 +94,7 @@ } }, { - "ordinal": 8, + "ordinal": 9, "name": "maybe_prompt_token: tokens::tokens::prompts::PromptToken", "type_info": { "type": "VarString", @@ -94,7 +104,7 @@ } }, { - "ordinal": 9, + "ordinal": 10, "name": "creator_set_visibility: enums::common::visibility::Visibility", "type_info": { "type": "String", @@ -104,7 +114,7 @@ } }, { - "ordinal": 10, + "ordinal": 11, "name": "maybe_model_weights_token: tokens::tokens::model_weights::ModelWeightToken", "type_info": { "type": "VarString", @@ -114,7 +124,7 @@ } }, { - "ordinal": 11, + "ordinal": 12, "name": "maybe_model_weights_title", "type_info": { "type": "VarString", @@ -124,7 +134,7 @@ } }, { - "ordinal": 12, + "ordinal": 13, "name": "maybe_model_weights_type: enums::by_table::model_weights::weights_types::WeightsType", "type_info": { "type": "VarString", @@ -134,7 +144,7 @@ } }, { - "ordinal": 13, + "ordinal": 14, "name": "maybe_model_weights_category: enums::by_table::model_weights::weights_category::WeightsCategory", "type_info": { "type": "VarString", @@ -144,7 +154,7 @@ } }, { - "ordinal": 14, + "ordinal": 15, "name": "maybe_model_cover_image_public_bucket_hash", "type_info": { "type": "VarString", @@ -154,7 +164,7 @@ } }, { - "ordinal": 15, + "ordinal": 16, "name": "maybe_model_cover_image_public_bucket_prefix", "type_info": { "type": "VarString", @@ -164,7 +174,7 @@ } }, { - "ordinal": 16, + "ordinal": 17, "name": "maybe_model_cover_image_public_bucket_extension", "type_info": { "type": "VarString", @@ -174,7 +184,7 @@ } }, { - "ordinal": 17, + "ordinal": 18, "name": "maybe_model_weight_creator_user_token: tokens::tokens::users::UserToken", "type_info": { "type": "VarString", @@ -184,7 +194,7 @@ } }, { - "ordinal": 18, + "ordinal": 19, "name": "maybe_model_weight_creator_username", "type_info": { "type": "VarString", @@ -194,7 +204,7 @@ } }, { - "ordinal": 19, + "ordinal": 20, "name": "maybe_model_weight_creator_display_name", "type_info": { "type": "VarString", @@ -204,7 +214,7 @@ } }, { - "ordinal": 20, + "ordinal": 21, "name": "maybe_model_weight_creator_gravatar_hash", "type_info": { "type": "String", @@ -214,7 +224,7 @@ } }, { - "ordinal": 21, + "ordinal": 22, "name": "public_bucket_directory_hash", "type_info": { "type": "VarString", @@ -224,7 +234,7 @@ } }, { - "ordinal": 22, + "ordinal": 23, "name": "maybe_public_bucket_prefix", "type_info": { "type": "VarString", @@ -234,7 +244,7 @@ } }, { - "ordinal": 23, + "ordinal": 24, "name": "maybe_public_bucket_extension", "type_info": { "type": "VarString", @@ -244,7 +254,7 @@ } }, { - "ordinal": 24, + "ordinal": 25, "name": "maybe_ratings_positive_count", "type_info": { "type": "Long", @@ -254,7 +264,7 @@ } }, { - "ordinal": 25, + "ordinal": 26, "name": "maybe_ratings_negative_count", "type_info": { "type": "Long", @@ -264,7 +274,7 @@ } }, { - "ordinal": 26, + "ordinal": 27, "name": "maybe_bookmark_count", "type_info": { "type": "Long", @@ -274,7 +284,7 @@ } }, { - "ordinal": 27, + "ordinal": 28, "name": "created_at", "type_info": { "type": "Timestamp", @@ -284,7 +294,7 @@ } }, { - "ordinal": 28, + "ordinal": 29, "name": "updated_at", "type_info": { "type": "Timestamp", @@ -307,6 +317,7 @@ true, true, true, + true, false, true, true, @@ -329,5 +340,5 @@ false ] }, - "hash": "726a6c62dc20e8372e950a10ecd5a438769448c7ce315bc8743b728e80179420" + "hash": "b161ade93a0b6c8b7bd01ad627659903da1fa4ca6c0c9925b2a4e2bc791d4083" } diff --git a/crates/lib/mysql_queries/src/queries/media_files/get_media_file.rs b/crates/lib/mysql_queries/src/queries/media_files/get_media_file.rs index 03620cdb28..03f76d4567 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/get_media_file.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/get_media_file.rs @@ -7,6 +7,7 @@ use anyhow::anyhow; use chrono::{DateTime, Utc}; use sqlx::MySqlPool; +use enums::by_table::media_files::media_file_subtype::MediaFileSubtype; use enums::by_table::media_files::media_file_type::MediaFileType; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; @@ -23,6 +24,7 @@ pub struct MediaFile { pub token: MediaFileToken, pub media_type: MediaFileType, + pub maybe_media_subtype: Option, // TODO: Bucket hash bits. @@ -87,6 +89,7 @@ pub struct MediaFileRaw { pub token: MediaFileToken, pub media_type: MediaFileType, + pub maybe_media_subtype: Option, // TODO: Bucket hash bits. @@ -157,6 +160,7 @@ pub async fn get_media_file( Ok(Some(MediaFile { token: record.token, media_type: record.media_type, + maybe_media_subtype: record.maybe_media_subtype, maybe_batch_token: record.maybe_batch_token, maybe_text_transcript: record.maybe_text_transcript, maybe_creator_user_token: record.maybe_creator_user_token, @@ -198,6 +202,7 @@ SELECT m.token as `token: tokens::tokens::media_files::MediaFileToken`, m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`, + m.maybe_media_subtype as `maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype`, users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`, users.username as maybe_creator_username, @@ -269,6 +274,7 @@ SELECT m.token as `token: tokens::tokens::media_files::MediaFileToken`, m.media_type as `media_type: enums::by_table::media_files::media_file_type::MediaFileType`, + m.maybe_media_subtype as `maybe_media_subtype: enums::by_table::media_files::media_file_subtype::MediaFileSubtype`, users.token as `maybe_creator_user_token: tokens::tokens::users::UserToken`, users.username as maybe_creator_username, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/get_media_file.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/get_media_file.rs index 4942b14e2f..885b6a2d00 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/get_media_file.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/get_media_file.rs @@ -10,6 +10,7 @@ use log::warn; use utoipa::ToSchema; use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use enums::by_table::media_files::media_file_subtype::MediaFileSubtype; use enums::by_table::media_files::media_file_type::MediaFileType; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; @@ -45,7 +46,11 @@ pub struct MediaFileInfo { /// Type of media will dictate which fields are populated and what /// the frontend should display (eg. video player vs audio player). pub media_type: MediaFileType, - + + /// If the media file has a subtype, we'll report it. + /// This is mostly used for Bevy engine files. + pub maybe_media_subtype: Option, + /// If the file was generated as part of a batch, this is the token for the batch. pub maybe_batch_token: Option, @@ -260,6 +265,7 @@ async fn modern_media_file_lookup( media_file: MediaFileInfo { token: result.token, media_type: result.media_type, + maybe_media_subtype: result.maybe_media_subtype, maybe_batch_token: result.maybe_batch_token, public_bucket_path, maybe_text_transcript: result.maybe_text_transcript, @@ -346,6 +352,7 @@ async fn emulate_media_file_with_legacy_tts_result_lookup( media_file: MediaFileInfo { token, media_type: MediaFileType::Audio, // NB: Always audio + maybe_media_subtype: None, maybe_batch_token: None, public_bucket_path, maybe_model_weight_info: Some(GetMediaFileModelInfo { From 4147c7d9a0f11ff313a06055e5cebb9b72d3d376 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:52:37 -0500 Subject: [PATCH 19/22] add to documentation --- crates/service/web/storyteller_web/src/api_doc.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/service/web/storyteller_web/src/api_doc.rs b/crates/service/web/storyteller_web/src/api_doc.rs index d71c7dfa54..cecaf6c8cd 100644 --- a/crates/service/web/storyteller_web/src/api_doc.rs +++ b/crates/service/web/storyteller_web/src/api_doc.rs @@ -3,6 +3,7 @@ use utoipa::OpenApi; use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use enums::by_table::media_files::media_file_subtype::MediaFileSubtype; use enums::by_table::media_files::media_file_type::MediaFileType; use enums::by_table::model_weights::{weights_category::WeightsCategory, weights_types::WeightsType}; use enums::by_table::prompts::prompt_type::PromptType; @@ -113,6 +114,7 @@ use crate::http_server::web_utils::response_success_helpers::*; MediaFileOriginCategory, MediaFileOriginModelType, MediaFileOriginProductCategory, + MediaFileSubtype, MediaFileType, PromptType, WeightsCategory, From c85e774948a65b93b34da3f7eae9863eea786a83 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 03:36:58 -0500 Subject: [PATCH 20/22] migration for user feature flags --- .../down.sql | 5 +++++ .../up.sql | 7 +++++++ _sql/migrations_squashed/user_tables/users.sql | 12 ++++++++++++ 3 files changed, 24 insertions(+) create mode 100644 _sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/down.sql create mode 100644 _sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/up.sql diff --git a/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/down.sql b/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/down.sql new file mode 100644 index 0000000000..7a1e189e08 --- /dev/null +++ b/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE users DROP COLUMN maybe_feature_flags; diff --git a/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/up.sql b/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/up.sql new file mode 100644 index 0000000000..bc3b2a0abc --- /dev/null +++ b/_sql/migrations/2024-03-05-082026_alter_users_add_feature_flags/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE users + ADD COLUMN maybe_feature_flags VARCHAR(255) DEFAULT NULL + AFTER user_role_slug; diff --git a/_sql/migrations_squashed/user_tables/users.sql b/_sql/migrations_squashed/user_tables/users.sql index 0969d89844..d99f9e4d73 100644 --- a/_sql/migrations_squashed/user_tables/users.sql +++ b/_sql/migrations_squashed/user_tables/users.sql @@ -22,6 +22,18 @@ CREATE TABLE users ( -- The role assigned to the user confers permissions. user_role_slug VARCHAR(16) NOT NULL, + -- An optional comma separated list of flags for user accounts + -- If a user has a flag, then they're enabled for that given feature. + -- There are better ways to do this, but this hack will do for now. + -- This should replace `can_access_studio`. + -- + -- Possible flag values (multiple is okay) : + -- - "explore_media" + -- - "studio" + maybe_feature_flags VARCHAR(255) DEFAULT NULL, + + -- DEPRECATED + -- TODO(bt, 2024-03-05): Remove after new flags are in place. -- Rollout flag for Storyteller Studio -- Default to false until rollout is 100%, then remove the column. can_access_studio BOOLEAN NOT NULL DEFAULT false, From 4fe1e1d106e43580202565ed866d09350f0c3354 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 03:43:32 -0500 Subject: [PATCH 21/22] user feature flag enum --- .../migrations_squashed/user_tables/users.sql | 1 + crates/_types/enums/src/by_table/mod.rs | 1 + crates/_types/enums/src/by_table/users/mod.rs | 1 + .../src/by_table/users/user_feature_flag.rs | 133 ++++++++++++++++++ 4 files changed, 136 insertions(+) create mode 100644 crates/_types/enums/src/by_table/users/mod.rs create mode 100644 crates/_types/enums/src/by_table/users/user_feature_flag.rs diff --git a/_sql/migrations_squashed/user_tables/users.sql b/_sql/migrations_squashed/user_tables/users.sql index d99f9e4d73..f82eefefdd 100644 --- a/_sql/migrations_squashed/user_tables/users.sql +++ b/_sql/migrations_squashed/user_tables/users.sql @@ -30,6 +30,7 @@ CREATE TABLE users ( -- Possible flag values (multiple is okay) : -- - "explore_media" -- - "studio" + -- - "video_style_transfer" maybe_feature_flags VARCHAR(255) DEFAULT NULL, -- DEPRECATED diff --git a/crates/_types/enums/src/by_table/mod.rs b/crates/_types/enums/src/by_table/mod.rs index d4a62bd7af..96b5275c14 100644 --- a/crates/_types/enums/src/by_table/mod.rs +++ b/crates/_types/enums/src/by_table/mod.rs @@ -17,6 +17,7 @@ pub mod tts_models; pub mod usages; pub mod user_bookmarks; pub mod user_ratings; +pub mod users; pub mod voice_conversion_models; pub mod voice_conversion_results; pub mod zs_voices; diff --git a/crates/_types/enums/src/by_table/users/mod.rs b/crates/_types/enums/src/by_table/users/mod.rs new file mode 100644 index 0000000000..0027d1743a --- /dev/null +++ b/crates/_types/enums/src/by_table/users/mod.rs @@ -0,0 +1 @@ +pub mod user_feature_flag; \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/users/user_feature_flag.rs b/crates/_types/enums/src/by_table/users/user_feature_flag.rs new file mode 100644 index 0000000000..5520b303b4 --- /dev/null +++ b/crates/_types/enums/src/by_table/users/user_feature_flag.rs @@ -0,0 +1,133 @@ +use std::collections::BTreeSet; + +use serde::Deserialize; +use serde::Serialize; +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// Used in the `users` table in a `VARCHAR` field (stored as comma separated set). +/// +/// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] +#[serde(rename_all = "snake_case")] +pub enum UserFeatureFlag { + /// Grants a user the ability to list media + /// (For now, it's hidden until we get an NSFW filter.) + ExploreMedia, + + /// Access to studio features + Studio, + + /// Access to video style transfer + VideoStyleTransfer, +} + +// TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(UserFeatureFlag); +impl_mysql_enum_coders!(UserFeatureFlag); +impl_mysql_from_row!(UserFeatureFlag); + +/// NB: Legacy API for older code. +impl UserFeatureFlag { + pub fn to_str(&self) -> &'static str { + match self { + Self::ExploreMedia => "explore_media", + Self::Studio => "studio", + Self::VideoStyleTransfer => "video_style_transfer", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "explore_media" => Ok(Self::ExploreMedia), + "studio" => Ok(Self::Studio), + "video_style_transfer" => Ok(Self::VideoStyleTransfer), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::ExploreMedia, + Self::Studio, + Self::VideoStyleTransfer, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::test_helpers::assert_serialization; + use crate::by_table::users::user_feature_flag::UserFeatureFlag; + + mod explicit_checks { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(UserFeatureFlag::ExploreMedia, "explore_media"); + assert_serialization(UserFeatureFlag::Studio, "studio"); + assert_serialization(UserFeatureFlag::VideoStyleTransfer, "video_style_transfer"); + } + + #[test] + fn test_to_str() { + assert_eq!(UserFeatureFlag::ExploreMedia.to_str(), "explore_media"); + assert_eq!(UserFeatureFlag::Studio.to_str(), "studio"); + assert_eq!(UserFeatureFlag::VideoStyleTransfer.to_str(), "video_style_transfer"); + } + + #[test] + fn test_from_str() { + assert_eq!(UserFeatureFlag::from_str("explore_media").unwrap(), UserFeatureFlag::ExploreMedia); + assert_eq!(UserFeatureFlag::from_str("studio").unwrap(), UserFeatureFlag::Studio); + assert_eq!(UserFeatureFlag::from_str("video_style_transfer").unwrap(), UserFeatureFlag::VideoStyleTransfer); + assert!(UserFeatureFlag::from_str("foo").is_err()); + } + + #[test] + fn all_variants() { + let mut variants = UserFeatureFlag::all_variants(); + assert_eq!(variants.len(), 3); + assert_eq!(variants.pop_first(), Some(UserFeatureFlag::ExploreMedia)); + assert_eq!(variants.pop_first(), Some(UserFeatureFlag::Studio)); + assert_eq!(variants.pop_first(), Some(UserFeatureFlag::VideoStyleTransfer)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(UserFeatureFlag::all_variants().len(), UserFeatureFlag::iter().len()); + } + + #[test] + fn round_trip() { + for variant in UserFeatureFlag::all_variants() { + assert_eq!(variant, UserFeatureFlag::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, UserFeatureFlag::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, UserFeatureFlag::from_str(&format!("{:?}", variant)).unwrap()); + } + } + + #[test] + fn serialized_length_ok_for_database() { + const MAX_LENGTH : usize = 32; + for variant in UserFeatureFlag::all_variants() { + let serialized = variant.to_str(); + assert!(serialized.len() > 0, "variant {:?} is too short", variant); + assert!(serialized.len() <= MAX_LENGTH, "variant {:?} is too long", variant); + } + } + } +} From 58d9a6826d7bfd515503bfd26e6fb1eca3ca3c9d Mon Sep 17 00:00:00 2001 From: salt Date: Thu, 7 Mar 2024 13:59:59 -0800 Subject: [PATCH 22/22] comfyui sidecar --- crates/service/job/inference_job/Cargo.toml | 2 +- .../config/inference-job.development.env | 4 +- .../comfy_ui/comfy_ui_dependencies.rs | 70 ++++++++++++++----- .../comfy_ui/comfy_ui_inference_command.rs | 6 ++ .../src/job_specific_dependencies.rs | 4 +- crates/service/job/inference_job/src/main.rs | 2 +- .../inference_job/src/util/instrumentation.rs | 5 ++ 7 files changed, 71 insertions(+), 22 deletions(-) diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 1c20087956..f8cbc47307 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -85,7 +85,7 @@ mime = "0.3.17" rand = { version = "0.8.5", features = [] } jsonpath_lib = "0.3.0" walkdir = "2.4.0" -reqwest = "0.11.22" +reqwest = { version = "0.11.22", features = ["blocking"] } opentelemetry = {version = "0.21.0", features = ["metrics"]} opentelemetry-otlp = { version = "0.14.0", features = ["metrics"] } opentelemetry_sdk = { version = "0.21.2", features = ["metrics", "rt-tokio"] } diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index 957a4f5132..88dac0d340 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -147,7 +147,8 @@ MOCAPNET_INFERENCE_EXECUTABLE_OR_COMMAND="pythonVirtualEnvironment/bin/python -m # # # ================================= # COMFY_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/" -COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunner.py" +COMFY_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python ../ComfyLauncher/ComfyRunnerForSidecar.py" +COMFY_LAUNCH_COMMAND="venv/bin/python /home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/main.py" # directory to where files should be downloaded to (default to just the comfy directory) COMFY_MOUNTS_DIRECTORY="/home/salt/clone/storyteller/storyteller-ml/workflows/comfy/ComfyUI/" @@ -187,3 +188,4 @@ HF_DATASETS_CACHE="/tmp/huggingface" # nltk also caches data. NLTK_DATA="/tmp/nltk_data" +SKIP_TELEMETRY="true" \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs index 906cdbf82d..f2f527a2b7 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_dependencies.rs @@ -1,6 +1,10 @@ use std::path::PathBuf; -use std::process::Command; +use tokio::process::Command; +use std::time::{Duration, Instant}; +use actix_web::http::StatusCode; use errors::AnyhowResult; +use reqwest::blocking::Client; +use tokio::time::sleep; use crate::job::job_types::workflow::comfy_ui::comfy_ui_inference_command::ComfyInferenceCommand; use crate::util::common_commands::ffmpeg_logo_watermark_command::FfmpegLogoWatermarkCommand; @@ -12,24 +16,56 @@ pub struct ComfyDependencies { } impl ComfyDependencies { - pub fn setup() -> AnyhowResult { - // run the setup scripts for the comfy dependencies + pub async fn setup() -> AnyhowResult { let inference_command = ComfyInferenceCommand::from_env()?; - // check to see if the flag file exists - if PathBuf::from("comfy_setup_complete").exists() { - println!("Comfy setup already complete"); - } else { - // if the flag file does not exist, run the setup script - // run setup script - let comfy_setup_script = inference_command.clone().comfy_setup_script; - let output = Command::new("python3") - .arg(comfy_setup_script) - .output() - .expect("failed to run comfyui setup script"); - println!("stdout: {}", String::from_utf8_lossy(&output.stdout)); - } + let comfy_setup_script = inference_command.clone().comfy_setup_script; + let output = Command::new("python3") + .arg(&comfy_setup_script) + .output() + .await.expect("failed to run comfyui setup script"); + println!("stdout: {}", String::from_utf8_lossy(&output.stdout)); + + // Start the comfy start script as a daemon process + let launch_command = format!( + "cd {} && {}", + inference_command.clone().comfy_root_code_directory.display(), + inference_command.clone().comfy_launch_command.display() + ); + + let _ = Command::new("sh") + .arg("-c") + .arg(&launch_command) + .stdout(std::process::Stdio::inherit()) // Inherit the stdout to stream the output + .spawn() + .expect("failed to start comfy start script"); + + let mut success = false; + + let client = reqwest::Client::new(); + let start = Instant::now(); + let timeout = Duration::from_secs(60); + while Instant::now().duration_since(start) < timeout { + let response = client.get("http://127.0.0.1:8188/prompt").send().await; + match response { + Ok(r) if r.status() == StatusCode::OK => { + println!("Received 200 response from http://127.0.0.1:8188/prompt"); + success = true; + break; + }, + _ => { + println!("Did not receive 200 response, retrying..."); + sleep(Duration::from_secs(5)).await; + } + } + } + + if !success { + println!("Timeout reached without receiving a 200 response."); + panic!("Comfy start failed"); + } + Ok(Self { - inference_command: inference_command, + inference_command, dependency_tokens: RequiredModels::init(), ffmpeg_watermark_command: FfmpegLogoWatermarkCommand::from_env()?, }) diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs index a4e66177c7..a0ec7c627f 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/comfy_ui_inference_command.rs @@ -46,6 +46,8 @@ pub struct ComfyInferenceCommand { pub(crate) comfy_setup_script: PathBuf, + pub(crate) comfy_launch_command: PathBuf, + /// Config file to use config_path: Option, @@ -100,6 +102,9 @@ impl ComfyInferenceCommand { let comfy_setup_script = easyenv::get_env_pathbuf_required( "COMFY_SETUP_SCRIPT")?; + let comfy_launch_command = easyenv::get_env_pathbuf_required( + "COMFY_LAUNCH_COMMAND")?; + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( "COMFY_INFERENCE_MAYBE_VENV_COMMAND"); @@ -130,6 +135,7 @@ impl ComfyInferenceCommand { mounts_directory, processing_script, comfy_setup_script, + comfy_launch_command, maybe_virtual_env_activation_command, maybe_docker_options, maybe_execution_timeout, diff --git a/crates/service/job/inference_job/src/job_specific_dependencies.rs b/crates/service/job/inference_job/src/job_specific_dependencies.rs index 758110d284..7d33b29665 100644 --- a/crates/service/job/inference_job/src/job_specific_dependencies.rs +++ b/crates/service/job/inference_job/src/job_specific_dependencies.rs @@ -34,7 +34,7 @@ pub struct JobSpecificDependencies { impl JobSpecificDependencies { - pub fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult { + pub async fn setup_for_jobs(scoped_execution: &ScopedExecution) -> AnyhowResult { let mut maybe_rvc_v2_dependencies = None; let mut maybe_sad_talker_dependencies = None; let mut maybe_svc_dependencies = None; @@ -101,7 +101,7 @@ impl JobSpecificDependencies { if scoped_execution.can_run_job(InferenceModelType::ComfyUi) { print_with_space("Setting ComfyUI dependencies..."); - maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup()?); + maybe_comfy_ui_dependencies = Some(ComfyDependencies::setup().await?); } if scoped_execution.can_run_job(InferenceModelType::ConvertFbxToGltf) { diff --git a/crates/service/job/inference_job/src/main.rs b/crates/service/job/inference_job/src/main.rs index f418d7f387..72d091154b 100644 --- a/crates/service/job/inference_job/src/main.rs +++ b/crates/service/job/inference_job/src/main.rs @@ -237,7 +237,7 @@ async fn main() -> AnyhowResult<()> { let meter = opentelemetry::global::meter("inference-job"); - let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution)?; + let job_specific_dependencies = JobSpecificDependencies::setup_for_jobs(&scoped_execution).await?; let scoped_tempdir_for_downloads = ScopedTempDirCreator::for_directory( easyenv::get_env_pathbuf_or_default( diff --git a/crates/service/job/inference_job/src/util/instrumentation.rs b/crates/service/job/inference_job/src/util/instrumentation.rs index d269a4d696..ee91e1b16a 100644 --- a/crates/service/job/inference_job/src/util/instrumentation.rs +++ b/crates/service/job/inference_job/src/util/instrumentation.rs @@ -21,6 +21,11 @@ pub fn init_otel_metrics_pipeline( ) -> Result<(), opentelemetry::metrics::MetricsError> { info!("Setting up otel metrics pipeline..."); + // check env var to see if we should skip telemetry + if std::env::var("SKIP_TELEMETRY").is_ok() { + warn!("Skipping telemetry setup due to SKIP_TELEMETRY env var"); + return Ok(()); + } let provider = opentelemetry_otlp::new_pipeline() .metrics(opentelemetry_sdk::runtime::Tokio) // TODO: 1. read host from env 2. Single pod of otel-collector is probably not good enough, run daemonset?