From e428f0ae9d0041657eea9575fcc159b1920ffe05 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 19:34:29 -0500 Subject: [PATCH 01/32] 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 b45a80142a9fbab7671b466abbe400b58f1bda61 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 19:46:05 -0500 Subject: [PATCH 02/32] 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 ea4887a41c1655e613d13cf4be01f2659ca1a126 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 23:06:56 -0500 Subject: [PATCH 03/32] 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 e023a27c181e366b6c48164f0808f323a9ec7dc9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 29 Feb 2024 23:19:55 -0500 Subject: [PATCH 04/32] 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 05/32] 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 f5999aef88..b5f187a02e 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 @@ -198,10 +198,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 23f7f0a539097c86b5518dd6a7578c523971523c Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 02:43:37 -0500 Subject: [PATCH 06/32] 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 2972d5407792ca9c852a0eb94ea65c4ac896a87d Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:14:24 -0500 Subject: [PATCH 07/32] 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 26f37629cae75124db779caa5e221b1f004651e3 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:23:08 -0500 Subject: [PATCH 08/32] 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 51dfddf78927d278f19da4be4b58d0ed2a6bdedd Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 03:31:21 -0500 Subject: [PATCH 09/32] 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 10/32] 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 11/32] 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 9fd9cf5cc93ed345adea687b5742b8be96c88e14 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 1 Mar 2024 05:25:29 -0500 Subject: [PATCH 12/32] add watermark to comfy --- .../comfy_ui/comfy_ui_dependencies.rs | 4 ++ .../workflow/comfy_ui/process_job.rs | 70 +++++++++++++++---- 2 files changed, 62 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 fee7faac54..5e80cd9c99 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,10 +1,13 @@ use std::path::PathBuf; + 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, } @@ -13,6 +16,7 @@ impl ComfyDependencies { Ok(Self { inference_command: ComfyInferenceCommand::from_env()?, 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 13/32] 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 e43580d6272a0fd40c51e8ce7babaa0aab6fa122 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Sun, 3 Mar 2024 22:43:10 -0500 Subject: [PATCH 14/32] 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 da4d7bea308f616d22483db77a1df3857adb2cdb Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:32:00 -0500 Subject: [PATCH 15/32] 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 16/32] 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 40d88f9fdd986add9ca31b29d55504190f324b71 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:51:38 -0500 Subject: [PATCH 17/32] 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 97e7c4a3cd47695ff90b2c565a3a4efac568a3da Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 4 Mar 2024 03:52:37 -0500 Subject: [PATCH 18/32] 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 a254aa451980c42c2fbf874773e1c1ffb3afef81 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 03:36:58 -0500 Subject: [PATCH 19/32] 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 fdd44d3ecdc10f201d86ea497eac26178ccbf721 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 03:43:32 -0500 Subject: [PATCH 20/32] 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 b6c0e7738e8456eb674239a36917693edc1f4b8f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 14:26:03 -0500 Subject: [PATCH 21/32] read new permission flags --- .../get_user_session_by_token.rs | 22 +++++++++++++++++ .../users/src/utils/session_checker.rs | 6 +++-- .../users/src/utils/user_session_extended.rs | 24 +++++++++++++++++++ 3 files changed, 50 insertions(+), 2 deletions(-) diff --git a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs index d51ad37a66..1e59488964 100644 --- a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs @@ -2,11 +2,14 @@ #![forbid(unused_imports)] #![forbid(unused_mut)] +use std::collections::HashSet; + use anyhow::anyhow; use log::warn; use sqlx::{Executor, MySql}; use sqlx::pool::PoolConnection; +use enums::by_table::users::user_feature_flag::UserFeatureFlag; use enums::common::visibility::Visibility; use errors::AnyhowResult; use tokens::tokens::users::UserToken; @@ -39,6 +42,9 @@ pub struct SessionUserRecord { // ===== FEATURE FLAGS ===== // + // Optional comma-separated list of parseable `UserFeatureFlag` enum features + pub maybe_feature_flags: Option, + pub can_access_studio: bool, // ===== ROLE ===== // @@ -79,6 +85,18 @@ impl SessionUserRecord { pub fn get_strongly_typed_user_token(&self) -> UserToken { UserToken::new_from_str(&self.user_token) } + + pub fn get_flags(&self) -> HashSet { + match self.maybe_feature_flags.as_deref() { + None => HashSet::new(), + Some(feature_flags) => { + feature_flags + .split(",") + .filter_map(|flag| UserFeatureFlag::from_str(flag).ok()) + .collect() + } + } + } } pub async fn get_user_session_by_token_pooled_connection( @@ -121,6 +139,7 @@ SELECT users.is_banned, users.can_access_studio, + users.maybe_feature_flags, user_roles.can_use_tts, user_roles.can_use_w2l, @@ -183,6 +202,7 @@ WHERE user_sessions.token = ? is_banned: i8_to_bool(raw_user_record.is_banned), can_access_studio: i8_to_bool(raw_user_record.can_access_studio), + maybe_feature_flags: raw_user_record.maybe_feature_flags, // Usage can_use_tts: nullable_i8_to_bool_default_false(raw_user_record.can_use_tts), @@ -250,6 +270,8 @@ struct SessionUserRawDbRecord { // Feature / Rollout Flags can_access_studio: i8, + maybe_feature_flags: Option, + // NB: These are `Option` due to the JOIN not being compile-time assured. // Usage can_use_tts: Option, diff --git a/crates/service/plugins/users/src/utils/session_checker.rs b/crates/service/plugins/users/src/utils/session_checker.rs index b9bcab82b7..6da3796858 100644 --- a/crates/service/plugins/users/src/utils/session_checker.rs +++ b/crates/service/plugins/users/src/utils/session_checker.rs @@ -16,7 +16,7 @@ use redis_caching::redis_ttl_cache::{RedisTtlCache, RedisTtlCacheConnection}; use redis_common::redis_cache_keys::RedisCacheKeys; use crate::cookies::session::session_cookie_manager::SessionCookieManager; -use crate::utils::user_session_extended::{UserSessionExtended, UserSessionPreferences, UserSessionPremiumPlanInfo, UserSessionRoleAndPermissions, UserSessionSubscriptionPlan, UserSessionUserDetails}; +use crate::utils::user_session_extended::{UserSessionExtended, UserSessionFeatureFlags, UserSessionPreferences, UserSessionPremiumPlanInfo, UserSessionRoleAndPermissions, UserSessionSubscriptionPlan, UserSessionUserDetails}; #[derive(Clone)] pub struct SessionChecker { @@ -168,7 +168,6 @@ impl SessionChecker { } - // ==================== UserSessionExtended ==================== //#[deprecated = "Use the PoolConnection method instead of the MySqlPool one."] @@ -258,6 +257,9 @@ impl SessionChecker { can_ban_users: user_session.can_ban_users, can_delete_users: user_session.can_delete_users, }, + feature_flags: UserSessionFeatureFlags { + maybe_feature_flags: user_session.maybe_feature_flags, + } })) } diff --git a/crates/service/plugins/users/src/utils/user_session_extended.rs b/crates/service/plugins/users/src/utils/user_session_extended.rs index fe8f2bb2f6..5fa0b2a34c 100644 --- a/crates/service/plugins/users/src/utils/user_session_extended.rs +++ b/crates/service/plugins/users/src/utils/user_session_extended.rs @@ -1,5 +1,8 @@ +use std::collections::HashSet; + use chrono::{DateTime, Utc}; +use enums::by_table::users::user_feature_flag::UserFeatureFlag; use enums::common::visibility::Visibility; #[derive(Clone, Default)] @@ -9,6 +12,7 @@ pub struct UserSessionExtended { pub premium: UserSessionPremiumPlanInfo, pub preferences: UserSessionPreferences, pub role: UserSessionRoleAndPermissions, + pub feature_flags: UserSessionFeatureFlags, } #[derive(Clone, Default)] @@ -96,3 +100,23 @@ pub struct UserSessionRoleAndPermissions { pub can_ban_users: bool, pub can_delete_users: bool, } + +#[derive(Default, Clone)] +pub struct UserSessionFeatureFlags { + // Optional comma-separated list of parseable `UserFeatureFlag` enum features + pub maybe_feature_flags: Option, +} + +impl UserSessionFeatureFlags { + pub fn get_flags(&self) -> HashSet { + match self.maybe_feature_flags.as_deref() { + None => HashSet::new(), + Some(feature_flags) => { + feature_flags + .split(",") + .filter_map(|flag| UserFeatureFlag::from_str(flag).ok()) + .collect() + } + } + } +} From 1bd5ceb9e672b8a0a81a65d9f454ddd47e96b5fe Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Tue, 5 Mar 2024 14:45:35 -0500 Subject: [PATCH 22/32] new permission accessors wired up --- .../get_user_session_by_token.rs | 15 --- crates/service/plugins/users/src/utils/mod.rs | 3 +- .../users/src/utils/session_checker.rs | 3 +- .../users/src/utils/user_session_extended.rs | 24 +--- .../src/utils/user_session_feature_flags.rs | 110 ++++++++++++++++++ .../src/util/allowed_studio_access.rs | 30 ++++- 6 files changed, 142 insertions(+), 43 deletions(-) create mode 100644 crates/service/plugins/users/src/utils/user_session_feature_flags.rs diff --git a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs index 1e59488964..425857e8b8 100644 --- a/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_sessions/get_user_session_by_token.rs @@ -2,14 +2,11 @@ #![forbid(unused_imports)] #![forbid(unused_mut)] -use std::collections::HashSet; - use anyhow::anyhow; use log::warn; use sqlx::{Executor, MySql}; use sqlx::pool::PoolConnection; -use enums::by_table::users::user_feature_flag::UserFeatureFlag; use enums::common::visibility::Visibility; use errors::AnyhowResult; use tokens::tokens::users::UserToken; @@ -85,18 +82,6 @@ impl SessionUserRecord { pub fn get_strongly_typed_user_token(&self) -> UserToken { UserToken::new_from_str(&self.user_token) } - - pub fn get_flags(&self) -> HashSet { - match self.maybe_feature_flags.as_deref() { - None => HashSet::new(), - Some(feature_flags) => { - feature_flags - .split(",") - .filter_map(|flag| UserFeatureFlag::from_str(flag).ok()) - .collect() - } - } - } } pub async fn get_user_session_by_token_pooled_connection( diff --git a/crates/service/plugins/users/src/utils/mod.rs b/crates/service/plugins/users/src/utils/mod.rs index d49f148287..c22814e9d2 100644 --- a/crates/service/plugins/users/src/utils/mod.rs +++ b/crates/service/plugins/users/src/utils/mod.rs @@ -1,3 +1,4 @@ pub mod email_to_gravatar; pub mod session_checker; -pub mod user_session_extended; \ No newline at end of file +pub mod user_session_extended; +pub mod user_session_feature_flags; \ No newline at end of file diff --git a/crates/service/plugins/users/src/utils/session_checker.rs b/crates/service/plugins/users/src/utils/session_checker.rs index 6da3796858..d5ece432e1 100644 --- a/crates/service/plugins/users/src/utils/session_checker.rs +++ b/crates/service/plugins/users/src/utils/session_checker.rs @@ -16,7 +16,8 @@ use redis_caching::redis_ttl_cache::{RedisTtlCache, RedisTtlCacheConnection}; use redis_common::redis_cache_keys::RedisCacheKeys; use crate::cookies::session::session_cookie_manager::SessionCookieManager; -use crate::utils::user_session_extended::{UserSessionExtended, UserSessionFeatureFlags, UserSessionPreferences, UserSessionPremiumPlanInfo, UserSessionRoleAndPermissions, UserSessionSubscriptionPlan, UserSessionUserDetails}; +use crate::utils::user_session_extended::{UserSessionExtended, UserSessionPreferences, UserSessionPremiumPlanInfo, UserSessionRoleAndPermissions, UserSessionSubscriptionPlan, UserSessionUserDetails}; +use crate::utils::user_session_feature_flags::UserSessionFeatureFlags; #[derive(Clone)] pub struct SessionChecker { diff --git a/crates/service/plugins/users/src/utils/user_session_extended.rs b/crates/service/plugins/users/src/utils/user_session_extended.rs index 5fa0b2a34c..00e655e25a 100644 --- a/crates/service/plugins/users/src/utils/user_session_extended.rs +++ b/crates/service/plugins/users/src/utils/user_session_extended.rs @@ -1,10 +1,9 @@ -use std::collections::HashSet; - use chrono::{DateTime, Utc}; -use enums::by_table::users::user_feature_flag::UserFeatureFlag; use enums::common::visibility::Visibility; +use crate::utils::user_session_feature_flags::UserSessionFeatureFlags; + #[derive(Clone, Default)] pub struct UserSessionExtended { pub user_token: String, @@ -101,22 +100,3 @@ pub struct UserSessionRoleAndPermissions { pub can_delete_users: bool, } -#[derive(Default, Clone)] -pub struct UserSessionFeatureFlags { - // Optional comma-separated list of parseable `UserFeatureFlag` enum features - pub maybe_feature_flags: Option, -} - -impl UserSessionFeatureFlags { - pub fn get_flags(&self) -> HashSet { - match self.maybe_feature_flags.as_deref() { - None => HashSet::new(), - Some(feature_flags) => { - feature_flags - .split(",") - .filter_map(|flag| UserFeatureFlag::from_str(flag).ok()) - .collect() - } - } - } -} diff --git a/crates/service/plugins/users/src/utils/user_session_feature_flags.rs b/crates/service/plugins/users/src/utils/user_session_feature_flags.rs new file mode 100644 index 0000000000..80298d944e --- /dev/null +++ b/crates/service/plugins/users/src/utils/user_session_feature_flags.rs @@ -0,0 +1,110 @@ + +use std::collections::HashSet; +use enums::by_table::users::user_feature_flag::UserFeatureFlag; + +#[derive(Default, Clone)] +pub struct UserSessionFeatureFlags { + // Optional comma-separated list of parseable `UserFeatureFlag` enum features + pub maybe_feature_flags: Option, +} + +impl UserSessionFeatureFlags { + pub fn get_flags(&self) -> HashSet { + match self.maybe_feature_flags.as_deref() { + None => HashSet::new(), + Some(feature_flags) => { + feature_flags + .split(",") + .map(|flag| flag.trim()) + .filter_map(|flag| UserFeatureFlag::from_str(flag).ok()) + .collect() + } + } + } + + // TODO(bt,2024-03-05): Caching + pub fn has_permission_unoptimized(&self, permission: UserFeatureFlag) -> bool { + self.get_flags().contains(&permission) + } +} + +#[cfg(test)] +mod tests { + use enums::by_table::users::user_feature_flag::UserFeatureFlag; + use crate::utils::user_session_feature_flags::UserSessionFeatureFlags; + + #[test] + fn test_no_flags() { + let flags = UserSessionFeatureFlags::default(); + assert_eq!(flags.get_flags().len(), 0); + + for flag in UserFeatureFlag::all_variants() { + assert_eq!(flags.has_permission_unoptimized(flag), false); + } + + // TODO(bt,2024-03-05): Expose strum to callers in test packages. + //for flag in UserFeatureFlag::iter() { + // assert_eq!(flags.has_permission_unoptimized(flag), false); + //} + } + + #[test] + fn test_single_feature() { + let flags = UserSessionFeatureFlags { + maybe_feature_flags: Some("studio".to_string()), + }; + + assert_eq!(flags.get_flags().len(), 1); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::Studio), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer), false); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia), false); + } + + #[test] + fn test_all_features() { + let flags = UserSessionFeatureFlags { + maybe_feature_flags: Some("studio,video_style_transfer,explore_media".to_string()), + }; + + assert_eq!(flags.get_flags().len(), 3); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::Studio), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia), true); + } + + #[test] + fn test_duplication() { + let flags = UserSessionFeatureFlags { + maybe_feature_flags: Some("studio,studio,studio,studio".to_string()), + }; + + assert_eq!(flags.get_flags().len(), 1); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::Studio), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer), false); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia), false); + } + + #[test] + fn test_spacing() { + let flags = UserSessionFeatureFlags { + maybe_feature_flags: Some(" studio, video_style_transfer , , , ".to_string()), + }; + + assert_eq!(flags.get_flags().len(), 2); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::Studio), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer), true); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia), false); + } + + #[test] + fn test_invalid_features_and_typos() { + let flags = UserSessionFeatureFlags { + maybe_feature_flags: Some("invalid,,typo,stdo,STUDIO".to_string()), + }; + + assert_eq!(flags.get_flags().len(), 0); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::Studio), false); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer), false); + assert_eq!(flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia), false); + } +} \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/util/allowed_studio_access.rs b/crates/service/web/storyteller_web/src/util/allowed_studio_access.rs index d30449d272..47b62e6bac 100644 --- a/crates/service/web/storyteller_web/src/util/allowed_studio_access.rs +++ b/crates/service/web/storyteller_web/src/util/allowed_studio_access.rs @@ -1,5 +1,7 @@ +use enums::by_table::users::user_feature_flag::UserFeatureFlag; use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; use users_component::utils::user_session_extended::UserSessionExtended; +use users_component::utils::user_session_feature_flags::UserSessionFeatureFlags; use crate::server_state::StaticFeatureFlags; @@ -20,24 +22,44 @@ pub trait UserSessionStudioFlag { impl UserSessionStudioFlag for UserSessionExtended { fn can_access_studio(&self) -> bool { - self.role.can_access_studio + self.role.can_access_studio || self.feature_flags.has_permission_unoptimized(UserFeatureFlag::Studio) } } impl UserSessionStudioFlag for &UserSessionExtended { fn can_access_studio(&self) -> bool { - self.role.can_access_studio + self.role.can_access_studio || self.feature_flags.has_permission_unoptimized(UserFeatureFlag::Studio) } } impl UserSessionStudioFlag for SessionUserRecord { fn can_access_studio(&self) -> bool { - self.can_access_studio + if self.can_access_studio { + return true + } + + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: self.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::Studio) } } impl UserSessionStudioFlag for &SessionUserRecord { fn can_access_studio(&self) -> bool { - self.can_access_studio + if self.can_access_studio { + return true + } + + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: self.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::Studio) } } From 979c12e5f88b11efd7e2ac1e5e10616636c798a6 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 02:24:47 -0500 Subject: [PATCH 23/32] codegen db --- ...2a1db4ed4552d53cb4f31ff4b08de01691f7.json} | 51 +++++++++++-------- 1 file changed, 31 insertions(+), 20 deletions(-) rename .sqlx/{query-423109f2dcee23a558c43ecaab0eb2836714d5316620daaa1127b447cad59212.json => query-1c32bdb79ed2bc1b8988bbb6ea3f2a1db4ed4552d53cb4f31ff4b08de01691f7.json} (85%) diff --git a/.sqlx/query-423109f2dcee23a558c43ecaab0eb2836714d5316620daaa1127b447cad59212.json b/.sqlx/query-1c32bdb79ed2bc1b8988bbb6ea3f2a1db4ed4552d53cb4f31ff4b08de01691f7.json similarity index 85% rename from .sqlx/query-423109f2dcee23a558c43ecaab0eb2836714d5316620daaa1127b447cad59212.json rename to .sqlx/query-1c32bdb79ed2bc1b8988bbb6ea3f2a1db4ed4552d53cb4f31ff4b08de01691f7.json index c8eac2bc0a..fb0c810712 100644 --- a/.sqlx/query-423109f2dcee23a558c43ecaab0eb2836714d5316620daaa1127b447cad59212.json +++ b/.sqlx/query-1c32bdb79ed2bc1b8988bbb6ea3f2a1db4ed4552d53cb4f31ff4b08de01691f7.json @@ -1,6 +1,6 @@ { "db_name": "MySQL", - "query": "\nSELECT\n users.token as user_token,\n users.username,\n users.display_name,\n\n users.email_address,\n users.email_confirmed,\n users.email_gravatar_hash,\n\n users.maybe_stripe_customer_id,\n users.maybe_loyalty_program_key,\n\n users.disable_gravatar,\n users.auto_play_audio_preference,\n users.auto_play_video_preference,\n users.preferred_tts_result_visibility as `preferred_tts_result_visibility: enums::common::visibility::Visibility`,\n users.preferred_w2l_result_visibility as `preferred_w2l_result_visibility: enums::common::visibility::Visibility`,\n\n users.user_role_slug,\n users.is_banned,\n\n users.can_access_studio,\n\n user_roles.can_use_tts,\n user_roles.can_use_w2l,\n user_roles.can_delete_own_tts_results,\n user_roles.can_delete_own_w2l_results,\n user_roles.can_delete_own_account,\n\n user_roles.can_upload_tts_models,\n user_roles.can_upload_w2l_templates,\n user_roles.can_delete_own_tts_models,\n user_roles.can_delete_own_w2l_templates,\n\n user_roles.can_approve_w2l_templates,\n user_roles.can_edit_other_users_profiles,\n user_roles.can_edit_other_users_tts_models,\n user_roles.can_edit_other_users_w2l_templates,\n user_roles.can_delete_other_users_tts_models,\n user_roles.can_delete_other_users_tts_results,\n user_roles.can_delete_other_users_w2l_templates,\n user_roles.can_delete_other_users_w2l_results,\n user_roles.can_ban_users,\n user_roles.can_delete_users\n\nFROM users\nLEFT OUTER JOIN user_sessions\n ON users.token = user_sessions.user_token\nLEFT OUTER JOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE user_sessions.token = ?\n AND user_sessions.deleted_at IS NULL\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\n ", + "query": "\nSELECT\n users.token as user_token,\n users.username,\n users.display_name,\n\n users.email_address,\n users.email_confirmed,\n users.email_gravatar_hash,\n\n users.maybe_stripe_customer_id,\n users.maybe_loyalty_program_key,\n\n users.disable_gravatar,\n users.auto_play_audio_preference,\n users.auto_play_video_preference,\n users.preferred_tts_result_visibility as `preferred_tts_result_visibility: enums::common::visibility::Visibility`,\n users.preferred_w2l_result_visibility as `preferred_w2l_result_visibility: enums::common::visibility::Visibility`,\n\n users.user_role_slug,\n users.is_banned,\n\n users.can_access_studio,\n users.maybe_feature_flags,\n\n user_roles.can_use_tts,\n user_roles.can_use_w2l,\n user_roles.can_delete_own_tts_results,\n user_roles.can_delete_own_w2l_results,\n user_roles.can_delete_own_account,\n\n user_roles.can_upload_tts_models,\n user_roles.can_upload_w2l_templates,\n user_roles.can_delete_own_tts_models,\n user_roles.can_delete_own_w2l_templates,\n\n user_roles.can_approve_w2l_templates,\n user_roles.can_edit_other_users_profiles,\n user_roles.can_edit_other_users_tts_models,\n user_roles.can_edit_other_users_w2l_templates,\n user_roles.can_delete_other_users_tts_models,\n user_roles.can_delete_other_users_tts_results,\n user_roles.can_delete_other_users_w2l_templates,\n user_roles.can_delete_other_users_w2l_results,\n user_roles.can_ban_users,\n user_roles.can_delete_users\n\nFROM users\nLEFT OUTER JOIN user_sessions\n ON users.token = user_sessions.user_token\nLEFT OUTER JOIN user_roles\n ON users.user_role_slug = user_roles.slug\nWHERE user_sessions.token = ?\n AND user_sessions.deleted_at IS NULL\n AND users.user_deleted_at IS NULL\n AND users.mod_deleted_at IS NULL\n ", "describe": { "columns": [ { @@ -165,6 +165,16 @@ }, { "ordinal": 16, + "name": "maybe_feature_flags", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 17, "name": "can_use_tts", "type_info": { "type": "Tiny", @@ -174,7 +184,7 @@ } }, { - "ordinal": 17, + "ordinal": 18, "name": "can_use_w2l", "type_info": { "type": "Tiny", @@ -184,7 +194,7 @@ } }, { - "ordinal": 18, + "ordinal": 19, "name": "can_delete_own_tts_results", "type_info": { "type": "Tiny", @@ -194,7 +204,7 @@ } }, { - "ordinal": 19, + "ordinal": 20, "name": "can_delete_own_w2l_results", "type_info": { "type": "Tiny", @@ -204,7 +214,7 @@ } }, { - "ordinal": 20, + "ordinal": 21, "name": "can_delete_own_account", "type_info": { "type": "Tiny", @@ -214,7 +224,7 @@ } }, { - "ordinal": 21, + "ordinal": 22, "name": "can_upload_tts_models", "type_info": { "type": "Tiny", @@ -224,7 +234,7 @@ } }, { - "ordinal": 22, + "ordinal": 23, "name": "can_upload_w2l_templates", "type_info": { "type": "Tiny", @@ -234,7 +244,7 @@ } }, { - "ordinal": 23, + "ordinal": 24, "name": "can_delete_own_tts_models", "type_info": { "type": "Tiny", @@ -244,7 +254,7 @@ } }, { - "ordinal": 24, + "ordinal": 25, "name": "can_delete_own_w2l_templates", "type_info": { "type": "Tiny", @@ -254,7 +264,7 @@ } }, { - "ordinal": 25, + "ordinal": 26, "name": "can_approve_w2l_templates", "type_info": { "type": "Tiny", @@ -264,7 +274,7 @@ } }, { - "ordinal": 26, + "ordinal": 27, "name": "can_edit_other_users_profiles", "type_info": { "type": "Tiny", @@ -274,7 +284,7 @@ } }, { - "ordinal": 27, + "ordinal": 28, "name": "can_edit_other_users_tts_models", "type_info": { "type": "Tiny", @@ -284,7 +294,7 @@ } }, { - "ordinal": 28, + "ordinal": 29, "name": "can_edit_other_users_w2l_templates", "type_info": { "type": "Tiny", @@ -294,7 +304,7 @@ } }, { - "ordinal": 29, + "ordinal": 30, "name": "can_delete_other_users_tts_models", "type_info": { "type": "Tiny", @@ -304,7 +314,7 @@ } }, { - "ordinal": 30, + "ordinal": 31, "name": "can_delete_other_users_tts_results", "type_info": { "type": "Tiny", @@ -314,7 +324,7 @@ } }, { - "ordinal": 31, + "ordinal": 32, "name": "can_delete_other_users_w2l_templates", "type_info": { "type": "Tiny", @@ -324,7 +334,7 @@ } }, { - "ordinal": 32, + "ordinal": 33, "name": "can_delete_other_users_w2l_results", "type_info": { "type": "Tiny", @@ -334,7 +344,7 @@ } }, { - "ordinal": 33, + "ordinal": 34, "name": "can_ban_users", "type_info": { "type": "Tiny", @@ -344,7 +354,7 @@ } }, { - "ordinal": 34, + "ordinal": 35, "name": "can_delete_users", "type_info": { "type": "Tiny", @@ -392,8 +402,9 @@ true, true, true, + true, true ] }, - "hash": "423109f2dcee23a558c43ecaab0eb2836714d5316620daaa1127b447cad59212" + "hash": "1c32bdb79ed2bc1b8988bbb6ea3f2a1db4ed4552d53cb4f31ff4b08de01691f7" } From 0c6d10a5e44cf0d83c42e8707e7be89bccd78633 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 02:42:34 -0500 Subject: [PATCH 24/32] add feature flags to session info --- .../service/plugins/users/src/endpoints/session_info_handler.rs | 2 ++ 1 file changed, 2 insertions(+) diff --git a/crates/service/plugins/users/src/endpoints/session_info_handler.rs b/crates/service/plugins/users/src/endpoints/session_info_handler.rs index 0ee0a5dc49..090e969eb5 100644 --- a/crates/service/plugins/users/src/endpoints/session_info_handler.rs +++ b/crates/service/plugins/users/src/endpoints/session_info_handler.rs @@ -47,6 +47,7 @@ pub struct SessionUserInfo { // Feature / rollout flags: pub can_access_studio: bool, + pub maybe_feature_flags: Option, // Premium plans: pub fakeyou_plan: FakeYouPlan, @@ -166,6 +167,7 @@ pub async fn session_info_handler( // Rollout / feature flags: can_access_studio: session_data.can_access_studio, + maybe_feature_flags: session_data.maybe_feature_flags, // Premium plans: fakeyou_plan: FakeYouPlan::Free, From 8a08793ff035bb1a9ed908fd6efe1e83f8156828 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 02:46:19 -0500 Subject: [PATCH 25/32] feature flags as hash set instead --- .../users/src/endpoints/session_info_handler.rs | 10 ++++++++-- .../users/src/utils/user_session_feature_flags.rs | 7 +++++++ 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/crates/service/plugins/users/src/endpoints/session_info_handler.rs b/crates/service/plugins/users/src/endpoints/session_info_handler.rs index 090e969eb5..498705cae4 100644 --- a/crates/service/plugins/users/src/endpoints/session_info_handler.rs +++ b/crates/service/plugins/users/src/endpoints/session_info_handler.rs @@ -3,6 +3,7 @@ #![forbid(unused_mut)] #![forbid(unused_variables)] +use std::collections::HashSet; use std::fmt; use actix_web::{HttpRequest, HttpResponse, web}; @@ -12,11 +13,13 @@ use log::warn; use sqlx::MySqlPool; use utoipa::ToSchema; +use enums::by_table::users::user_feature_flag::UserFeatureFlag; use http_server_common::response::response_error_helpers::to_simple_json_error; use crate::common_responses::user_details_lite::UserDetailsLight; use crate::cookies::anonymous_visitor_tracking::avt_cookie_manager::AvtCookieManager; use crate::utils::session_checker::SessionChecker; +use crate::utils::user_session_feature_flags::UserSessionFeatureFlags; #[derive(Serialize, Copy, Clone, ToSchema)] #[serde(rename_all = "snake_case")] @@ -47,7 +50,7 @@ pub struct SessionUserInfo { // Feature / rollout flags: pub can_access_studio: bool, - pub maybe_feature_flags: Option, + pub maybe_feature_flags: HashSet, // Premium plans: pub fakeyou_plan: FakeYouPlan, @@ -150,6 +153,9 @@ pub async fn session_info_handler( match maybe_user_session { None => {} Some(session_data) => { + let feature_flags = + UserSessionFeatureFlags::new(session_data.maybe_feature_flags.as_deref()); + if !session_data.is_banned { // NB: Banned users can't be logged in logged_in = true; @@ -167,7 +173,7 @@ pub async fn session_info_handler( // Rollout / feature flags: can_access_studio: session_data.can_access_studio, - maybe_feature_flags: session_data.maybe_feature_flags, + maybe_feature_flags: feature_flags.get_flags(), // Premium plans: fakeyou_plan: FakeYouPlan::Free, diff --git a/crates/service/plugins/users/src/utils/user_session_feature_flags.rs b/crates/service/plugins/users/src/utils/user_session_feature_flags.rs index 80298d944e..8ac9ade6a2 100644 --- a/crates/service/plugins/users/src/utils/user_session_feature_flags.rs +++ b/crates/service/plugins/users/src/utils/user_session_feature_flags.rs @@ -9,6 +9,13 @@ pub struct UserSessionFeatureFlags { } impl UserSessionFeatureFlags { + + pub fn new(maybe_feature_flags: Option<&str>) -> Self { + Self { + maybe_feature_flags: maybe_feature_flags.map(|s| s.to_string()), + } + } + pub fn get_flags(&self) -> HashSet { match self.maybe_feature_flags.as_deref() { None => HashSet::new(), From 037b2017c5ef02dad71e3ddb48bf220a74e51bb4 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 03:00:27 -0500 Subject: [PATCH 26/32] feature flag integrations --- .../media_files/list/list_media_files.rs | 5 +- .../endpoints/workflows/enqueue_comfy_ui.rs | 8 +-- .../service/web/storyteller_web/src/main.rs | 1 + .../web/storyteller_web/src/server_state.rs | 5 ++ .../src/util/allowed_explore_media_access.rs | 52 ++++++++++++++++++ .../allowed_video_style_transfer_access.rs | 55 +++++++++++++++++++ .../web/storyteller_web/src/util/mod.rs | 4 +- 7 files changed, 123 insertions(+), 7 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs create mode 100644 crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs index 833b7d26ae..ccd87d5a3d 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list/list_media_files.rs @@ -24,6 +24,7 @@ use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; +use crate::util::allowed_explore_media_access::allowed_explore_media_access; use crate::util::allowed_studio_access::allowed_studio_access; #[derive(Deserialize, ToSchema, IntoParams)] @@ -136,8 +137,8 @@ pub async fn list_media_files_handler( // ==================== FEATURE FLAG CHECK ==================== // - if !allowed_studio_access(maybe_user_session.as_ref(), &server_state.flags) { - warn!("Storyteller Studio access is not permitted for user"); + if !allowed_explore_media_access(maybe_user_session.as_ref()) { + warn!("Explore media access is not permitted for user"); return Err(ListMediaFilesError::NotAuthorized); } diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs index 3457b9cf63..3eb1bf8099 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs @@ -27,10 +27,10 @@ use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; -use crate::configs::plans::plan_category::{PlanCategory}; +use crate::configs::plans::plan_category::PlanCategory; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; -use crate::util::allowed_studio_access::allowed_studio_access; +use crate::util::allowed_video_style_transfer_access::allowed_video_style_transfer_access; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; /// Debug requests can get routed to special "debug-only" workers, which can @@ -144,8 +144,8 @@ pub async fn enqueue_comfy_ui_handler( // ==================== FEATURE FLAG CHECK ==================== // - if !allowed_studio_access(maybe_user_session.as_ref(), &server_state.flags) { - warn!("Storyteller Studio access is not permitted for user"); + if !allowed_video_style_transfer_access(maybe_user_session.as_ref(), &server_state.flags) { + warn!("Video style transfer access is not permitted for user"); return Err(EnqueueComfyError::NotAuthorized); } diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 88597e2f3e..2c5a809d66 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -361,6 +361,7 @@ async fn main() -> AnyhowResult<()> { switch_voice_conversion_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_VOICE_CONVERSION_TO_MODEL_WEIGHTS", false), switch_tts_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_TTS_TO_MODEL_WEIGHTS", false), force_session_studio_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_STUDIO_FLAG", false), + force_session_video_style_transfer_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_VST_FLAG", false), }; let third_party_url_redirector = ThirdPartyUrlRedirector::new(server_environment); diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index d4f7ed9249..e1c59c252c 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -273,6 +273,11 @@ pub struct StaticFeatureFlags { /// TEMPORARY: Require users to have a studio flag on their `users` table record to use studio features. /// This will also prevent anonymous users from using the service APIs marked as studio-only. pub force_session_studio_flags: bool, + + // TODO(2024-03-05): Remove when done. + /// TEMPORARY: Require users to have a flag on their `users` table record to use video style transfer. + /// This will also prevent anonymous users from using the service APIs marked as VST-only. + pub force_session_video_style_transfer_flags: bool, } /// Instead of top level service denial, these are bans against entities that instead return diff --git a/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs b/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs new file mode 100644 index 0000000000..b9d801edc4 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/allowed_explore_media_access.rs @@ -0,0 +1,52 @@ +use enums::by_table::users::user_feature_flag::UserFeatureFlag; +use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; +use users_component::utils::user_session_extended::UserSessionExtended; +use users_component::utils::user_session_feature_flags::UserSessionFeatureFlags; + +use crate::server_state::StaticFeatureFlags; + +/// Check whether we should allow the request access to explore media. +pub fn allowed_explore_media_access(maybe_session: Option) -> bool { + // NB: We're more careful here since this is sensitive until we have NLP checks. + maybe_session + .map(|ref session| session.can_access_explore_media()) + .unwrap_or(false) +} + +pub trait UserSessionExploreMediaFlag { + fn can_access_explore_media(&self) -> bool; +} + +impl UserSessionExploreMediaFlag for UserSessionExtended { + fn can_access_explore_media(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) + } +} + +impl UserSessionExploreMediaFlag for &UserSessionExtended { + fn can_access_explore_media(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) + } +} + +impl UserSessionExploreMediaFlag for SessionUserRecord { + fn can_access_explore_media(&self) -> bool { + user_session_has_feature(self) + } +} + +impl UserSessionExploreMediaFlag for &SessionUserRecord { + fn can_access_explore_media(&self) -> bool { + user_session_has_feature(self) + } +} + +fn user_session_has_feature(user_session: &SessionUserRecord) -> bool { + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: user_session.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::ExploreMedia) +} diff --git a/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs b/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs new file mode 100644 index 0000000000..2203d9a327 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/allowed_video_style_transfer_access.rs @@ -0,0 +1,55 @@ +use enums::by_table::users::user_feature_flag::UserFeatureFlag; +use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; +use users_component::utils::user_session_extended::UserSessionExtended; +use users_component::utils::user_session_feature_flags::UserSessionFeatureFlags; + +use crate::server_state::StaticFeatureFlags; + +/// Check whether we should allow the request access to video style transfer +pub fn allowed_video_style_transfer_access(maybe_session: Option, flags: &StaticFeatureFlags) -> bool { + if !flags.force_session_video_style_transfer_flags { + return true; + } + + maybe_session + .map(|ref session| session.can_access_video_style_transfer()) + .unwrap_or(false) +} + +pub trait UserSessionVideoStyleTransferFlag { + fn can_access_video_style_transfer(&self) -> bool; +} + +impl UserSessionVideoStyleTransferFlag for UserSessionExtended { + fn can_access_video_style_transfer(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) + } +} + +impl UserSessionVideoStyleTransferFlag for &UserSessionExtended { + fn can_access_video_style_transfer(&self) -> bool { + self.feature_flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) + } +} + +impl UserSessionVideoStyleTransferFlag for SessionUserRecord { + fn can_access_video_style_transfer(&self) -> bool { + user_session_has_feature(self) + } +} + +impl UserSessionVideoStyleTransferFlag for &SessionUserRecord { + fn can_access_video_style_transfer(&self) -> bool { + user_session_has_feature(self) + } +} + +fn user_session_has_feature(user_session: &SessionUserRecord) -> bool { + // TODO(bt, 2024-03-05): this is horrible. + // There should be a wrapper class between the query and the caller. + let flags = UserSessionFeatureFlags { + maybe_feature_flags: user_session.maybe_feature_flags.clone(), + }; + + flags.has_permission_unoptimized(UserFeatureFlag::VideoStyleTransfer) +} diff --git a/crates/service/web/storyteller_web/src/util/mod.rs b/crates/service/web/storyteller_web/src/util/mod.rs index 13b83e2474..1b4f55b6ed 100644 --- a/crates/service/web/storyteller_web/src/util/mod.rs +++ b/crates/service/web/storyteller_web/src/util/mod.rs @@ -1,5 +1,7 @@ +pub mod allowed_explore_media_access; pub mod allowed_studio_access; +pub mod allowed_video_style_transfer_access; pub mod delete_role_disambiguation; pub mod encrypted_sort_id; pub mod placeholder_images; -pub mod troll_user_bans; +pub mod troll_user_bans; \ No newline at end of file From d246797e6342f2a48ff7f5b557202e167d593070 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 03:15:38 -0500 Subject: [PATCH 27/32] remove old feature flags (100% rolled out) --- .../enqueue_infer_tts_handler.rs | 131 ++++++------------ .../models/list_voice_conversion_models.rs | 2 +- .../service/web/storyteller_web/src/main.rs | 2 - .../web/storyteller_web/src/server_state.rs | 10 -- 4 files changed, 42 insertions(+), 103 deletions(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs index d287528d28..3f7a00cfac 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/enqueue_infer_tts_handler/enqueue_infer_tts_handler.rs @@ -24,9 +24,7 @@ use http_server_common::request::get_request_ip::get_request_ip; use migration::text_to_speech::get_tts_model_for_enqueue_inference_migration::TtsModelForEnqueueInferenceMigrationWrapper; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated}; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; -use mysql_queries::queries::tts::tts_inference_jobs::insert_tts_inference_job::TtsInferenceJobInsertBuilder; use redis_common::redis_keys::RedisKeys; -use tokens::tokens::tts_inference_jobs::TtsInferenceJobToken; use tokens::tokens::users::UserToken; use tts_common::priority::{FAKEYOU_DEFAULT_VALID_API_TOKEN_PRIORITY_LEVEL, FAKEYOU_INVESTOR_PRIORITY_LEVEL}; use user_input_common::check_for_slurs::contains_slurs; @@ -55,10 +53,6 @@ const STORYTELLER_DEMO_COOKIE_NAME : &str = "storyteller_demo"; /// const DEBUG_HEADER_NAME : &str = "enable-debug-mode"; -/// Requests with this header can be forced onto the new jobs system. -/// It'll be used to help test TT2 on the new system before it is rolled out for everyone. -const NEW_JOB_SYSTEM_HEADER_NAME : &str = "new-job-system"; - /// The routing tag header can send workloads to particular k8s hosts. /// This is useful for catching the live logs or intercepting the job. /// NB: This is only for the new job system. @@ -234,9 +228,6 @@ pub async fn enqueue_infer_tts_handler( let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME) .is_some(); - let use_new_job_system_for_request = get_request_header_optional(&http_request, NEW_JOB_SYSTEM_HEADER_NAME) - .is_some(); - let maybe_routing_tag= get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME) .map(|routing_tag| routing_tag.trim().to_string()); @@ -377,95 +368,55 @@ pub async fn enqueue_infer_tts_handler( // ==================== ENQUEUE APPROPRIATE TTS TYPE ==================== // - let use_new_job_system = use_new_job_system_for_request || server_state.flags.enable_enqueue_generic_tts_job; - - let job_token; - let job_token_type; - let model_type = InferenceModelType::Tacotron2; let max_duration_seconds = plan.tts_max_duration_seconds(); - if use_new_job_system { - // This branch uses the `generic-inference-job` service and tables. - info!("Creating tts inference job record (new generic job system)..."); + // This branch uses the `generic-inference-job` service and tables. + info!("Creating tts inference job record (new generic job system)..."); - let maybe_creator_user_token_typed = maybe_user_token - .as_deref() - .map(|token| UserToken::new_from_str(token)); + let maybe_creator_user_token_typed = maybe_user_token + .as_deref() + .map(|token| UserToken::new_from_str(token)); - let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); + let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); - let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { - uuid_idempotency_token: &request.uuid_idempotency_token, - job_type: InferenceJobType::Tacotron2, - inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(model_type), - maybe_model_token: Some(request.tts_model_token.as_str()), - maybe_input_source_token: None, // NB: TTS doesn't have input media - maybe_input_source_token_type: None, // NB: TTS doesn't have input media - maybe_download_url: None, - maybe_cover_image_media_file_token: None, - maybe_raw_inference_text: Some(inference_text.as_str()), - maybe_max_duration_seconds: Some(max_duration_seconds), - maybe_inference_args: Some(GenericInferenceArgs { - inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech), - args: None, // NB: We don't need to encode any args yet. - }), - maybe_creator_user_token: maybe_creator_user_token_typed.as_ref(), - maybe_avt_token: maybe_avt_token.as_ref(), - creator_ip_address: &ip_address, - creator_set_visibility: set_visibility, - priority_level, - requires_keepalive: false, // TODO: We may want this to be the case in the future. - is_debug_request, - maybe_routing_tag: maybe_routing_tag.as_deref(), - mysql_pool: &server_state.mysql_pool, - }).await; + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + job_type: InferenceJobType::Tacotron2, + inference_category: InferenceCategory::TextToSpeech, + maybe_model_type: Some(model_type), + maybe_model_token: Some(request.tts_model_token.as_str()), + maybe_input_source_token: None, // NB: TTS doesn't have input media + maybe_input_source_token_type: None, // NB: TTS doesn't have input media + maybe_download_url: None, + maybe_cover_image_media_file_token: None, + maybe_raw_inference_text: Some(inference_text.as_str()), + maybe_max_duration_seconds: Some(max_duration_seconds), + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::TextToSpeech), + args: None, // NB: We don't need to encode any args yet. + }), + maybe_creator_user_token: maybe_creator_user_token_typed.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: set_visibility, + priority_level, + requires_keepalive: false, // TODO: We may want this to be the case in the future. + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; - let inference_job_token = match query_result { - Ok((inference_job_token, _id)) => inference_job_token, - Err(err) => { - warn!("New (generic) tts inference job creation DB error: {:?}", err); - return Err(InferTtsError::ServerError); - } - }; - - job_token = inference_job_token.to_string(); - job_token_type = InferenceJobTokenType::Generic; - - } else { - // This branch uses the `tts-inference-job` service and tables. - info!("Creating tts inference job record (legacy job system)..."); - - // This token is returned to the client. - job_token = TtsInferenceJobToken::generate().to_string(); - - let query_result = TtsInferenceJobInsertBuilder::new_for_fakeyou_request() - .set_job_token(&job_token) - .set_uuid_idempotency_token(&request.uuid_idempotency_token) - .set_model_token(&request.tts_model_token) - .set_raw_inference_text(&inference_text) - .set_maybe_creator_user_token(maybe_user_token.as_deref()) - .set_creator_ip_address(&ip_address) - .set_creator_set_visibility(set_visibility.to_str()) - .set_priority_level(priority_level) - .set_max_duration_seconds(max_duration_seconds) - .set_is_from_api(is_from_api) - .set_is_debug_request(is_debug_request) - .insert(&server_state.mysql_pool) - .await; - - match query_result { - Ok(_) => {}, - Err(err) => { - warn!("New tts inference job creation DB error: {:?}", err); - return Err(InferTtsError::ServerError); - } + let inference_job_token = match query_result { + Ok((inference_job_token, _id)) => inference_job_token, + Err(err) => { + warn!("New (generic) tts inference job creation DB error: {:?}", err); + return Err(InferTtsError::ServerError); } + }; - job_token_type = InferenceJobTokenType::LegacyTts; - } + let job_token = inference_job_token.to_string(); server_state.firehose_publisher.enqueue_tts_inference(maybe_user_token.as_deref(), &job_token, &request.tts_model_token) .await @@ -476,8 +427,8 @@ pub async fn enqueue_infer_tts_handler( let response = InferTtsSuccessResponse { success: true, - inference_job_token: job_token.to_string(), - inference_job_token_type: job_token_type, + inference_job_token: job_token, + inference_job_token_type: InferenceJobTokenType::Generic, }; let body = serde_json::to_string(&response) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/models/list_voice_conversion_models.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/models/list_voice_conversion_models.rs index 68b206f143..aeaf7cd6f9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/models/list_voice_conversion_models.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_conversion/models/list_voice_conversion_models.rs @@ -134,7 +134,7 @@ pub async fn list_voice_conversion_models_handler( })?; // TODO: Fail open in case the DB is down. Pull from expired cache if query fails. - let models = get_all_models(&mut mysql_connection, server_state.flags.switch_voice_conversion_to_model_weights) + let models = get_all_models(&mut mysql_connection, true) .await .map_err(|e| { error!("Error querying database: {:?}", e); diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 2c5a809d66..a61474c0d9 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -357,8 +357,6 @@ async fn main() -> AnyhowResult<()> { troll_ban_user_percent: easyenv::get_env_num("FF_TROLL_BANNED_USER_PERCENT", 0)?, // Temporary flags - enable_enqueue_generic_tts_job: easyenv::get_env_bool_or_default("FF_ENABLE_ENQUEUE_GENERIC_TTS_JOB", false), - switch_voice_conversion_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_VOICE_CONVERSION_TO_MODEL_WEIGHTS", false), switch_tts_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_TTS_TO_MODEL_WEIGHTS", false), force_session_studio_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_STUDIO_FLAG", false), force_session_video_style_transfer_flags: easyenv::get_env_bool_or_default("FF_FORCE_SESSION_VST_FLAG", false), diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index e1c59c252c..8c2cd25211 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -254,16 +254,6 @@ pub struct StaticFeatureFlags { /// This should be a number over 100. pub troll_ban_user_percent: u8, - // TODO(2023-03-20): Remove temporary flag when done. - // NB(bt,2023-12-18): This is rolled out (=="TRUE") - /// TEMPORARY: Control enqueuing TTS jobs to the generic job worker. - pub enable_enqueue_generic_tts_job: bool, - - // TODO(2023-12-18): Remove temporary flag when done. - /// TEMPORARY: Move voice control model listing over to `model_weights` from `voice_conversion_models` - /// This will control all downstream enqueuing, jobs, etc. - pub switch_voice_conversion_to_model_weights: bool, - // TODO(2024-01-13): Remove temporary flag when done. /// TEMPORARY: Move voice control model listing over to `model_weights` from `tts_models` /// This will control all downstream enqueuing, jobs, etc. From 003500e6560485c9763fe1358a094f3ddb659293 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 03:57:57 -0500 Subject: [PATCH 28/32] new user feature flag rollouts --- _sql/actions/user_access/new_enroll.sql | 55 +++++++++++++++++++++++++ 1 file changed, 55 insertions(+) create mode 100644 _sql/actions/user_access/new_enroll.sql diff --git a/_sql/actions/user_access/new_enroll.sql b/_sql/actions/user_access/new_enroll.sql new file mode 100644 index 0000000000..36242bd843 --- /dev/null +++ b/_sql/actions/user_access/new_enroll.sql @@ -0,0 +1,55 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Staff (1) +-- teddanson is victor's alt +UPDATE users +SET + can_access_studio = false, + maybe_feature_flags = 'explore_media,studio,video_style_transfer' +WHERE username IN ( + 'bflat', + 'brandon', + 'crossproduct', + 'crossproduct1', + 'echelon', + 'olivicmic', + 'saltacc', + 'teddanson', + 'vegito1089', + 'wilwong', + 'zzz_last_item' +); + +-- Staff (2) +-- el_cid_93 is for testing (who is this??) +-- endtimes is @sugarbro (testing) +-- tammieteller is Tammie (Pebblebed) +UPDATE users +SET + can_access_studio = false, + maybe_feature_flags = 'studio,video_style_transfer' +WHERE username IN ( + 'el_cid_93', + 'endtimes', + 'fyscott', + 'gateway', + 'jags111', + 'justinjohn0306', + 'mrvintage', + 'rewin123', + 'tammieteller', + 'yae_ph', + 'zzz_last_item' +); + +-- Early access to VST only +UPDATE users +SET + can_access_studio = false, + maybe_feature_flags = 'video_style_transfer' +WHERE username IN ( + 'zzz_last_item' +); + From f40d106b59e62be6c6e7b1eb96ce5c3f6d5447f9 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 6 Mar 2024 21:02:18 -0500 Subject: [PATCH 29/32] onboard fakeyou users --- .../user_access/enroll_users_in_studio.sql | 36 ------------------- _sql/actions/user_access/new_enroll.sql | 15 +++++--- 2 files changed, 11 insertions(+), 40 deletions(-) delete mode 100644 _sql/actions/user_access/enroll_users_in_studio.sql diff --git a/_sql/actions/user_access/enroll_users_in_studio.sql b/_sql/actions/user_access/enroll_users_in_studio.sql deleted file mode 100644 index c066acf8f2..0000000000 --- a/_sql/actions/user_access/enroll_users_in_studio.sql +++ /dev/null @@ -1,36 +0,0 @@ --- noinspection SqlDialectInspectionForFile --- noinspection SqlNoDataSourceInspectionForFile --- noinspection SqlResolveForFile - --- Staff --- teddanson is victor's alt --- endtimes is @sugarbro (testing) --- el_cid_93 is for testing --- SELECT token, username FROM users WHERE username IN ( -UPDATE users SET can_access_studio = true WHERE username IN ( - 'bflat', - 'brandon', - 'crossproduct', - 'crossproduct1', - 'echelon', - 'el_cid_93', - 'endtimes', - 'fyscott', - 'gateway', - 'jags111', - 'mrvintage', - 'olivicmic', - 'rewin123', - 'saltacc', - 'teddanson', - 'vegito1089', - 'wilwong', - 'yae_ph', - 'zzz_last_item' -); - -UPDATE users SET can_access_studio = true WHERE username IN ( - 'endtimes', - 'tropicalfun', - 'zzz_last_item' -); diff --git a/_sql/actions/user_access/new_enroll.sql b/_sql/actions/user_access/new_enroll.sql index 36242bd843..3c755913b7 100644 --- a/_sql/actions/user_access/new_enroll.sql +++ b/_sql/actions/user_access/new_enroll.sql @@ -3,7 +3,7 @@ -- noinspection SqlResolveForFile -- Staff (1) --- teddanson is victor's alt +-- * teddanson is victor's alt UPDATE users SET can_access_studio = false, @@ -23,9 +23,9 @@ WHERE username IN ( ); -- Staff (2) --- el_cid_93 is for testing (who is this??) --- endtimes is @sugarbro (testing) --- tammieteller is Tammie (Pebblebed) +-- * el_cid_93 is for testing (who is this??) +-- * endtimes is @sugarbro (testing) +-- * tammieteller is Tammie (Pebblebed) UPDATE users SET can_access_studio = false, @@ -40,6 +40,7 @@ WHERE username IN ( 'mrvintage', 'rewin123', 'tammieteller', + 'tropicalfun', 'yae_ph', 'zzz_last_item' ); @@ -50,6 +51,12 @@ SET can_access_studio = false, maybe_feature_flags = 'video_style_transfer' WHERE username IN ( + 'kenjoplays', + 'ofccccccc', + 'sonicgt2', + 'stewiegroffin', + 'tanooki426', + 'wawoul', 'zzz_last_item' ); From 5c0473d407755a08ba059c2d161d17e035c25595 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 7 Mar 2024 00:48:45 -0500 Subject: [PATCH 30/32] stubs for status alert endpoint --- .../web/storyteller_web/src/api_doc.rs | 10 +- .../src/http_server/endpoints/service/mod.rs | 3 +- .../endpoints/service/status_alert_handler.rs | 98 +++++++++++++++++++ .../service/web/storyteller_web/src/routes.rs | 8 ++ 4 files changed, 116 insertions(+), 3 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs diff --git a/crates/service/web/storyteller_web/src/api_doc.rs b/crates/service/web/storyteller_web/src/api_doc.rs index cecaf6c8cd..e476f35eeb 100644 --- a/crates/service/web/storyteller_web/src/api_doc.rs +++ b/crates/service/web/storyteller_web/src/api_doc.rs @@ -44,6 +44,7 @@ use crate::http_server::endpoints::media_files::upload::upload_media_file_handle use crate::http_server::endpoints::media_files::upload_video::upload_error::VideoMediaFileUploadError; use crate::http_server::endpoints::media_files::upload_video::upload_video_media_file_handler::*; use crate::http_server::endpoints::prompts::get_prompt::*; +use crate::http_server::endpoints::service::status_alert_handler::*; use crate::http_server::endpoints::user_bookmarks::batch_get_user_bookmarks_handler::*; use crate::http_server::endpoints::user_bookmarks::create_user_bookmark_handler::*; use crate::http_server::endpoints::user_bookmarks::delete_user_bookmark_handler::*; @@ -79,6 +80,7 @@ use crate::http_server::web_utils::response_success_helpers::*; crate::http_server::endpoints::media_files::upload::upload_media_file_handler::upload_media_file_handler, crate::http_server::endpoints::media_files::upload_video::upload_video_media_file_handler::upload_video_media_file_handler, crate::http_server::endpoints::prompts::get_prompt::get_prompt_handler, + crate::http_server::endpoints::service::status_alert_handler::status_alert_handler, crate::http_server::endpoints::user_bookmarks::batch_get_user_bookmarks_handler::batch_get_user_bookmarks_handler, crate::http_server::endpoints::user_bookmarks::create_user_bookmark_handler::create_user_bookmark_handler, crate::http_server::endpoints::user_bookmarks::delete_user_bookmark_handler::delete_user_bookmark_handler, @@ -233,10 +235,12 @@ use crate::http_server::web_utils::response_success_helpers::*; SetUserRatingError, SetUserRatingRequest, SetUserRatingResponse, + StatusAlertCategory, + StatusAlertError, + StatusAlertInfo, + StatusAlertResponse, StatusDetailsResponse, StorytellerStreamPlan, - VideoMediaFileUploadError, - UploadVideoMediaSuccessResponse, TerminateInferenceJobError, TerminateInferenceJobPathInfo, TerminateInferenceJobSuccessResponse, @@ -244,6 +248,7 @@ use crate::http_server::web_utils::response_success_helpers::*; UpdateWeightPathInfo, UpdateWeightRequest, UploadMediaSuccessResponse, + UploadVideoMediaSuccessResponse, UserBookmarkDetailsForUserList, UserBookmarkEntityType, UserBookmarkForEntityListItem, @@ -253,6 +258,7 @@ use crate::http_server::web_utils::response_success_helpers::*; UserProfileUserBadge, UserRatingEntityType, UserRatingValue, + VideoMediaFileUploadError, Weight, WeightsData, ZsDatasetRecord, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/service/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/service/mod.rs index a8069876cf..c53223a497 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/service/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/service/mod.rs @@ -1,2 +1,3 @@ pub mod health_check_handler; -pub mod public_info_handler; \ No newline at end of file +pub mod public_info_handler; +pub mod status_alert_handler; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs new file mode 100644 index 0000000000..1a5343f05d --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs @@ -0,0 +1,98 @@ +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, ResponseError, web}; +use actix_web::http::StatusCode; +use utoipa::ToSchema; + +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; + +use crate::server_state::ServerState; + +#[derive(Serialize, ToSchema)] +pub struct StatusAlertResponse { + pub success: bool, + + /// If there's an alert, this will be set. + /// The sub keys are optional, but at least one of them will be set. + /// i.e. we can have an alert with no message or no predefined category. + pub maybe_alert: Option, +} + +#[derive(Serialize, ToSchema)] +pub struct StatusAlertInfo { + /// If an alert is set, this might be a key for a common i18n message to show. + pub maybe_category: Option, + + /// We can set an optional message to show to the frontend. + /// + /// This field may or may not be set. If it's set, this is custom text that should + /// be displayed to the user. For example, we might tell the user when we'll be + /// back online, etc. Both this field and `maybe_category` are optional. + /// + /// If `maybe_category` is present, it should be a key for a predefined and + /// internationalized (i18n) message stored on the frontend. + /// + /// Either or both of these fields may be set. + pub maybe_message: Option, +} + +#[derive(Serialize, ToSchema, Copy, Clone)] +#[serde(rename_all = "snake_case")] +pub enum StatusAlertCategory { + /// System is down for maintenance + DownForMaintenance, +} + +#[derive(Debug, Serialize, ToSchema)] +pub enum StatusAlertError { + ServerError, +} + +impl ResponseError for StatusAlertError { + fn status_code(&self) -> StatusCode { + match *self { + StatusAlertError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl std::fmt::Display for StatusAlertError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +#[utoipa::path( + get, + path = "/v1/status_alert_check", + responses( + ( + status = 200, + description = "Check service status for frontend alert messages", + body = StatusAlertResponse, + ), + (status = 500, description = "Server error", body = StatusAlertError) + ), +)] +pub async fn status_alert_handler( + _http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let response = StatusAlertResponse { + success: true, + maybe_alert: None, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| StatusAlertError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index d2d4965f9e..8406e3fb81 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -85,6 +85,7 @@ use crate::http_server::endpoints::moderation::users::list_users::list_users_han use crate::http_server::endpoints::prompts::get_prompt::get_prompt_handler; use crate::http_server::endpoints::service::health_check_handler::get_health_check_handler; use crate::http_server::endpoints::service::public_info_handler::get_public_info_handler; +use crate::http_server::endpoints::service::status_alert_handler::status_alert_handler; use crate::http_server::endpoints::stats::get_unified_queue_stats::get_unified_queue_stats_handler; use crate::http_server::endpoints::stubs::app_model_downloads::get_app_model_downloads_handler; use crate::http_server::endpoints::stubs::app_news::get_app_news_handler; @@ -270,6 +271,7 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> .into_app(); // =================== BVH from Workflow ==================== + let mut app = RouteBuilder::from_app(app) .add_post("/v1/conversion/enqueue_bvh_to_workflow", enqueue_bvh_to_workflow_handler) .into_app(); @@ -280,6 +282,12 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> .add_get("/v1/stats/queues", get_unified_queue_stats_handler) .into_app(); + // ==================== Service Status ==================== + + let mut app = RouteBuilder::from_app(app) + .add_get("/v1/status_alert_check", status_alert_handler) + .into_app(); + // ==================== COMPONENTS ==================== app = add_suggested_api_v1_account_creation_and_session_routes(app); // /create_account, /session, /login, /logout From 4981a739d9b53296eca9940955d0b21457602d74 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 7 Mar 2024 01:05:51 -0500 Subject: [PATCH 31/32] Feature flags to populate status alert message/category --- .../endpoints/service/status_alert_handler.rs | 31 ++++++++++++++++++- .../service/web/storyteller_web/src/main.rs | 6 +++- .../web/storyteller_web/src/server_state.rs | 6 ++++ 3 files changed, 41 insertions(+), 2 deletions(-) diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs index 1a5343f05d..e0bf03036f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/service/status_alert_handler.rs @@ -84,9 +84,30 @@ pub async fn status_alert_handler( server_state: web::Data> ) -> Result { + let maybe_category = server_state + .flags + .maybe_status_alert_category + .as_deref() + .map(|category| category_to_enum(&category)) + .flatten(); + + let maybe_message = server_state + .flags + .maybe_status_alert_custom_message + .as_deref() + .map(|message| message.trim().to_string()); + + let maybe_alert = match (maybe_category, maybe_message) { + (None, None) => None, + (category, message) => Some(StatusAlertInfo { + maybe_category: category, + maybe_message: message, + }), + }; + let response = StatusAlertResponse { success: true, - maybe_alert: None, + maybe_alert, }; let body = serde_json::to_string(&response) @@ -96,3 +117,11 @@ pub async fn status_alert_handler( .content_type("application/json") .body(body)) } + +fn category_to_enum(category: &str) -> Option { + let key = category.trim(); + match key { + "down_for_maintenance" => Some(StatusAlertCategory::DownForMaintenance), + _ => None, + } +} diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index a61474c0d9..6940c36c94 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -340,7 +340,11 @@ async fn main() -> AnyhowResult<()> { .ok_or(anyhow!("invalid server environment"))?; let service_feature_flags = StaticFeatureFlags { - // Permanent (control plane / safety) flags + // Permanent (control plane / safety) flags : messaging + maybe_status_alert_category: easyenv::get_env_string_optional("FF_STATUS_ALERT_CATEGORY"), + maybe_status_alert_custom_message: easyenv::get_env_string_optional("FF_STATUS_ALERT_CUSTOM_MESSAGE"), + + // Permanent (control plane / safety) flags : disabling features global_429_pushback_filter_enabled: easyenv::get_env_bool_or_default("FF_GLOBAL_429_PUSHBACK_FILTER_ENABLED", false), disable_unified_queue_stats_endpoint: easyenv::get_env_bool_or_default("FF_DISABLE_QUEUE_STATS_ENDPOINT", false), disable_inference_queue_length_endpoint: easyenv::get_env_bool_or_default("FF_DISABLE_INFERENCE_QUEUE_LENGTH_ENDPOINT", false), diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index 8c2cd25211..b7789edb6f 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -223,6 +223,12 @@ pub struct StaticFeatureFlags { /// Used to bring the service back online slowly. pub global_429_pushback_filter_enabled: bool, + /// If we're suffering an outage, set the alert category to display a predefined message to users. + pub maybe_status_alert_category: Option, + + /// If we're suffering an outage, set custom text for the alert message. + pub maybe_status_alert_custom_message: Option, + /// Disable the live `/v1/stats/queues` endpoint for all users and serve a static value instead. pub disable_unified_queue_stats_endpoint: bool, From 2a89943f45031278b9f032c3be16d3a21a67ccd5 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 7 Mar 2024 03:36:42 -0500 Subject: [PATCH 32/32] Mocap job: read video file dimensions --- Cargo.lock | 1 + crates/lib/files/videos/src/get_mp4_info.rs | 27 +++++++- crates/service/job/inference_job/Cargo.toml | 5 +- .../job_types/mocap/mocap_net/process_job.rs | 66 ++++++++++++++----- .../upload_video_media_file_handler.rs | 4 +- 5 files changed, 78 insertions(+), 25 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 70f3cb9a12..d6f1270264 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -3241,6 +3241,7 @@ dependencies = [ "tokens", "tokio", "tts_common", + "videos", "walkdir", ] diff --git a/crates/lib/files/videos/src/get_mp4_info.rs b/crates/lib/files/videos/src/get_mp4_info.rs index 2a0327bf6c..11bd2f60fc 100644 --- a/crates/lib/files/videos/src/get_mp4_info.rs +++ b/crates/lib/files/videos/src/get_mp4_info.rs @@ -6,30 +6,47 @@ use mp4::TrackType; use errors::AnyhowResult; +#[derive(Debug, Clone)] pub struct Mp4Info { /// Framerate of the longest video track. pub framerate: f64, + /// Duration of the entire mp4. pub duration_millis: u128, + + /// Width of the longest video track. + pub width: u16, + + /// Height of the longest video track. + pub height: u16, } -pub fn get_mp4_info_for_bytes(file_bytes: &[u8], file_size: u64) -> AnyhowResult { +pub fn get_mp4_info_for_bytes(file_bytes: &[u8]) -> AnyhowResult { + get_mp4_info_for_bytes_and_len(file_bytes, file_bytes.len()) +} + +pub fn get_mp4_info_for_bytes_and_len(file_bytes: &[u8], file_size: usize) -> AnyhowResult { let reader = BufReader::new(Cursor::new(file_bytes)); - get_mp4_info(reader, file_size) + get_mp4_info(reader, file_size as u64) } pub fn get_mp4_info(reader: T, file_size: u64) -> AnyhowResult { let mp4 = mp4::Mp4Reader::read_header(reader, file_size)?; - let mut framerate = 0.0; let mut longest_duration = Duration::from_secs(0); + let mut framerate = 0.0; + let mut width = 0; + let mut height = 0; + for track in mp4.tracks().values() { match track.track_type() { Ok(TrackType::Video) => { if track.duration() > longest_duration { longest_duration = track.duration(); framerate = track.frame_rate(); + width = track.width(); + height = track.height(); } } Err(err) => { @@ -43,6 +60,8 @@ pub fn get_mp4_info(reader: T, file_size: u64) -> AnyhowResult) -> Result return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))), + Some(args) => args, + }; + + let mc_args = match poly_args { + Mc(args) => args, + _ => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))), + }; + // ==================== TEMP DIR ==================== // let work_temp_dir = format!("temp_mocapnet_inference_{}", job.id.0); @@ -86,6 +106,22 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result) -> Result return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))), - Some(args) => args, - }; - - let mc_args = match poly_args { - Mc(args) => args, - _ => return Err(ProcessSingleJobError::Other(anyhow!("Mocap args not found"))), - }; - // ==================== RUN INFERENCE SCRIPT ==================== // let stderr_output_file = work_temp_dir.path().join("stderr.txt"); let inference_start_time = Instant::now(); @@ -130,8 +151,8 @@ pub async fn process_job(args: MocapNetProcessJobArgs<'_>) -> Result) -> Result Option { + let video_file = File::open(video_filename).ok()?; + let file_length = video_file.metadata().ok()?.len(); + let reader = BufReader::new(video_file); + let info = get_mp4_info(reader, file_length).ok()?; + Some(info) +} \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs index 3a01c23f09..6eac332b77 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload_video/upload_video_media_file_handler.rs @@ -18,7 +18,7 @@ use mimetypes::mimetype_to_extension::mimetype_to_extension; use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType}; use tokens::tokens::media_files::MediaFileToken; -use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes}; +use videos::get_mp4_info::{get_mp4_info, get_mp4_info_for_bytes, get_mp4_info_for_bytes_and_len}; use crate::http_server::endpoints::engine::create_scene_handler::{CreateSceneError, CreateSceneSuccessResponse}; use crate::http_server::endpoints::media_files::upload_video::drain_multipart_request::drain_multipart_request; @@ -176,7 +176,7 @@ pub async fn upload_video_media_file_handler( // ==================== FRAME RATE ==================== // - let mp4_info = get_mp4_info_for_bytes(file_bytes.as_ref(), file_size_bytes as u64) + let mp4_info = get_mp4_info_for_bytes(file_bytes.as_ref()) .map_err(|err| { warn!("Error reading mp4 info: {:?}", err); VideoMediaFileUploadError::ServerError