diff --git a/.env b/.env index 7545992311..4ec2066144 100644 --- a/.env +++ b/.env @@ -39,16 +39,20 @@ W2L_END_BUMP_FILENAME=vocodes-short-end-bump.mp4 W2L_DOWNLOAD_SCRIPT_NAME=vocodes_process_upload.py -TTS_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts +TTS_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-ml/tts TTS_INFERENCE_SCRIPT_NAME=vocodes_inference.py -TACOTRON_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts -TACOTRON_VIRTUAL_ENV_ACTIVATION_COMMAND="source python-tacotron/bin/activate" +TACOTRON_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-ml/tts +TACOTRON_VIRTUAL_ENV_ACTIVATION_COMMAND="source python/bin/activate" TACOTRON_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_tacotron.py" -TALKNET_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-web/models/tts +TALKNET_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-ml/tts TALKNET_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_talknet.py" +HIFIGAN_ROOT_CODE_DIRECTORY=/home/bt/dev/storyteller/storyteller-ml/tts +HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND="source python/bin/activate" +HIFIGAN_MODEL_CHECK_SCRIPT_NAME="vocodes_model_check_hifigan.py" + TTS_INFERENCE_SIDECAR_HOSTNAME=localhost:8000 #TTS_WAVEGLOW_VOCODER_MODEL_FILENAME=waveglow_256channels_universal_v5.pt # THIS SOUNDS HORRID diff --git a/Cargo.lock b/Cargo.lock index 808d284926..bdc3926fe1 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1297,6 +1297,7 @@ dependencies = [ "log", "once_cell", "regex", + "reusable_types", "serde", "serde_derive", "serde_json", @@ -1700,6 +1701,38 @@ dependencies = [ "version_check", ] +[[package]] +name = "generic-download-job" +version = "0.0.1" +dependencies = [ + "anyhow", + "chrono", + "clap", + "config", + "container_common", + "database_queries", + "dotenv", + "easyenv", + "google_drive_common", + "hostname", + "hyper 0.14.14", + "jobs_common", + "lfu", + "log", + "newrelic-telemetry", + "r2d2_redis", + "reusable_types", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "storage_buckets_common", + "subprocess", + "tempdir", + "tokio 1.13.0", + "tts_common", +] + [[package]] name = "getopts" version = "0.2.21" @@ -3225,6 +3258,19 @@ dependencies = [ "quick-error", ] +[[package]] +name = "reusable_types" +version = "0.0.1" +dependencies = [ + "anyhow", + "chrono", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "toml", +] + [[package]] name = "ring" version = "0.16.20" @@ -3849,12 +3895,14 @@ dependencies = [ "memory_caching", "native-tls", "newrelic-telemetry", + "once_cell", "pulldown-cmark", "r2d2_redis", "rand 0.8.4", "redis_common", "regex", "reqwest", + "reusable_types", "ring 0.17.0-alpha.10", "rust-s3", "serde", diff --git a/Cargo.toml b/Cargo.toml index ba0c482263..64686d9906 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -10,12 +10,14 @@ members = [ "crates/lib/jobs_common", "crates/lib/memory_caching", "crates/lib/redis_common", + "crates/lib/reusable_types", "crates/lib/storage_buckets_common", "crates/lib/tts_common", "crates/lib/twitch_common", "crates/lib/user_input_common", # Application + "crates/service/generic_download_job", "crates/service/storyteller_web", "crates/service/tts_download_job", "crates/service/tts_inference_job", diff --git a/_migrations/2022-08-05-202017_generic_download_jobs/down.sql b/_migrations/2022-08-05-202017_generic_download_jobs/down.sql new file mode 100644 index 0000000000..38e2110e12 --- /dev/null +++ b/_migrations/2022-08-05-202017_generic_download_jobs/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table generic_download_jobs; \ No newline at end of file diff --git a/_migrations/2022-08-05-202017_generic_download_jobs/up.sql b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql new file mode 100644 index 0000000000..17a883e0a6 --- /dev/null +++ b/_migrations/2022-08-05-202017_generic_download_jobs/up.sql @@ -0,0 +1,110 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE generic_download_jobs ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + -- This is so the in-progress results can be looked up by the UI. + token VARCHAR(32) NOT NULL, + + -- Idempotency token + -- This is so the frontend client doesn't submit duplicate jobs. + uuid_idempotency_token VARCHAR(36) NOT NULL, + + -- ========== SUCCESS CASE ========== + + -- If the job completes successfully, this is the result token. + -- This is only populated on a successful result. + on_success_downloaded_entity_token VARCHAR(32) DEFAULT NULL, + + -- The type of the object will vary based on the type of the upload, + -- and we may include heuristics that auto-detect types in the future + on_success_downloaded_entity_type VARCHAR(32) DEFAULT NULL, + + -- ========== UPLOAD DETAILS ========== + + -- A "well defined" type of upload + -- The first upload type will be "hifigan_model" + download_type VARCHAR(32) NOT NULL, + + -- Assume the same URL can be uploaded more than once (in case of errors or + -- people sharing models) + + download_url VARCHAR(512) NOT NULL, + + -- The user-supplied title of the upload + -- This might have different meanings for different upload types + -- The "name" of the voice model, which might be complicated. + -- If maybe_subject_token (etc.) is set, then it's authoritative instead. + title VARCHAR(255) NOT NULL, + + -- ========== CREATOR DETAILS AND PREFERENCES ========== + + -- Foreign key to user + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address VARCHAR(40) NOT NULL, + + -- The creator can set a desired visibility for their data. + -- This does not always apply to every upload type. + -- Additionally, some upload types may require moderator approval prior + -- to being publicly listed, and this field has no bearing on that. + -- NB: DO NOT CHANGE THE ORDER OF THE ENUM VALUES DURING SCHEMA MIGRATIONS. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== JOB SYSTEM DETAILS ========== + + -- Jobs begin as "pending", then transition to other states. + -- + -- * Pending = job is ready to go + -- * Started = job is running + -- * Complete_Success = job is done (success) + -- * Complete_Failure = job is done (failure) + -- * Attempt_Failed = job failed but may retry. + -- * Dead = job failed permanently. + -- + -- Pending -> Started -> Complete_Success + -- |-> Complete_Failure + -- \-> Attempt_Failed -> Started -> { Complete, Failed, Dead } + status ENUM( + 'pending', + 'started', + 'complete_success', + 'complete_failure', + 'attempt_failed', + 'dead') NOT NULL DEFAULT 'pending', + + -- We can track this against a "max_attempt_count" + attempt_count INT(3) NOT NULL DEFAULT 0, + + -- If there is a failure, tell the user why. + failure_reason VARCHAR(512) DEFAULT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- Failed jobs will set a next attempt time. + -- Subsequent tries can increase the timeout. + -- Failures because of permissions require human intervention => [retry_at=null]. + -- Failures because of invalid files are dead => [status=dead]. + retry_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + UNIQUE KEY (uuid_idempotency_token), + KEY fk_on_download_type (download_type), + KEY fk_creator_user_token (creator_user_token), + KEY index_status (status), + KEY index_creator_ip_address (creator_ip_address) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-08-05-202110_vocoder_models/down.sql b/_migrations/2022-08-05-202110_vocoder_models/down.sql new file mode 100644 index 0000000000..200f3abf59 --- /dev/null +++ b/_migrations/2022-08-05-202110_vocoder_models/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP TABLE vocoder_models; \ No newline at end of file diff --git a/_migrations/2022-08-05-202110_vocoder_models/up.sql b/_migrations/2022-08-05-202110_vocoder_models/up.sql new file mode 100644 index 0000000000..1ee7c74c58 --- /dev/null +++ b/_migrations/2022-08-05-202110_vocoder_models/up.sql @@ -0,0 +1,128 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +CREATE TABLE vocoder_models ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- Effective "primary key" (PUBLIC) + token VARCHAR(32) NOT NULL, + + -- The architecture of the vocoder. + -- Currently supported values: 'hifigan' + vocoder_type VARCHAR(32) NOT NULL, + + -- ========== MUTABLE METADATA ========== + + -- The "name" of the vocoder, which might be complicated. + title VARCHAR(255) NOT NULL, + + -- The description of the model in markdown. + description_markdown TEXT NOT NULL, + + -- Generated HTML (not user-editable). + description_rendered_html TEXT NOT NULL, + + -- ========== MUTABLE MOD-ONLY METADATA ========== + + -- Whether the vocoder is recommended for use + is_staff_recommended BOOLEAN NOT NULL DEFAULT FALSE, + + -- ========== IMMUTABLE PROVENANCE AND METADATA ========== + + -- Where the file was originally downloaded (if it was downloaded) + original_download_url VARCHAR(512) DEFAULT NULL, + + -- The filename that was used at upload time (if available) + original_filename VARCHAR(255) DEFAULT NULL, + + -- File characteristics + file_size_bytes INT(10) NOT NULL DEFAULT 0, + + -- ========== CREATOR DETAILS ========== + + -- The person that created the template. + creator_user_token VARCHAR(32) NOT NULL, + + -- For abuse tracking. + -- Wide enough for IPv4/6 + creator_ip_address_creation VARCHAR(40) NOT NULL, + creator_ip_address_last_update VARCHAR(40) NOT NULL, + + -- (THIS MIGHT NOT BE USED) + -- NB: DO NOT SORT! + -- THIS MUST MATCH THE RESPECTIVE JOBS TABLE. + creator_set_visibility ENUM( + 'public', + 'hidden', + 'private' + ) NOT NULL DEFAULT 'public', + + -- ========== BUCKET STORAGE ========== + + -- The pytorch model + -- For now, this will be a hash of the file contents. + -- NB: NOT UNIQUE! We can allow duplicate uploads. + private_bucket_hash CHAR(64) NOT NULL, + + -- The "full url" of the bucket path, not including the hostname + private_bucket_object_name VARCHAR(255) NOT NULL, + + -- ========== MODERATION DETAILS ========== + + -- Mods have locked the ability for non-authors (ie. people who didn't upload this) + -- to assign the vocoder for use with their TTS models. + -- TTS models that already incorporate the vocoder will not be affected. + -- The vocoder will remain visible. + -- This is independent of other visibility controls (is_mod_disabled_from_author_use, + -- deleted_at, and creator_set_visibility) + is_mod_disabled_from_public_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- Mods have locked the ability for the author (ie. the person that uploaded this) + -- to assign the vocoder for use with their TTS models. + -- TTS models that already incorporate the vocoder will not be affected. + -- The vocoder will remain visible. + -- This is independent of other visibility controls (is_mod_disabled_from_public_use, + -- deleted_at, and creator_set_visibility) + is_mod_disabled_from_author_use BOOLEAN NOT NULL DEFAULT FALSE, + + -- Extremely popular models may be locked from deletion or modification by users. + -- If true, only moderators can edit or delete the vocoder. + is_mod_author_editing_locked BOOLEAN NOT NULL DEFAULT FALSE, + + -- If a moderator has comments. + maybe_mod_comments VARCHAR(255) DEFAULT NULL, + + -- The last moderator that made changes. + maybe_mod_user_token VARCHAR(32) DEFAULT NULL, + + -- ========== VECTOR CLOCK ========== + + -- Incremented with every update. + version INT NOT NULL DEFAULT 0, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- If this is removed by a mod. + -- It completely disappears from the system. + user_deleted_at TIMESTAMP NULL, + mod_deleted_at TIMESTAMP NULL, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token), + KEY fk_creator_user_token (creator_user_token), + KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY index_creator_ip_address_creation (creator_ip_address_creation), + KEY index_creator_ip_address_last_update (creator_ip_address_last_update), + KEY index_creator_set_visibility (creator_set_visibility), + KEY index_private_bucket_hash (private_bucket_hash), + KEY index_is_mod_disabled_from_public_use (is_mod_disabled_from_public_use), + KEY index_is_mod_disabled_from_author_use (is_mod_disabled_from_author_use), + KEY index_is_mod_author_editing_locked (is_mod_author_editing_locked) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/down.sql b/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/down.sql new file mode 100644 index 0000000000..a37b7208ed --- /dev/null +++ b/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/down.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +DROP INDEX `fk_maybe_custom_vocoder_token` ON `tts_models`; + +ALTER TABLE tts_models DROP COLUMN maybe_custom_vocoder_token; diff --git a/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/up.sql b/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/up.sql new file mode 100644 index 0000000000..79dad50a69 --- /dev/null +++ b/_migrations/2022-08-10-055529_alter_table_tts_models_add_vocoder/up.sql @@ -0,0 +1,11 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- If set, we'll use a custom vocoder (from table vocoder_models) instead of a default vocoder. +-- If not set, we'll use `maybe_default_pretrained_vocoder` (or the website default). +ALTER TABLE tts_models +ADD COLUMN maybe_custom_vocoder_token VARCHAR(32) DEFAULT NULL + AFTER maybe_default_pretrained_vocoder; + +CREATE INDEX `fk_maybe_custom_vocoder_token` ON `tts_models` (`maybe_custom_vocoder_token`); diff --git a/_migrations_squashed/tts_tables/tts_models.sql b/_migrations_squashed/tts_tables/tts_models.sql index 04b55aca2a..acaaa5f122 100644 --- a/_migrations_squashed/tts_tables/tts_models.sql +++ b/_migrations_squashed/tts_tables/tts_models.sql @@ -53,6 +53,10 @@ CREATE TABLE tts_models ( -- If not set, use the website default. (Currently HifiGan-SS) maybe_default_pretrained_vocoder VARCHAR(64) DEFAULT NULL, + -- If set, we'll use a custom vocoder (from table vocoder_models) instead of a default vocoder. + -- If not set, we'll use `maybe_default_pretrained_vocoder` (or the website default). + maybe_custom_vocoder_token VARCHAR(32) DEFAULT NULL, + -- Optional Pointer to a newer version of the voice -- If there's a newer version, we can disable this one. -- maybe_updated_model_token VARCHAR(32) DEFAULT NULL, @@ -201,6 +205,7 @@ CREATE TABLE tts_models ( UNIQUE KEY (token), KEY fk_creator_user_token (creator_user_token), KEY fk_maybe_mod_user_token (maybe_mod_user_token), + KEY fk_maybe_custom_vocoder_token (maybe_custom_vocoder_token), KEY index_creator_ip_address_creation (creator_ip_address_creation), KEY index_creator_ip_address_last_update (creator_ip_address_last_update), KEY index_creator_set_visibility (creator_set_visibility), diff --git a/crates/lib/config/src/bad_urls.rs b/crates/lib/config/src/bad_urls.rs index cf4e277130..13be46019b 100644 --- a/crates/lib/config/src/bad_urls.rs +++ b/crates/lib/config/src/bad_urls.rs @@ -1,36 +1,11 @@ //! Probably should be in its own "static_config" crate. -use url::{Host, Url}; +use crate::is_bad_download_url::is_bad_download_url; /// Reports if the URL is inappropriate for downloading as a TTS model, vocoder, etc. pub fn is_bad_tts_model_download_url(url: &str) -> anyhow::Result { - - if url.trim().is_empty() { - // We can't use empty URLs. - return Ok(true); - } - - let url = Url::parse(url)?; - - match url.host() { - Some(Host::Domain(domain)) => { - let domain = domain.to_lowercase(); - - // TODO: If this list grows, store it in a static hashset - let bad_host = domain.contains("fb.watch") - || domain.contains("tiktok.com") - || domain.contains("vm.tiktok.com") // NB: This hostname is known for never disconnecting, which freezes the job - || domain.contains("youtu.be") - || domain.contains("youtube.com"); - - if bad_host { - return Ok(true) - } - } - _ => {}, - } - - Ok(false) + // TODO: Just use this. + is_bad_download_url(url) } #[cfg(test)] diff --git a/crates/lib/config/src/is_bad_download_url.rs b/crates/lib/config/src/is_bad_download_url.rs new file mode 100644 index 0000000000..b3fe096c98 --- /dev/null +++ b/crates/lib/config/src/is_bad_download_url.rs @@ -0,0 +1,53 @@ +//! Probably should be in its own "static_config" crate. + +use url::{Host, Url}; + +/// Reports if the URL is inappropriate for downloading as a TTS model, vocoder, etc. +pub fn is_bad_download_url(url: &str) -> anyhow::Result { + + if url.trim().is_empty() { + // We can't use empty URLs. + return Ok(true); + } + + let url = Url::parse(url)?; + + match url.host() { + Some(Host::Domain(domain)) => { + let domain = domain.to_lowercase(); + + // TODO: If this list grows, store it in a static hashset + let bad_host = domain.contains("fb.watch") + || domain.contains("tiktok.com") + || domain.contains("vm.tiktok.com") // NB: This hostname is known for never disconnecting, which freezes the job + || domain.contains("youtu.be") + || domain.contains("youtube.com"); + + if bad_host { + return Ok(true) + } + } + _ => {}, + } + + Ok(false) +} + +#[cfg(test)] +mod tests { + use crate::is_bad_download_url::is_bad_download_url; + + #[test] + fn bad_download_url() { + assert_eq!(is_bad_download_url("").unwrap(), true); + assert_eq!(is_bad_download_url(" ").unwrap(), true); + assert_eq!(is_bad_download_url("https://vm.tiktok.com/ZMNYjT7Xy/?k=1 ").unwrap(), true); // NB: We get lots of these + assert_eq!(is_bad_download_url("https://m.youtube.com/watch?v=HY-vzGBiAZo").unwrap(), true); // NB: We get lots of these + } + + #[test] + fn good_tts_model_url() { + assert_eq!(is_bad_download_url("https://drive.google.com/file/d/1-1kEoX4HGCwJm4R9cZhSVByWmUoQVGVm/view").unwrap(), false); + assert_eq!(is_bad_download_url("https://drive.google.com/file/d/1SofQhvSkDY-vi_zuBfHivBbJo4CqhJeH/view?usp=sharing").unwrap(), false); + } +} diff --git a/crates/lib/config/src/lib.rs b/crates/lib/config/src/lib.rs index 8952cd1747..8feecf589b 100644 --- a/crates/lib/config/src/lib.rs +++ b/crates/lib/config/src/lib.rs @@ -1,3 +1,4 @@ pub mod bad_urls; pub mod common_env; +pub mod is_bad_download_url; pub mod shared_constants; diff --git a/crates/lib/database_queries/Cargo.toml b/crates/lib/database_queries/Cargo.toml index ab092d51d2..1b2bd703bc 100644 --- a/crates/lib/database_queries/Cargo.toml +++ b/crates/lib/database_queries/Cargo.toml @@ -17,6 +17,7 @@ path = "src/lib.rs" # Internal config = { path = "../config" } container_common = { path = "../container_common" } +reusable_types = { path = "../reusable_types" } # External anyhow = "1.0.40" diff --git a/crates/lib/database_queries/src/mediators/badge_granter.rs b/crates/lib/database_queries/src/mediators/badge_granter.rs index d771744beb..d3f199d60e 100644 --- a/crates/lib/database_queries/src/mediators/badge_granter.rs +++ b/crates/lib/database_queries/src/mediators/badge_granter.rs @@ -12,7 +12,7 @@ pub enum UserBadge { // Granted for early vocodes users. EarlyUser, - // Granted for uploading models + // Granted for uploading tts models TtsModelUploader1, TtsModelUploader5, TtsModelUploader10, @@ -23,6 +23,17 @@ pub enum UserBadge { TtsModelUploader200, TtsModelUploader250, + // Granted for vocoder models + VocoderModelUploader1, + VocoderModelUploader5, + VocoderModelUploader10, + VocoderModelUploader20, + VocoderModelUploader50, + VocoderModelUploader100, + VocoderModelUploader150, + VocoderModelUploader200, + VocoderModelUploader250, + // Granted for uploading templates W2lTemplateUploader1, W2lTemplateUploader10, @@ -63,6 +74,15 @@ impl UserBadge { UserBadge::TtsModelUploader150 => "tts_model_uploader_150", UserBadge::TtsModelUploader200 => "tts_model_uploader_200", UserBadge::TtsModelUploader250 => "tts_model_uploader_250", + UserBadge::VocoderModelUploader1 => "vocoder_model_uploader_1", + UserBadge::VocoderModelUploader5 => "vocoder_model_uploader_5", + UserBadge::VocoderModelUploader10 => "vocoder_model_uploader_10", + UserBadge::VocoderModelUploader20 => "vocoder_model_uploader_20", + UserBadge::VocoderModelUploader50 => "vocoder_model_uploader_50", + UserBadge::VocoderModelUploader100 => "vocoder_model_uploader_100", + UserBadge::VocoderModelUploader150 => "vocoder_model_uploader_150", + UserBadge::VocoderModelUploader200 => "vocoder_model_uploader_200", + UserBadge::VocoderModelUploader250 => "vocoder_model_uploader_250", UserBadge::W2lTemplateUploader1 => "w2l_template_uploader_1", UserBadge::W2lTemplateUploader10 => "w2l_template_uploader_10", UserBadge::W2lTemplateUploader50 => "w2l_template_uploader_50", @@ -206,6 +226,51 @@ impl BadgeGranter { Ok(()) } + /// This needs to be called *after* successful upload. + pub async fn maybe_grant_vocoder_model_uploads_badge(&self, user_token: &str) -> AnyhowResult<()> { + let count = self.count_vocoder_models_uploaded(user_token).await?; + + let mut maybe_badge = None; + + if count >= 250 { + maybe_badge = Some(UserBadge::VocoderModelUploader250); + } else if count >= 200 { + maybe_badge = Some(UserBadge::VocoderModelUploader200); + } else if count >= 150 { + maybe_badge = Some(UserBadge::VocoderModelUploader150); + } else if count >= 100 { + maybe_badge = Some(UserBadge::VocoderModelUploader100); + } else if count >= 50 { + maybe_badge = Some(UserBadge::VocoderModelUploader50); + } else if count >= 20 { + maybe_badge = Some(UserBadge::VocoderModelUploader20); + } else if count >= 10 { + maybe_badge = Some(UserBadge::VocoderModelUploader10); + } else if count >= 5 { + maybe_badge = Some(UserBadge::VocoderModelUploader5); + } else if count >= 1 { + maybe_badge = Some(UserBadge::VocoderModelUploader1); + } + + let badge = match maybe_badge { + Some(badge) => badge, + None => return Ok(()), + }; + + if self.has_badge(user_token, badge).await? { + return Ok(()) + } + + let _record_id = self.insert( + badge, + user_token, + ).await?; + + self.firehose_publisher.publish_user_badge_granted(user_token, badge.to_db_value()) + .await?; + + Ok(()) + } // ======================================================================= @@ -267,6 +332,28 @@ LIMIT 1 self.handle_count_query(maybe_result) } + async fn count_vocoder_models_uploaded( + &self, + user_token: &str, + ) -> AnyhowResult { + // NB: This could get expensive! + let maybe_result = sqlx::query_as!( + CountRecord, + r#" +SELECT count(*) as count +FROM vocoder_models +WHERE + creator_user_token = ? +LIMIT 1 + "#, + user_token + ) + .fetch_one(&self.mysql_pool) + .await; + + self.handle_count_query(maybe_result) + } + async fn count_w2l_templates_uploaded( &self, user_token: &str, diff --git a/crates/lib/database_queries/src/mediators/firehose_publisher.rs b/crates/lib/database_queries/src/mediators/firehose_publisher.rs index 2499a2b144..203554fafa 100644 --- a/crates/lib/database_queries/src/mediators/firehose_publisher.rs +++ b/crates/lib/database_queries/src/mediators/firehose_publisher.rs @@ -24,6 +24,9 @@ enum FirehoseEvent { W2lInferenceStarted, W2lInferenceCompleted, + GenericDownloadStarted, + GenericDownloadCompleted, + TwitterMention, TwitterRetweet, DiscordJoin, @@ -51,6 +54,8 @@ impl FirehoseEvent { FirehoseEvent::DiscordMessage => "discord_message", FirehoseEvent::TwitchSubscribe => "twitch_subscribe", FirehoseEvent::TwitchFollow => "twitch_follow", + FirehoseEvent::GenericDownloadStarted => "generic_download_started", + FirehoseEvent::GenericDownloadCompleted => "generic_download_completed", } } } @@ -172,6 +177,27 @@ impl FirehosePublisher { Ok(()) } + pub async fn enqueue_generic_download(&self, user_token: &str, job_token: &str) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::GenericDownloadStarted, + Some(user_token), + None, + Some(job_token) + ).await?; + Ok(()) + } + + // NB: Entity token is optional. + pub async fn publish_generic_download_finished(&self, user_token: &str, entity_token: Option<&str>) -> AnyhowResult<()> { + let _record_id = self.insert( + FirehoseEvent::GenericDownloadCompleted, + Some(user_token), + entity_token, + entity_token, + ).await?; + Ok(()) + } + // ======================================================================= async fn insert( diff --git a/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs b/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs new file mode 100644 index 0000000000..ae5eba70f2 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/get_generic_download_job_status.rs @@ -0,0 +1,59 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use log::warn; +use sqlx::MySqlPool; +use container_common::anyhow_result::AnyhowResult; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; + +pub struct GenericDownloadJobStatus { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_downloaded_entity_type: Option, + pub maybe_downloaded_entity_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +/// Look up job status. +/// Returns Ok(None) when the record cannot be found. +pub async fn get_generic_download_job_status(job_token: &str, mysql_pool: &MySqlPool) + -> AnyhowResult> +{ + let maybe_status = sqlx::query_as!( + GenericDownloadJobStatus, + r#" +SELECT + jobs.token as job_token, + + jobs.status, + jobs.attempt_count, + jobs.on_success_downloaded_entity_token as maybe_downloaded_entity_token, + jobs.on_success_downloaded_entity_type as maybe_downloaded_entity_type, + + jobs.created_at, + jobs.updated_at + +FROM generic_download_jobs as jobs + +WHERE jobs.token = ? + "#, + job_token + ) + .fetch_one(mysql_pool) + .await; + + match maybe_status { + Ok(record) => Ok(Some(record)), + Err(err) => match err { + sqlx::Error::RowNotFound => Ok(None), + _ => { + warn!("error querying job record: {:?}", err); + Err(anyhow!("error querying job record: {:?}", err)) + } + } + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs b/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs new file mode 100644 index 0000000000..ed9c936022 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/insert_generic_download_job.rs @@ -0,0 +1,58 @@ +use anyhow::anyhow; +use sqlx::MySqlPool; +use container_common::anyhow_result::AnyhowResult; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; + +pub struct Args <'a> { + pub job_token: &'a str, + pub uuid_idempotency_token: &'a str, + pub download_type: GenericDownloadType, + pub download_url: &'a str, + pub title: &'a str, + pub creator_user_token: &'a str, + pub creator_ip_address: &'a str, + pub creator_set_visibility: EntityVisibility, + pub mysql_pool: &'a MySqlPool, +} + +pub async fn insert_generic_download_job(args: Args<'_>) -> AnyhowResult { + + let query = sqlx::query!( + r#" +INSERT INTO generic_download_jobs +SET + token = ?, + uuid_idempotency_token = ?, + download_type = ?, + download_url = ?, + title = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + status = "pending" + "#, + args.job_token, + args.uuid_idempotency_token, + args.download_type.to_str(), + args.download_url, + args.title, + args.creator_user_token, + args.creator_ip_address, + args.creator_set_visibility.to_str(), + ); + + let query_result = query.execute(args.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("error inserting new generic download job: {:?}", err)); + } + }; + + Ok(record_id) +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs b/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs new file mode 100644 index 0000000000..a2ade0e9f9 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/_keys.rs @@ -0,0 +1,5 @@ +//! Primary and foreign keys + +/// Since jobs are internal, we'll use IDs as primary keys. +#[derive(Clone, Copy, PartialOrd, PartialEq, Eq, Ord, Debug)] +pub struct GenericDownloadJobId(pub i64); diff --git a/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs b/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs new file mode 100644 index 0000000000..ddbaeb86a8 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/list_available_generic_download_jobs.rs @@ -0,0 +1,123 @@ +use anyhow::anyhow; +use chrono::Utc; +use container_common::anyhow_result::AnyhowResult; +use crate::column_types::job_status::JobStatus; +use crate::queries::generic_download::job::_keys::GenericDownloadJobId; +use crate::tokens::Tokens; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; +use sqlx::MySqlPool; +use std::path::Path; + +/// table: generic_download_jobs +#[derive(Debug)] +pub struct AvailableDownloadJob { + pub id: GenericDownloadJobId, + pub download_job_token: String, + + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: EntityVisibility, + + pub download_type: GenericDownloadType, + pub download_url: String, + pub title: String, + + pub status: JobStatus, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} + +pub async fn list_available_generic_download_jobs(pool: &MySqlPool, num_records: u32) + -> AnyhowResult> +{ + let job_records = sqlx::query_as!( + AvailableDownloadJobRawInternal, + r#" +SELECT + id, + token AS download_job_token, + + creator_user_token, + creator_ip_address, + creator_set_visibility as `creator_set_visibility: reusable_types::entity_visibility::EntityVisibility`, + + download_type as `download_type: reusable_types::generic_download_type::GenericDownloadType`, + download_url, + title, + + status as `status: crate::column_types::job_status::JobStatus`, + attempt_count, + failure_reason, + + created_at, + updated_at, + retry_at +FROM generic_download_jobs +WHERE + ( + status IN ("pending", "attempt_failed") + ) + AND + ( + retry_at IS NULL + OR + retry_at < CURRENT_TIMESTAMP + ) + ORDER BY id ASC + LIMIT ? + "#, + num_records, + ) + .fetch_all(pool) + .await?; + + let job_records = job_records.into_iter() + .map(|record : AvailableDownloadJobRawInternal| { + AvailableDownloadJob { + id: GenericDownloadJobId(record.id), + download_job_token: record.download_job_token, + creator_ip_address: record.creator_ip_address, + creator_user_token: record.creator_user_token, + creator_set_visibility: record.creator_set_visibility, + download_type: record.download_type, + download_url: record.download_url, + title: record.title, + status: record.status, + attempt_count: record.attempt_count, + failure_reason: record.failure_reason, + created_at: record.created_at, + updated_at: record.updated_at, + retry_at: record.retry_at, + } + }) + .collect::>(); + + Ok(job_records) +} + +#[derive(Debug)] +struct AvailableDownloadJobRawInternal { + pub id: i64, + pub download_job_token: String, + + pub creator_user_token: String, + pub creator_ip_address: String, + pub creator_set_visibility: EntityVisibility, + + pub download_type: GenericDownloadType, + pub download_url: String, + pub title: String, + + pub status: JobStatus, + pub attempt_count: i32, + pub failure_reason: Option, + + pub created_at: chrono::DateTime, + pub updated_at: chrono::DateTime, + pub retry_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs new file mode 100644 index 0000000000..fdc6de7e67 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_done.rs @@ -0,0 +1,39 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use sqlx::MySqlPool; + +pub async fn mark_generic_download_job_done( + pool: &MySqlPool, + job: &AvailableDownloadJob, + success: bool, + maybe_entity_token: Option<&str>, + maybe_entity_type: Option<&str>, +) -> AnyhowResult<()> +{ + let status = if success { "complete_success" } else { "complete_failure" }; + + let query_result = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = ?, + on_success_downloaded_entity_token = ?, + on_success_downloaded_entity_type = ?, + failure_reason = NULL, + retry_at = NULL +WHERE id = ? + "#, + status, + maybe_entity_token, + maybe_entity_type, + job.id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => Err(anyhow!("error with query: {:?}", err)), + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs new file mode 100644 index 0000000000..b6de45c92c --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_failure.rs @@ -0,0 +1,45 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use sqlx::MySqlPool; +use sqlx; + +/// Mark a single download job failure. The job may be re-run. +pub async fn mark_generic_download_job_failure( + pool: &MySqlPool, + job: &AvailableDownloadJob, + failure_reason: &str, + max_attempts: i32 +) -> AnyhowResult<()> { + + // statuses: "attempt_failed", "complete_failure", "dead" + let mut next_status = "attempt_failed"; + + if job.attempt_count >= max_attempts { + // NB: Job attempt count is incremented at start + next_status = "dead"; + } + + let query_result = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = ?, + failure_reason = ?, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + next_status, + failure_reason.to_string(), + job.id.0, + ) + .execute(pool) + .await; + + match query_result { + Err(err) => { + Err(anyhow!("error with query: {:?}", err)) + }, + Ok(_r) => Ok(()), + } +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs new file mode 100644 index 0000000000..be5ab74789 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mark_generic_download_job_pending_and_grab_lock.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use sqlx::MySqlPool; +use sqlx; +use crate::queries::generic_download::job::_keys::GenericDownloadJobId; + +// TODO(2022-02-09): It's unclear how stuck "pending" jobs don't become perma-dead +// I think this logic is flawed and the job system needs to pick up stuck "pending" jobs. + +pub struct GenericDownloadJobLockRecord { + id: i64, + status: String, +} + +pub async fn mark_generic_download_job_pending_and_grab_lock( + pool: &MySqlPool, + job_id: GenericDownloadJobId +) -> AnyhowResult { + + // NB: We use transactions and "SELECT ... FOR UPDATE" to simulate mutexes. + let mut transaction = pool.begin().await?; + + let maybe_record = sqlx::query_as!( + GenericDownloadJobLockRecord, + r#" +SELECT + id, + status +FROM generic_download_jobs +WHERE id = ? +FOR UPDATE + "#, + job_id.0, + ) + .fetch_one(&mut transaction) + .await; + + let record : GenericDownloadJobLockRecord = match maybe_record { + Ok(record) => record, + Err(err) => { + match err { + sqlx::Error::RowNotFound => { + return Err(anyhow!("could not job")); + }, + _ => { + return Err(anyhow!("query error")); + } + } + } + }; + + let can_transact = match record.status.as_ref() { + "pending" => true, // It's okay for us to take the lock. + "attempt_failed" => true, // We can retry. + "started" => false, // Job in progress (another job beat us, and we can't take the lock) + "complete_success" => false, // Job already complete + "complete_failure" => false, // Job already complete (permanently dead; no need to retry) + "dead" => false, // Job already complete (permanently dead; retries exhausted) + _ => false, // Future-proof + }; + + if !can_transact { + transaction.rollback().await?; + return Ok(false); + } + + let _acquire_lock = sqlx::query!( + r#" +UPDATE generic_download_jobs +SET + status = 'started', + attempt_count = attempt_count + 1, + retry_at = NOW() + interval 2 minute +WHERE id = ? + "#, + job_id.0, + ) + .execute(&mut transaction) + .await?; + + transaction.commit().await?; + + Ok(true) +} diff --git a/crates/lib/database_queries/src/queries/generic_download/job/mod.rs b/crates/lib/database_queries/src/queries/generic_download/job/mod.rs new file mode 100644 index 0000000000..ac425ffd26 --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/job/mod.rs @@ -0,0 +1,5 @@ +pub mod _keys; +pub mod list_available_generic_download_jobs; +pub mod mark_generic_download_job_done; +pub mod mark_generic_download_job_failure; +pub mod mark_generic_download_job_pending_and_grab_lock; \ No newline at end of file diff --git a/crates/lib/database_queries/src/queries/generic_download/mod.rs b/crates/lib/database_queries/src/queries/generic_download/mod.rs new file mode 100644 index 0000000000..4feea997bc --- /dev/null +++ b/crates/lib/database_queries/src/queries/generic_download/mod.rs @@ -0,0 +1,3 @@ +pub mod get_generic_download_job_status; +pub mod insert_generic_download_job; +pub mod job; diff --git a/crates/lib/database_queries/src/queries/mod.rs b/crates/lib/database_queries/src/queries/mod.rs index eb9ea3bf70..9c75dc80b4 100644 --- a/crates/lib/database_queries/src/queries/mod.rs +++ b/crates/lib/database_queries/src/queries/mod.rs @@ -1,9 +1,11 @@ pub mod api_tokens; +pub mod generic_download; pub mod health_check; pub mod ip_bans; pub mod model_categories; pub mod tts; pub mod twitch; pub mod users; +pub mod vocoder; pub mod voice_clone_requests; pub mod w2l; diff --git a/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs b/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs index 558f800cfd..e99c1e1317 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_download_jobs/tts_download_job_queries.rs @@ -1,7 +1,3 @@ -//! NB: This seems required for sqlx to generate the cached queries. -//! Sqlx's prepare needs a *single* binary to work against, so we need to -//! include these in the main binary to generate all the queries. - use anyhow::anyhow; use chrono::Utc; use container_common::anyhow_result::AnyhowResult; diff --git a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs index 6d9286b5cc..96e14dcc01 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_inference_jobs/list_available_tts_inference_jobs.rs @@ -1,11 +1,14 @@ -use anyhow::anyhow; +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use chrono::{Utc, DateTime}; use container_common::anyhow_result::AnyhowResult; use crate::column_types::job_status::JobStatus; use crate::column_types::record_visibility::RecordVisibility; use crate::helpers::boolean_converters::i8_to_bool; use crate::queries::tts::tts_inference_jobs::_keys::TtsInferenceJobId; -use log::{warn, info}; use sqlx::MySqlPool; /// table: tts_inference_jobs @@ -31,9 +34,9 @@ pub struct AvailableTtsInferenceJob { pub attempt_count: i32, pub failure_reason: Option, - pub created_at: chrono::DateTime, - pub updated_at: chrono::DateTime, - pub retry_at: Option>, + pub created_at: DateTime, + pub updated_at: DateTime, + pub retry_at: Option>, } /// Query jobs that are ready to run diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/edit_tts_model_details.rs b/crates/lib/database_queries/src/queries/tts/tts_models/edit_tts_model_details.rs index ed539f2854..3d850e2720 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_models/edit_tts_model_details.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_models/edit_tts_model_details.rs @@ -14,6 +14,7 @@ pub async fn edit_tts_model_details_as_author( ietf_primary_language_subtag: &str, creator_set_visibility: RecordVisibility, maybe_default_pretrained_vocoder: Option, + maybe_custom_vocoder_token: Option<&str>, text_pipeline_type: Option<&str>, // Author fields author_ip_address: &str, @@ -24,6 +25,7 @@ pub async fn edit_tts_model_details_as_author( UPDATE tts_models SET maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ?, text_pipeline_type = ?, title = ?, description_markdown = ?, @@ -37,6 +39,7 @@ WHERE token = ? LIMIT 1 "#, maybe_default_pretrained_vocoder.map(|v| v.to_str()), + maybe_custom_vocoder_token, text_pipeline_type, title, description_markdown, @@ -61,6 +64,7 @@ pub async fn edit_tts_model_details_as_mod( ietf_primary_language_subtag: &str, creator_set_visibility: RecordVisibility, maybe_default_pretrained_vocoder: Option, + maybe_custom_vocoder_token: Option<&str>, text_pipeline_type: Option<&str>, // moderator fields moderator_user_token: &str, @@ -71,6 +75,7 @@ pub async fn edit_tts_model_details_as_mod( UPDATE tts_models SET maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ?, text_pipeline_type = ?, title = ?, description_markdown = ?, @@ -84,6 +89,7 @@ WHERE token = ? LIMIT 1 "#, maybe_default_pretrained_vocoder.map(|v| v.to_str()), + maybe_custom_vocoder_token, text_pipeline_type, title, description_markdown, diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs index 63e80ac80b..70933e4478 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model.rs @@ -1,3 +1,8 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use anyhow::anyhow; use chrono::{DateTime, Utc}; use container_common::anyhow_result::AnyhowResult; @@ -16,10 +21,11 @@ pub struct TtsModelRecord { /// NB: text_pipeline_type may not always be present in the database. pub text_pipeline_type: Option, - - pub maybe_default_pretrained_vocoder: Option, pub text_preprocessing_algorithm: String, + pub maybe_custom_vocoder: Option, + pub maybe_default_pretrained_vocoder: Option, + pub creator_user_token: String, pub creator_username: String, pub creator_display_name: String, @@ -48,6 +54,15 @@ pub struct TtsModelRecord { pub maybe_moderator_fields: Option, } +pub struct CustomVocoderFields { + pub vocoder_token: String, + pub vocoder_title: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, +} + /// "Moderator-only fields" that we wouldn't want to expose to ordinary users. /// It's the web endpoint controller's responsibility to clear these for non-mods. pub struct TtsModelModeratorFields { @@ -98,8 +113,21 @@ pub async fn get_tts_model_by_token( model_token: model.model_token, tts_model_type: model.tts_model_type, text_pipeline_type: model.text_pipeline_type, - maybe_default_pretrained_vocoder: maybe_vocoder, text_preprocessing_algorithm: model.text_preprocessing_algorithm, + maybe_default_pretrained_vocoder: maybe_vocoder, + maybe_custom_vocoder: match model.maybe_custom_vocoder_token { + // NB: We're relying on a single field's presence to infer that the others vocoder fields + // are also there. If for some reason they aren't, fail open. + None => None, + Some(vocoder_token) => Some(CustomVocoderFields { + vocoder_token, + vocoder_title: model.maybe_custom_vocoder_title.unwrap_or("".to_string()), + creator_user_token: model.maybe_custom_vocoder_creator_user_token.unwrap_or("".to_string()), + creator_username: model.maybe_custom_vocoder_creator_username.unwrap_or("".to_string()), + creator_display_name: model.maybe_custom_vocoder_creator_display_name.unwrap_or("".to_string()), + creator_gravatar_hash: model.maybe_custom_vocoder_creator_gravatar_hash.unwrap_or("".to_string()), + }) + }, creator_user_token: model.creator_user_token, creator_username: model.creator_username, creator_display_name: model.creator_display_name, @@ -168,6 +196,13 @@ SELECT tts.is_locked_from_use, tts.is_locked_from_user_modification, + tts.maybe_custom_vocoder_token, + vocoder.title as maybe_custom_vocoder_title, + vocoder_user.token as maybe_custom_vocoder_creator_user_token, + vocoder_user.username as maybe_custom_vocoder_creator_username, + vocoder_user.display_name as maybe_custom_vocoder_creator_display_name, + vocoder_user.email_gravatar_hash as maybe_custom_vocoder_creator_gravatar_hash, + tts.created_at, tts.updated_at, @@ -179,6 +214,12 @@ SELECT FROM tts_models as tts JOIN users ON users.token = tts.creator_user_token + +LEFT OUTER JOIN vocoder_models AS vocoder + ON vocoder.token = tts.maybe_custom_vocoder_token +LEFT OUTER JOIN users AS vocoder_user + ON vocoder_user.token = vocoder.creator_user_token + WHERE tts.token = ? "#, tts_model_token @@ -224,6 +265,13 @@ SELECT tts.is_locked_from_use, tts.is_locked_from_user_modification, + tts.maybe_custom_vocoder_token, + vocoder.title as maybe_custom_vocoder_title, + vocoder_user.token as maybe_custom_vocoder_creator_user_token, + vocoder_user.username as maybe_custom_vocoder_creator_username, + vocoder_user.display_name as maybe_custom_vocoder_creator_display_name, + vocoder_user.email_gravatar_hash as maybe_custom_vocoder_creator_gravatar_hash, + tts.created_at, tts.updated_at, @@ -235,6 +283,12 @@ SELECT FROM tts_models as tts JOIN users ON users.token = tts.creator_user_token + +LEFT OUTER JOIN vocoder_models AS vocoder + ON vocoder.token = tts.maybe_custom_vocoder_token +LEFT OUTER JOIN users AS vocoder_user + ON vocoder_user.token = vocoder.creator_user_token + WHERE tts.token = ? AND tts.user_deleted_at IS NULL @@ -280,6 +334,14 @@ struct InternalTtsModelRecordRaw { pub created_at: DateTime, pub updated_at: DateTime, + // Joined custom vocoder fields + pub maybe_custom_vocoder_token: Option, + pub maybe_custom_vocoder_title: Option, + pub maybe_custom_vocoder_creator_user_token: Option, + pub maybe_custom_vocoder_creator_username: Option, + pub maybe_custom_vocoder_creator_display_name: Option, + pub maybe_custom_vocoder_creator_gravatar_hash: Option, + // Moderator fields pub creator_ip_address_creation: String, pub creator_ip_address_last_update: String, diff --git a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs index c12910499a..f3e840e224 100644 --- a/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs +++ b/crates/lib/database_queries/src/queries/tts/tts_models/get_tts_model_for_inference.rs @@ -1,8 +1,13 @@ -use anyhow::anyhow; +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use chrono::{Utc, DateTime}; use log::warn; use sqlx::MySqlPool; use sqlx; +use reusable_types::vocoder_type::VocoderType; // TODO: Can probably just reuse another query. @@ -14,6 +19,18 @@ pub struct TtsModelForInferenceRecord { /// NB: text_pipeline_type may not always be present in the database. pub text_pipeline_type: Option, + /// [vocoders 1] + /// This is the new type of vocoder configuration. Users can choose a custom trained + /// vocoder to associate with their model. The tokens reference the `vocoder_models` + /// table. + pub maybe_custom_vocoder: Option, + + /// [vocoders 2] + /// This is the old type of vocoder configuration, which leverages old pretrained + /// vocoders that we manually uploaded. There aren't many good options for users to + /// choose here, so this should be treated as a legacy option going forward. We'll + /// likely be stuck with this configuration for some time, however, due to the large + /// collection of legacy models we have. pub maybe_default_pretrained_vocoder: Option, pub creator_user_token: String, @@ -28,6 +45,13 @@ pub struct TtsModelForInferenceRecord { pub user_deleted_at: Option>, } +#[derive(Clone)] +pub struct CustomVocoderFields { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + pub vocoder_title: String, + pub vocoder_private_bucket_hash: String, +} #[derive(Clone, Debug)] pub enum TtsModelForInferenceError { @@ -57,13 +81,20 @@ pub async fn get_tts_model_for_inference( // NB: Lookup failure is Err(RowNotFound). // NB: Since this is publicly exposed, we don't query sensitive data. let maybe_model = sqlx::query_as!( - TtsModelForInferenceRecord, + InternalTtsModelForInferenceRecordRaw, r#" SELECT tts.token as model_token, tts.tts_model_type, tts.text_pipeline_type, + tts.maybe_default_pretrained_vocoder, + + tts.maybe_custom_vocoder_token, + vocoder.vocoder_type as `maybe_custom_vocoder_type: reusable_types::vocoder_type::VocoderType`, + vocoder.title as maybe_custom_vocoder_title, + vocoder.private_bucket_hash as maybe_custom_vocoder_private_bucket_hash, + tts.creator_user_token, users.username as creator_username, users.display_name as creator_display_name, @@ -73,9 +104,15 @@ SELECT tts.updated_at, tts.user_deleted_at, tts.mod_deleted_at + FROM tts_models as tts + JOIN users -ON users.token = tts.creator_user_token + ON users.token = tts.creator_user_token + +LEFT OUTER JOIN vocoder_models AS vocoder + ON vocoder.token = tts.maybe_custom_vocoder_token + WHERE tts.token = ? "#, &model_token @@ -83,7 +120,7 @@ WHERE tts.token = ? .fetch_one(pool) .await; // TODO: This will return error if it doesn't exist - let model : TtsModelForInferenceRecord = match maybe_model { + let model : InternalTtsModelForInferenceRecordRaw = match maybe_model { Ok(model) => model, Err(err) => { match err { @@ -104,5 +141,61 @@ WHERE tts.token = ? return Err(TtsModelForInferenceError::ModelDeleted); } - Ok(model) + Ok(TtsModelForInferenceRecord { + model_token: model.model_token, + tts_model_type: model.tts_model_type, + text_pipeline_type: model.text_pipeline_type, + maybe_custom_vocoder: match model.maybe_custom_vocoder_token { + // NB: We're relying on a single field's presence to infer that the others vocoder fields + // are also there. If for some reason they aren't, fail open. + None => None, + Some(vocoder_token) => Some(CustomVocoderFields { + vocoder_token, + vocoder_type: model.maybe_custom_vocoder_type.ok_or( + TtsModelForInferenceError::DatabaseError { reason: "custom_vocoder_type field error".to_string() })?, + vocoder_title: model.maybe_custom_vocoder_title.ok_or( + TtsModelForInferenceError::DatabaseError { reason: "custom_vocoder_title field error".to_string() })?, + vocoder_private_bucket_hash: model.maybe_custom_vocoder_private_bucket_hash.ok_or( + TtsModelForInferenceError::DatabaseError { reason: "vocoder_private_bucket_hash field error".to_string() })?, + }) + }, + maybe_default_pretrained_vocoder: model.maybe_default_pretrained_vocoder, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + title: model.title, + private_bucket_hash: model.private_bucket_hash, + created_at: model.created_at, + updated_at: model.updated_at, + mod_deleted_at: model.mod_deleted_at, + user_deleted_at: model.user_deleted_at, + }) } + +struct InternalTtsModelForInferenceRecordRaw { + model_token: String, + tts_model_type: String, + + text_pipeline_type: Option, + + // Joined custom vocoder fields + maybe_custom_vocoder_token: Option, + maybe_custom_vocoder_type: Option, + maybe_custom_vocoder_title: Option, + maybe_custom_vocoder_private_bucket_hash: Option, + + // Legacy vocoder config + maybe_default_pretrained_vocoder: Option, + + creator_user_token: String, + creator_username: String, + creator_display_name: String, + + title: String, + private_bucket_hash: String, + created_at: DateTime, + updated_at: DateTime, + mod_deleted_at: Option>, + user_deleted_at: Option>, +} + diff --git a/crates/lib/database_queries/src/queries/vocoder/get_vocoder_model.rs b/crates/lib/database_queries/src/queries/vocoder/get_vocoder_model.rs new file mode 100644 index 0000000000..9ad4683a8b --- /dev/null +++ b/crates/lib/database_queries/src/queries/vocoder/get_vocoder_model.rs @@ -0,0 +1,243 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::warn; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::vocoder_type::VocoderType; +use sqlx::MySqlPool; + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// The moderator-only fields are good practice, though. + +// NB: Do not publicly expose this type. +pub struct VocoderModelRecord { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + + pub is_staff_recommended: bool, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub creator_set_visibility: EntityVisibility, + + pub created_at: DateTime, + pub updated_at: DateTime, + + pub maybe_moderator_fields: Option, +} + +// NB: Do not publicly expose this type. +/// "Moderator-only fields" that we wouldn't want to expose to ordinary users. +/// It's the web endpoint controller's responsibility to clear these for non-mods. +pub struct VocoderModelRecordModFields { + pub creator_is_banned: bool, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + + pub is_mod_disabled_from_public_use: bool, + pub is_mod_disabled_from_author_use: bool, + pub is_mod_author_editing_locked: bool, + + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// The moderator-only fields are good practice, though. + +pub async fn get_vocoder_model_by_token( + vocoder_model_token: &str, + can_see_deleted: bool, + mysql_pool: &MySqlPool +) -> AnyhowResult> { + + let maybe_record = if can_see_deleted { + select_including_deleted(vocoder_model_token, mysql_pool).await + } else { + select_without_deleted(vocoder_model_token, mysql_pool).await + }; + + let model : InternalVocoderModelRecordRaw = match maybe_record { + Ok(model) => model, + Err(ref err) => match err { + sqlx::Error::RowNotFound => { + warn!("vocoder model not found: {:?}", &err); + return Ok(None); + }, + _ => { + warn!("vocoder model query error: {:?}", &err); + return Err(anyhow!("database error")); + } + } + }; + + Ok(Some(VocoderModelRecord { + vocoder_token: model.vocoder_token, + vocoder_type: model.vocoder_type, + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + title: model.title, + description_markdown: model.description_markdown, + description_rendered_html: model.description_rendered_html, + is_staff_recommended: i8_to_bool(model.is_staff_recommended), + creator_set_visibility: model.creator_set_visibility, + created_at: model.created_at, + updated_at: model.updated_at, + maybe_moderator_fields: Some(VocoderModelRecordModFields { + creator_is_banned: i8_to_bool(model.creator_is_banned), + creator_ip_address_creation: model.creator_ip_address_creation, + creator_ip_address_last_update: model.creator_ip_address_last_update, + is_mod_disabled_from_public_use: i8_to_bool(model.is_mod_disabled_from_public_use), + is_mod_disabled_from_author_use: i8_to_bool(model.is_mod_disabled_from_author_use), + is_mod_author_editing_locked: i8_to_bool(model.is_mod_author_editing_locked), + user_deleted_at: model.user_deleted_at, + mod_deleted_at: model.mod_deleted_at, + }), + })) +} + +async fn select_including_deleted( + vocoder_model_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + InternalVocoderModelRecordRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + users.is_banned as creator_is_banned, + + vocoder.title, + vocoder.description_markdown, + vocoder.description_rendered_html, + vocoder.is_staff_recommended, + + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + + vocoder.creator_set_visibility as `creator_set_visibility: reusable_types::entity_visibility::EntityVisibility`, + + vocoder.created_at, + vocoder.updated_at, + + vocoder.creator_ip_address_creation, + vocoder.creator_ip_address_last_update, + vocoder.mod_deleted_at, + vocoder.user_deleted_at + +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token +WHERE + vocoder.token = ? + "#, + vocoder_model_token + ) + .fetch_one(mysql_pool) + .await +} + +async fn select_without_deleted( + vocoder_model_token: &str, + mysql_pool: &MySqlPool +) -> Result { + sqlx::query_as!( + InternalVocoderModelRecordRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + users.is_banned as creator_is_banned, + + vocoder.title, + vocoder.description_markdown, + vocoder.description_rendered_html, + vocoder.is_staff_recommended, + + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + + vocoder.creator_set_visibility as `creator_set_visibility: reusable_types::entity_visibility::EntityVisibility`, + + vocoder.created_at, + vocoder.updated_at, + + vocoder.creator_ip_address_creation, + vocoder.creator_ip_address_last_update, + vocoder.mod_deleted_at, + vocoder.user_deleted_at + +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token +WHERE + vocoder.token = ? + AND vocoder.user_deleted_at IS NULL + AND vocoder.mod_deleted_at IS NULL + "#, + vocoder_model_token + ) + .fetch_one(mysql_pool) + .await +} + +struct InternalVocoderModelRecordRaw { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + + pub is_staff_recommended: i8, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub creator_set_visibility: EntityVisibility, + + pub created_at: DateTime, + pub updated_at: DateTime, + + // Moderator fields + pub creator_is_banned: i8, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + + pub is_mod_disabled_from_public_use: i8, + pub is_mod_disabled_from_author_use: i8, + pub is_mod_author_editing_locked: i8, + + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/vocoder/insert_vocoder_model.rs b/crates/lib/database_queries/src/queries/vocoder/insert_vocoder_model.rs new file mode 100644 index 0000000000..f6152981a5 --- /dev/null +++ b/crates/lib/database_queries/src/queries/vocoder/insert_vocoder_model.rs @@ -0,0 +1,89 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use crate::tokens::Tokens; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::vocoder_type::VocoderType; +use sqlx::MySqlPool; +use std::path::Path; + +pub struct Args<'a, P: AsRef> { + pub vocoder_type: VocoderType, + + pub title: &'a str, + + pub original_download_url: &'a str, + pub original_filename: &'a str, + pub file_size_bytes: u64, + + pub creator_user_token: &'a str, + pub creator_ip_address: &'a str, + pub creator_set_visibility: EntityVisibility, + + pub private_bucket_hash: &'a str, + pub private_bucket_object_name: P, + + pub mysql_pool: &'a MySqlPool, +} + +pub async fn insert_vocoder_model>( + args: Args<'_, P>, +) -> AnyhowResult<(u64, String)> { + + let model_token = Tokens::new_vocoder_model()?; + + let private_bucket_object_name = &args.private_bucket_object_name + .as_ref() + .display() + .to_string(); + + let query_result = sqlx::query!( + r#" +INSERT INTO vocoder_models +SET + token = ?, + vocoder_type = ?, + title = ?, + description_markdown = '', + description_rendered_html = '', + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + creator_user_token = ?, + creator_ip_address_creation = ?, + creator_ip_address_last_update = ?, + creator_set_visibility = ?, + private_bucket_hash = ?, + private_bucket_object_name = ? + "#, + &model_token, + args.vocoder_type.to_str(), + &args.title, + &args.original_download_url, + &args.original_filename, + args.file_size_bytes, + &args.creator_user_token, + &args.creator_ip_address, + &args.creator_ip_address, + args.creator_set_visibility.to_str(), + args.private_bucket_hash, + private_bucket_object_name, + ) + .execute(args.mysql_pool) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + Ok((record_id, model_token)) +} diff --git a/crates/lib/database_queries/src/queries/vocoder/list_vocoder_models.rs b/crates/lib/database_queries/src/queries/vocoder/list_vocoder_models.rs new file mode 100644 index 0000000000..728eab0b95 --- /dev/null +++ b/crates/lib/database_queries/src/queries/vocoder/list_vocoder_models.rs @@ -0,0 +1,273 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use container_common::anyhow_result::AnyhowResult; +use crate::helpers::boolean_converters::i8_to_bool; +use log::{warn, info}; +use reusable_types::vocoder_type::VocoderType; +use sqlx::MySqlPool; + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// But I'm in a hurry... + +// NB: Do not publicly expose this type. +pub struct VocoderModelListItem { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub is_staff_recommended: bool, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub created_at: DateTime, + pub updated_at: DateTime, + + // Moderator fields + pub moderator_fields: VocoderModelListItemModFields, +} + +// NB: Do not publicly expose this type. +pub struct VocoderModelListItemModFields { + pub is_mod_disabled_from_public_use: bool, + pub is_mod_disabled_from_author_use: bool, + pub is_mod_author_editing_locked: bool, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +// FIXME: This is the old style of query scoping and shouldn't be copied. +// But I'm in a hurry... + +pub async fn list_vocoder_models( + mysql_pool: &MySqlPool, + scope_creator_username: Option<&str>, + can_see_deleted_and_hidden: bool, +) -> AnyhowResult> { + + let maybe_models = match scope_creator_username { + Some(username) => { + list_vocoder_models_creator_scoped(mysql_pool, username, can_see_deleted_and_hidden) + .await + }, + None => { + list_vocoder_models_for_all_creators(mysql_pool, can_see_deleted_and_hidden) + .await + }, + }; + + let models : Vec = match maybe_models { + Ok(models) => models, + Err(err) => return match err { + _RowNotFound => { + Ok(Vec::new()) + }, + _ => { + warn!("vocoder model list query error: {:?}", err); + Err(anyhow!("vocoder model list query error")) + } + } + }; + + Ok(models.into_iter() + .map(|model| { + VocoderModelListItem { + vocoder_token: model.vocoder_token, + vocoder_type: model.vocoder_type, + title: model.title, + is_staff_recommended: i8_to_bool(model.is_staff_recommended), + creator_user_token: model.creator_user_token, + creator_username: model.creator_username, + creator_display_name: model.creator_display_name, + creator_gravatar_hash: model.creator_gravatar_hash, + created_at: model.created_at, + updated_at: model.updated_at, + moderator_fields: VocoderModelListItemModFields { + is_mod_disabled_from_public_use: i8_to_bool(model.is_mod_disabled_from_public_use), + is_mod_disabled_from_author_use: i8_to_bool(model.is_mod_disabled_from_author_use), + is_mod_author_editing_locked: i8_to_bool(model.is_mod_author_editing_locked), + user_deleted_at: model.user_deleted_at, + mod_deleted_at: model.mod_deleted_at, + } + } + }) + .collect::>()) +} + +async fn list_vocoder_models_for_all_creators( + mysql_pool: &MySqlPool, + can_see_deleted_and_hidden: bool +) -> AnyhowResult> { + // NB: Sqlx doesn't like non-string literal queries, so we have to branch. + // Repeating the query is ugly, but that's the best we can do without a + // full-on query builder pattern. + let maybe_models = if !can_see_deleted_and_hidden { + info!("listing vocoder models for everyone; publicly visible only"); + sqlx::query_as!( + InternalVocoderModelListItemRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + vocoder.title, + vocoder.is_staff_recommended, + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + vocoder.created_at, + vocoder.updated_at, + vocoder.user_deleted_at, + vocoder.mod_deleted_at +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token +WHERE + vocoder.is_mod_disabled_from_public_use IS FALSE + AND vocoder.is_mod_disabled_from_author_use IS FALSE + AND vocoder.user_deleted_at IS NULL + AND vocoder.mod_deleted_at IS NULL + "#) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing tts models for everyone; all visibility states (undeleted only)"); + sqlx::query_as!( + InternalVocoderModelListItemRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + vocoder.title, + vocoder.is_staff_recommended, + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + vocoder.created_at, + vocoder.updated_at, + vocoder.user_deleted_at, + vocoder.mod_deleted_at +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token + "#) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_models) +} + +async fn list_vocoder_models_creator_scoped( + mysql_pool: &MySqlPool, + scope_creator_username: &str, + can_see_deleted_and_hidden: bool +) -> AnyhowResult> { + // NB: Sqlx doesn't like non-string literal queries, so we have to branch. + // Repeating the query is ugly, but that's the best we can do without a + // full-on query builder pattern. + let maybe_models = if !can_see_deleted_and_hidden { + info!("listing tts models for user; mod-approved only"); + sqlx::query_as!( + InternalVocoderModelListItemRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + vocoder.title, + vocoder.is_staff_recommended, + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + vocoder.created_at, + vocoder.updated_at, + vocoder.user_deleted_at, + vocoder.mod_deleted_at +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token +WHERE + users.username = ? + AND vocoder.is_mod_disabled_from_public_use IS FALSE + AND vocoder.is_mod_disabled_from_author_use IS FALSE + AND vocoder.user_deleted_at IS NULL + AND vocoder.mod_deleted_at IS NULL + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + } else { + info!("listing tts models for user; all"); + sqlx::query_as!( + InternalVocoderModelListItemRaw, + r#" +SELECT + vocoder.token as vocoder_token, + vocoder.vocoder_type as `vocoder_type: reusable_types::vocoder_type::VocoderType`, + vocoder.creator_user_token, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + vocoder.title, + vocoder.is_staff_recommended, + vocoder.is_mod_disabled_from_public_use, + vocoder.is_mod_disabled_from_author_use, + vocoder.is_mod_author_editing_locked, + vocoder.created_at, + vocoder.updated_at, + vocoder.user_deleted_at, + vocoder.mod_deleted_at +FROM vocoder_models as vocoder +JOIN users + ON users.token = vocoder.creator_user_token +WHERE + users.username = ? + "#, + scope_creator_username) + .fetch_all(mysql_pool) + .await? + }; + + Ok(maybe_models) +} + +struct InternalVocoderModelListItemRaw { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub is_staff_recommended: i8, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub created_at: DateTime, + pub updated_at: DateTime, + + // Moderator fields + pub is_mod_disabled_from_public_use: i8, + pub is_mod_disabled_from_author_use: i8, + pub is_mod_author_editing_locked: i8, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} diff --git a/crates/lib/database_queries/src/queries/vocoder/mod.rs b/crates/lib/database_queries/src/queries/vocoder/mod.rs new file mode 100644 index 0000000000..cb86bd42d7 --- /dev/null +++ b/crates/lib/database_queries/src/queries/vocoder/mod.rs @@ -0,0 +1,3 @@ +pub mod get_vocoder_model; +pub mod insert_vocoder_model; +pub mod list_vocoder_models; \ No newline at end of file diff --git a/crates/lib/database_queries/src/tokens.rs b/crates/lib/database_queries/src/tokens.rs index 44df715317..2b8ed58e04 100644 --- a/crates/lib/database_queries/src/tokens.rs +++ b/crates/lib/database_queries/src/tokens.rs @@ -13,6 +13,10 @@ impl Tokens { random_prefix_crockford_token("TM:", 15, false) } + pub fn new_vocoder_model() -> AnyhowResult { + random_prefix_crockford_token("VM:", 15, false) + } + pub fn new_tts_result() -> AnyhowResult { random_prefix_crockford_token("TR:", 32, false) } @@ -55,6 +59,10 @@ impl Tokens { random_prefix_crockford_token("JWINF:", 32, false) } + pub fn new_generic_upload_job() -> AnyhowResult { + random_prefix_crockford_token("JGUP:", 32, false) + } + pub fn new_firehose_event() -> AnyhowResult { random_prefix_crockford_token("EV:", 32, false) } diff --git a/crates/lib/jobs_common/src/redis_job_status_logger.rs b/crates/lib/jobs_common/src/redis_job_status_logger.rs index c61de40b06..89b9508824 100644 --- a/crates/lib/jobs_common/src/redis_job_status_logger.rs +++ b/crates/lib/jobs_common/src/redis_job_status_logger.rs @@ -14,8 +14,8 @@ pub struct RedisJobStatusLogger <'a> { impl <'a> RedisJobStatusLogger <'a> { pub const STATUS_KEY_TTL_SECONDS : usize = 60 * 60; - pub fn new_tts_download(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self { - let status_key = RedisKeys::tts_download_extra_status_info(w2l_job_token); + pub fn new_tts_download(redis: &'a mut RedisPool, tts_job_token: &str) -> Self { + let status_key = RedisKeys::tts_download_extra_status_info(tts_job_token); Self { redis, status_key } } @@ -24,6 +24,11 @@ impl <'a> RedisJobStatusLogger <'a> { Self { redis, status_key } } + pub fn new_generic_download(redis: &'a mut RedisPool, job_token: &str) -> Self { + let status_key = RedisKeys::generic_download_extra_status_info(job_token); + Self { redis, status_key } + } + pub fn new_w2l_download(redis: &'a mut RedisPool, w2l_job_token: &str) -> Self { let status_key = RedisKeys::w2l_download_extra_status_info(w2l_job_token); Self { redis, status_key } diff --git a/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs b/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs index d810e4a107..382857857e 100644 --- a/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs +++ b/crates/lib/jobs_common/src/semi_persistent_cache_dir.rs @@ -8,7 +8,8 @@ pub struct SemiPersistentCacheDir { cache_root: PathBuf, tts_synthesizer_model_root: PathBuf, - tts_pretrained_vocoder_model_root: PathBuf, + tts_pretrained_vocoder_model_root: PathBuf, // Pretrained, non-user uploaded vocoders + vocoder_model_root: PathBuf, // User-uploaded vocoders w2l_model_root: PathBuf, w2l_end_bump_root: PathBuf, @@ -33,6 +34,7 @@ impl SemiPersistentCacheDir { tts_synthesizer_model_root: cache_root.join("tts/synthesizer_models/"), tts_pretrained_vocoder_model_root: cache_root.join("tts/vocoder_models_pretrained/"), + vocoder_model_root: cache_root.join("tts/user_uploaded_vocoder_models/"), w2l_model_root: cache_root.join("w2l/models/"), w2l_end_bump_root: cache_root.join("w2l/end_bumps/"), @@ -75,6 +77,21 @@ impl SemiPersistentCacheDir { Ok(()) } + // ==================== USER UPLOADED VOCODER MODELS ==================== + + pub fn custom_vocoder_model_path(&self, model_filename: &str) -> PathBuf { + self.vocoder_model_root.join(model_filename) + } + + pub fn custom_vocoder_model_directory(&self) -> &Path { + &self.vocoder_model_root + } + + pub fn create_custom_vocoder_model_path(&self) -> AnyhowResult<()> { + let _ = fs::create_dir_all(self.custom_vocoder_model_directory())?; + Ok(()) + } + // ==================== W2L MODELS (there are only two of them) ==================== /// There are only two pretrained W2L models, so we won't run out of space. @@ -164,24 +181,4 @@ impl SemiPersistentCacheDir { // temp_dir.path().join(&filename) //} -} - -//#[cfg(test)] -//mod tests { -// use crate::util::semi_persistent_cache_dir::SemiPersistentCacheDir; -// use tempdir::TempDir; -// -// #[test] -// fn test_w2l_output_results_path() { -// let dirs = SemiPersistentCacheDir::default_paths(); -// let temp_dir = TempDir::new("test").unwrap(); -// -// let filename = dirs.w2l_output_results_path(&temp_dir, "TYPE:TOKEN") -// .to_str() -// .map(|s| s.to_string()) -// .unwrap(); -// -// assert!(&filename.ends_with("TYPETOKEN_result.mp4")); -// } -//} -// \ No newline at end of file +} \ No newline at end of file diff --git a/crates/lib/redis_common/src/redis_keys.rs b/crates/lib/redis_common/src/redis_keys.rs index 4970bf9ad7..c151d73c91 100644 --- a/crates/lib/redis_common/src/redis_keys.rs +++ b/crates/lib/redis_common/src/redis_keys.rs @@ -25,6 +25,12 @@ impl RedisKeys { format!("ttsDownloadExtraStatus:{}", inference_job_token) } + /// We write extra in-progress status information to keys. + /// These keys should have a TTL. + pub fn generic_download_extra_status_info(inference_job_token: &str) -> String { + format!("genericDownloadExtraStatus:{}", inference_job_token) + } + /// We write extra in-progress status information to keys. /// These keys should have a TTL. pub fn tts_inference_extra_status_info(inference_job_token: &str) -> String { diff --git a/crates/lib/reusable_types/Cargo.toml b/crates/lib/reusable_types/Cargo.toml new file mode 100644 index 0000000000..4dd4df5e04 --- /dev/null +++ b/crates/lib/reusable_types/Cargo.toml @@ -0,0 +1,31 @@ +[package] +name = "reusable_types" +publish = false +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] + +[lib] +name = "reusable_types" +path = "src/lib.rs" + +[dependencies] + +# Internal +# (none) + +# External +anyhow = "1.0.40" +chrono = { version = "0.4.19", features = ["serde"] } + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" +toml = "0.5.8" + +# MySql Database +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } diff --git a/crates/lib/reusable_types/src/entity_visibility.rs b/crates/lib/reusable_types/src/entity_visibility.rs new file mode 100644 index 0000000000..55fea828f7 --- /dev/null +++ b/crates/lib/reusable_types/src/entity_visibility.rs @@ -0,0 +1,66 @@ +/// Used in various database tables (as enums! careful!) and the HTTP API to convey +/// how the associated entity should be made visible to the public. +/// +/// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. +#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] +#[serde(rename_all = "lowercase")] +#[sqlx(rename_all = "lowercase")] +pub enum EntityVisibility { + /// Public entities are able to be listed in public lists. + /// It does not mean that they necessarily will be (eg. they could be "mod unapproved" or deleted). + Public, + /// Hidden entities are not shown in public lists, but the URL to them may be given out freely. + /// They are available to non-logged-in users as long as they have the URL. + Hidden, + /// Private entities should only be available to the creator, a list of approved users, and + /// website moderation staff. + Private, +} + +/// NB: Legacy API for older code. +impl EntityVisibility { + pub fn to_str(&self) -> &'static str { + match self { + Self::Public => "public", + Self::Hidden => "hidden", + Self::Private => "private", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "public" => Ok(Self::Public), + "hidden" => Ok(Self::Hidden), + "private" => Ok(Self::Private), + _ => Err(format!("invalid value: {:?}", value)), + } + } +} + +#[cfg(test)] +mod tests { + use crate::entity_visibility::EntityVisibility; + use crate::test_helpers::assert_serialization; + + #[test] + fn test_serialization() { + assert_serialization(EntityVisibility::Public, "public"); + assert_serialization(EntityVisibility::Hidden, "hidden"); + assert_serialization(EntityVisibility::Private, "private"); + } + + #[test] + fn test_to_str() { + assert_eq!(EntityVisibility::Public.to_str(), "public"); + assert_eq!(EntityVisibility::Hidden.to_str(), "hidden"); + assert_eq!(EntityVisibility::Private.to_str(), "private"); + } + + #[test] + fn test_from_str() { + assert_eq!(EntityVisibility::from_str("public").unwrap(), EntityVisibility::Public); + assert_eq!(EntityVisibility::from_str("hidden").unwrap(), EntityVisibility::Hidden); + assert_eq!(EntityVisibility::from_str("private").unwrap(), EntityVisibility::Private); + assert!(EntityVisibility::from_str("foo").is_err()); + } +} diff --git a/crates/lib/reusable_types/src/generic_download_type.rs b/crates/lib/reusable_types/src/generic_download_type.rs new file mode 100644 index 0000000000..d5dfaf7928 --- /dev/null +++ b/crates/lib/reusable_types/src/generic_download_type.rs @@ -0,0 +1,38 @@ +/// Our "generic downloads" pipeline supports a wide variety of ML models and other media. +/// Each type of download is identified by the following enum variants. +/// These types are present in the HTTP API and database columns as serialized here. +/// +/// DO NOT CHANGE VALUES WITHOUT A MIGRATION STRATEGY. +#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] +pub enum GenericDownloadType { + #[serde(rename = "hifigan")] + #[sqlx(rename = "hifigan")] + HifiGan, +} + +/// NB: Legacy API for older code. +impl GenericDownloadType { + pub fn to_str(&self) -> &'static str { + match self { + Self::HifiGan=> "hifigan", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "hifigan" => Ok(Self::HifiGan), + _ => Err(format!("invalid value: {:?}", value)), + } + } +} + +#[cfg(test)] +mod tests { + use crate::generic_download_type::GenericDownloadType; + use crate::test_helpers::assert_serialization; + + #[test] + fn test_serialization() { + assert_serialization(GenericDownloadType::HifiGan, "hifigan"); + } +} diff --git a/crates/lib/reusable_types/src/lib.rs b/crates/lib/reusable_types/src/lib.rs new file mode 100644 index 0000000000..cf64868e4e --- /dev/null +++ b/crates/lib/reusable_types/src/lib.rs @@ -0,0 +1,20 @@ +// Never allow these +#![forbid(private_in_public)] +#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things + +// Okay to toggle +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +#[macro_use] extern crate serde_derive; + +pub mod entity_visibility; +pub mod generic_download_type; +pub mod vocoder_type; + +#[cfg(test)] pub mod test_helpers; \ No newline at end of file diff --git a/crates/lib/reusable_types/src/test_helpers.rs b/crates/lib/reusable_types/src/test_helpers.rs new file mode 100644 index 0000000000..6d16bfddd5 --- /dev/null +++ b/crates/lib/reusable_types/src/test_helpers.rs @@ -0,0 +1,16 @@ +//! Only imported in tests + +use serde::Serialize; + +pub fn to_json(t: T) -> String { + toml::to_string(&t) + .expect("serialization error") + .replace("\"", "") // JSON values are quoted, so we remove quotes +} + +/// Assert that the Serialize is represented by the expected string value. +/// This is useful for ensuring stability of serialization. +pub fn assert_serialization(t: T, expected: &str) { + // TODO: See if there's a way to check the sqlx serialization + assert_eq!(&to_json(t), expected); +} diff --git a/crates/lib/reusable_types/src/vocoder_type.rs b/crates/lib/reusable_types/src/vocoder_type.rs new file mode 100644 index 0000000000..1581c99903 --- /dev/null +++ b/crates/lib/reusable_types/src/vocoder_type.rs @@ -0,0 +1,33 @@ +#[derive(Clone, Copy, Eq, PartialEq, Debug, Deserialize, Serialize, sqlx::Type)] +pub enum VocoderType { + #[serde(rename = "hifigan")] + #[sqlx(rename = "hifigan")] + HifiGan, +} + +/// NB: Legacy API for older code. +impl VocoderType { + pub fn to_str(&self) -> &'static str { + match self { + Self::HifiGan=> "hifigan", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "hifigan" => Ok(Self::HifiGan), + _ => Err(format!("invalid value: {:?}", value)), + } + } +} + +#[cfg(test)] +mod tests { + use crate::test_helpers::assert_serialization; + use crate::vocoder_type::VocoderType; + + #[test] + fn test_serialization() { + assert_serialization(VocoderType::HifiGan, "hifigan"); + } +} diff --git a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs index 4a99eff362..a9611564d6 100644 --- a/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs +++ b/crates/lib/storage_buckets_common/src/bucket_path_unifier.rs @@ -6,9 +6,10 @@ use std::path::PathBuf; /// different types of objects are stored. pub struct BucketPathUnifier { // TTS - pub tts_pretrained_vocoder_model_root: PathBuf, pub tts_synthesizer_model_root: PathBuf, pub tts_inference_output_root: PathBuf, + pub tts_pretrained_vocoder_model_root: PathBuf, // Pretrained, non-user uploaded vocoders + pub vocoder_model_root: PathBuf, // User-uploaded vocoders // W2L pub user_uploaded_w2l_templates_root: PathBuf, pub user_uploaded_audio_for_w2l_root: PathBuf, @@ -19,20 +20,13 @@ pub struct BucketPathUnifier { impl BucketPathUnifier { - // TODO - //pub fn from_env_vars() -> AnyhowResult { - // Ok(Self { - // user_uploaded_audio_for_w2l_root: , - // user_uploaded_w2l_templates_root: , - // }) - //} - pub fn default_paths() -> Self { Self { // TTS - tts_pretrained_vocoder_model_root: PathBuf::from("/tts_pretrained_vocoders"), tts_synthesizer_model_root: PathBuf::from("/user_uploaded_tts_synthesizers"), tts_inference_output_root: PathBuf::from("/tts_inference_output"), + tts_pretrained_vocoder_model_root: PathBuf::from("/tts_pretrained_vocoders"), // Pretrained, non-user uploaded vocoders + vocoder_model_root: PathBuf::from("/user_uploaded_vocoders"), // W2L user_uploaded_audio_for_w2l_root: PathBuf::from("/user_uploaded_w2l_audio"), user_uploaded_w2l_templates_root: PathBuf::from("/user_uploaded_w2l_templates"), @@ -44,12 +38,6 @@ impl BucketPathUnifier { // ==================== TTS MODELS (SYNTHESIZER + VOCODER) ==================== // - // TTS pretrained vocoder models. - // For now this will be limited, but once they're user-uploadable, we'll add more. - pub fn tts_pretrained_vocoders_path(&self, tts_vocoder_model_name: &str) -> PathBuf { - self.tts_pretrained_vocoder_model_root.join(tts_vocoder_model_name) - } - pub fn tts_synthesizer_path(&self, tts_synthesizer_file_hash: &str) -> PathBuf { let hashed_path = Self::hashed_directory_path(tts_synthesizer_file_hash); let model_filename = format!("{}.pt", &tts_synthesizer_file_hash); @@ -68,6 +56,25 @@ impl BucketPathUnifier { .join(model_filename) } + // ==================== VOCODER MODELS ==================== // + + /// TTS pretrained vocoder models. + /// For now this will be limited, but once they're user-uploadable, we'll add more. + /// NB: These are vocoders that we uploaded without the FakeYou upload system!! + pub fn tts_pretrained_vocoders_path(&self, tts_vocoder_model_name: &str) -> PathBuf { + self.tts_pretrained_vocoder_model_root.join(tts_vocoder_model_name) + } + + /// User-uploaded vocoders. + pub fn vocoder_path(&self, vocoder_file_hash: &str) -> PathBuf { + let hashed_path = Self::hashed_directory_path(vocoder_file_hash); + let model_filename = format!("{}.pt", &vocoder_file_hash); + + self.vocoder_model_root + .join(hashed_path) + .join(model_filename) + } + // ==================== TTS INFERENCE OUTPUT ==================== // /// This should include the string "vocodes" for downloaders. @@ -179,9 +186,10 @@ mod tests { fn get_instance() -> BucketPathUnifier { BucketPathUnifier { // TTS - tts_pretrained_vocoder_model_root: PathBuf::from("/test_path_tts_vocoders"), tts_synthesizer_model_root: PathBuf::from("/test_path_synthesizers"), tts_inference_output_root: PathBuf::from("/test_path_tts_output"), + tts_pretrained_vocoder_model_root: PathBuf::from("/test_path_tts_vocoders"), // Pretrained, non-user uploaded vocoders + vocoder_model_root: PathBuf::from("/test_path_user_vocoders"), // User-uploaded vocoders // W2L user_uploaded_w2l_templates_root: PathBuf::from("/test_path_w2l_templates"), user_uploaded_audio_for_w2l_root: PathBuf::from("/test_path_w2l_audio"), @@ -191,6 +199,7 @@ mod tests { } } + // NB: Pretrained, non-user uploaded vocoders #[test] fn test_tts_pretrained_vocoders_path() { let paths = get_instance(); @@ -198,6 +207,14 @@ mod tests { "/test_path_tts_vocoders/melgan.pth"); } + // NB: User-uploaded vocoders + #[test] + fn test_vocoders_path() { + let paths = get_instance(); + assert_eq!(paths.vocoder_path("foobar").to_str().unwrap(), + "/test_path_user_vocoders/f/o/o/foobar.pt"); + } + #[test] fn test_tts_synthesizer_path() { let paths = get_instance(); diff --git a/crates/service/generic_download_job/Cargo.toml b/crates/service/generic_download_job/Cargo.toml new file mode 100644 index 0000000000..83833637b5 --- /dev/null +++ b/crates/service/generic_download_job/Cargo.toml @@ -0,0 +1,62 @@ +[package] +name = "generic-download-job" +publish = false +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] + +[[bin]] +name = "generic-download-job" +path = "src/main.rs" + +[dependencies] + +# Internal +config = { path = "../../lib/config" } +container_common = { path = "../../lib/container_common" } +database_queries = { path = "../../lib/database_queries" } +google_drive_common = { path = "../../lib/google_drive_common" } +jobs_common = { path = "../../lib/jobs_common" } +reusable_types = { path = "../../lib/reusable_types" } +storage_buckets_common = { path = "../../lib/storage_buckets_common" } +tts_common = { path = "../../lib/tts_common" } + +# Vendored code +newrelic-telemetry = { path = "../../vendor/newrelic-telemetry-sdk-rust" } + +# External +anyhow = "1.0.40" +chrono = { version = "0.4.19", features = ["serde"] } +clap = "2.34.0" # TODO: This is really out of date, but has conflicts when updated +dotenv = "0.15.0" +easyenv = "0.0.6" +hostname = "0.3.1" +lfu = "0.2.5" +log = "0.4.14" +subprocess = "0.2.6" +tempdir = "0.3.7" + +# MySql Database +# "offline" supports typesafe builds without a live database +# "rustls" doesn't appear to work with IP addresses, per: +# https://github.com/launchbadge/sqlx/issues/1095#issuecomment-798050828 +sqlx = { version = "0.5", features = [ "mysql", "runtime-actix-native-tls", "offline", "chrono" ] } + +# Redis +r2d2_redis = "0.14.0" + +# Runtime +# NB: We switched to Tokio from `async-std` for broader ecosystem support. +tokio = { version = "1.6.0", features = ["macros"] } + +# NB: Hyper is just for HTTP client use in jobs. I'd use Actix client, but I need tokio runtime +# NB: Furthermore, the Tokio runtime is old, so I have to use an older Hyper. +hyper = { version = "0.14.14" } + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" +serde_json = "1.0.64" diff --git a/crates/service/generic_download_job/src/job_state.rs b/crates/service/generic_download_job/src/job_state.rs new file mode 100644 index 0000000000..c92c1817fd --- /dev/null +++ b/crates/service/generic_download_job/src/job_state.rs @@ -0,0 +1,41 @@ +use crate::HifiGanModelCheckCommand; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use sqlx::MySqlPool; +use std::path::PathBuf; +use storage_buckets_common::bucket_client::BucketClient; +use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; + +pub struct JobState { + pub download_temp_directory: PathBuf, + pub mysql_pool: MySqlPool, + + pub redis_pool: r2d2::Pool, + + pub bucket_client: BucketClient, + pub firehose_publisher: FirehosePublisher, + pub badge_granter: BadgeGranter, + pub google_drive_downloader: GoogleDriveDownloadCommand, + + pub bucket_path_unifier: BucketPathUnifier, + + pub hifigan_model_check_command: HifiGanModelCheckCommand, + + // Command to run + pub download_script: String, + // Root to store TTS results + pub bucket_root_tts_model_uploads: String, + + // Sleep between batches + pub job_batch_wait_millis: u64, + + // How long to wait between log lines + pub no_op_logger_millis: u64, + + // Max job attempts before failure. + // NB: This is an i32 so we don't need to convert to db column type. + pub job_max_attempts: i32, +} diff --git a/crates/service/generic_download_job/src/job_steps/job_results.rs b/crates/service/generic_download_job/src/job_steps/job_results.rs new file mode 100644 index 0000000000..ae282b7134 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/job_results.rs @@ -0,0 +1,9 @@ + +pub struct JobResults { + /// The token of the thing that got downloaded and saved. + /// Can be a token corresponding to any of multiple different tables (vocoder_models, tts_models, etc.) + pub entity_token: Option, + + /// This may differ from `GenericDownloadType` if the download is polymorphic. + pub entity_type: Option, +} \ No newline at end of file diff --git a/crates/service/generic_download_job/src/job_steps/mod.rs b/crates/service/generic_download_job/src/job_steps/mod.rs new file mode 100644 index 0000000000..9fffcb7af3 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/mod.rs @@ -0,0 +1,3 @@ +pub mod process_single_job; +pub mod process_hifigan_vocoder; +pub mod job_results; \ No newline at end of file diff --git a/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs b/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs new file mode 100644 index 0000000000..baef4fd4d7 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/process_hifigan_vocoder.rs @@ -0,0 +1,135 @@ +use anyhow::anyhow; +use container_common::anyhow_result::AnyhowResult; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_file_sha2::hash_file_sha2; +use crate::JobState; +use crate::job_steps::job_results::JobResults; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use database_queries::queries::tts::tts_download_jobs::tts_download_job_queries::insert_tts_model; +use database_queries::queries::vocoder::insert_vocoder_model::{Args, insert_vocoder_model}; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; +use reusable_types::vocoder_type::VocoderType; +use std::fs::File; +use std::io::Read; +use std::path::PathBuf; +use tempdir::TempDir; + +/// Returns the token of the entity. +pub async fn process_hifigan_vocoder<'a, 'b>( + job_state: &JobState, + job: &AvailableDownloadJob, + temp_dir: &TempDir, + download_filename: &str, + redis_logger: &'a mut RedisJobStatusLogger<'b>, +) -> AnyhowResult { + + // ==================== RUN MODEL CHECK ==================== // + + info!("Checking that model is valid..."); + + redis_logger.log_status("checking hifigan model")?; + + let file_path = PathBuf::from(download_filename.clone()); + + let output_metadata_fs_path = temp_dir.path().join("metadata.json"); + + let model_check_result = job_state.hifigan_model_check_command.execute( + &file_path, + &output_metadata_fs_path + ); + + if let Err(e) = model_check_result { + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(anyhow!("model check error: {:?}", e)); + } + + // ==================== CHECK ALL FILES EXIST AND GET METADATA ==================== // + + info!("Checking that metadata output file exists..."); + + check_file_exists(&output_metadata_fs_path)?; + + let file_metadata = match read_metadata_file(&output_metadata_fs_path) { + Ok(metadata) => metadata, + Err(e) => { + safe_delete_temp_file(&file_path); + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== UPLOAD MODEL FILE ==================== // + + info!("Uploading HifiGan vocoder to GCS..."); + + let private_bucket_hash = hash_file_sha2(&download_filename)?; + + info!("File hash: {}", private_bucket_hash); + + let model_bucket_path = job_state.bucket_path_unifier.vocoder_path(&private_bucket_hash); + + info!("Destination bucket path: {:?}", &model_bucket_path); + + redis_logger.log_status("uploading hifigan model")?; + + if let Err(e) = job_state.bucket_client.upload_filename(&model_bucket_path, &file_path).await { + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + + // ==================== DELETE DOWNLOADED FILE ==================== // + + // NB: We should be using a tempdir, but to make absolutely certain we don't overflow the disk... + safe_delete_temp_file(&output_metadata_fs_path); + safe_delete_temp_file(&file_path); + safe_delete_temp_directory(&temp_dir); + + // ==================== SAVE RECORDS ==================== // + + info!("Saving model record..."); + let (id, model_token) = insert_vocoder_model(Args { + vocoder_type: VocoderType::HifiGan, + title: &job.title, + original_download_url: &job.download_url, + original_filename: &download_filename, + file_size_bytes: file_metadata.file_size_bytes, + creator_user_token: &job.creator_user_token, + creator_ip_address: &job.creator_ip_address, + creator_set_visibility: job.creator_set_visibility, + private_bucket_hash: &private_bucket_hash, + private_bucket_object_name: &model_bucket_path, + mysql_pool: &job_state.mysql_pool + }).await?; + + job_state.badge_granter.maybe_grant_vocoder_model_uploads_badge(&job.creator_user_token) + .await + .map_err(|e| { + warn!("error maybe awarding badge: {:?}", e); + anyhow!("error maybe awarding badge") + })?; + + Ok(JobResults { + entity_token: Some(model_token), + entity_type: Some("hifigan".to_string()), // NB: This may be different from `GenericDownloadType` in the future! + }) +} + +#[derive(Deserialize)] +struct FileMetadata { + pub file_size_bytes: u64, +} + +fn read_metadata_file(filename: &PathBuf) -> AnyhowResult { + let mut file = File::open(filename)?; + let mut buffer = String::new(); + file.read_to_string(&mut buffer)?; + Ok(serde_json::from_str(&buffer)?) +} + diff --git a/crates/service/generic_download_job/src/job_steps/process_single_job.rs b/crates/service/generic_download_job/src/job_steps/process_single_job.rs new file mode 100644 index 0000000000..c8160957d0 --- /dev/null +++ b/crates/service/generic_download_job/src/job_steps/process_single_job.rs @@ -0,0 +1,103 @@ +use anyhow::anyhow; +use config::bad_urls::is_bad_tts_model_download_url; +use config::is_bad_download_url::is_bad_download_url; +use container_common::anyhow_result::AnyhowResult; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_file_sha2::hash_file_sha2; +use crate::JobState; +use crate::job_steps::process_hifigan_vocoder::process_hifigan_vocoder; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::AvailableDownloadJob; +use database_queries::queries::generic_download::job::mark_generic_download_job_done::mark_generic_download_job_done; +use database_queries::queries::generic_download::job::mark_generic_download_job_pending_and_grab_lock::mark_generic_download_job_pending_and_grab_lock; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{info, warn}; +use reusable_types::generic_download_type::GenericDownloadType; +use std::path::PathBuf; +use tempdir::TempDir; + +pub async fn process_single_job(job_state: &JobState, job: &AvailableDownloadJob) -> AnyhowResult<()> { + let mut redis = job_state.redis_pool.get()?; + let mut redis_logger = RedisJobStatusLogger::new_generic_download(&mut redis, &job.download_job_token); + + // ==================== ATTEMPT TO GRAB JOB LOCK ==================== // + + let lock_acquired = mark_generic_download_job_pending_and_grab_lock(&job_state.mysql_pool, job.id).await?; + + if !lock_acquired { + warn!("Could not acquire job lock for: {}", &job.id.0); + return Ok(()) + } + + // ==================== SETUP TEMP DIRS ==================== // + + let temp_dir = format!("temp_{}", job.id.0); + let temp_dir = TempDir::new(&temp_dir)?; + + // ==================== DOWNLOAD MODEL FILE ==================== // + + info!("Calling downloader..."); + + redis_logger.log_status("downloading URL")?; + + if is_bad_download_url(&job.download_url)? { + warn!("Bad download URL: `{}`", &job.download_url); + return Err(anyhow!("Bad download URL: `{}`", &job.download_url)); + } + + let download_filename = match job_state.google_drive_downloader.download_file(&job.download_url, &temp_dir).await { + Ok(filename) => filename, + Err(e) => { + safe_delete_temp_directory(&temp_dir); + return Err(e); + } + }; + + // ==================== HANDLE DIFFERENT DOWNLOAD TYPES ==================== // + + let mut entity_token : Option = None; + let mut entity_type : Option = None; + + match job.download_type { + GenericDownloadType::HifiGan => { + let results = process_hifigan_vocoder( + job_state, + job, + &temp_dir, + &download_filename, + &mut redis_logger, + ).await?; + entity_token = results.entity_token.clone(); + entity_type = results.entity_type.clone(); + } + } + + // ===================================================== + + info!("Marking job complete..."); + mark_generic_download_job_done( + &job_state.mysql_pool, + job, + true, + entity_token.as_deref(), + entity_type.as_deref(), + ).await?; + + info!("Saved model record: {} - {}", job.id.0, &job.download_job_token); + + job_state.firehose_publisher.publish_generic_download_finished( + &job.creator_user_token, + entity_token.as_deref()) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + anyhow!("error publishing event") + })?; + + redis_logger.log_status("done")?; + + info!("Job done: {}", job.id.0); + + Ok(()) +} diff --git a/crates/service/generic_download_job/src/main.rs b/crates/service/generic_download_job/src/main.rs new file mode 100644 index 0000000000..562717990c --- /dev/null +++ b/crates/service/generic_download_job/src/main.rs @@ -0,0 +1,238 @@ +#![forbid(private_in_public)] +#![forbid(unused_must_use)] +//#![forbid(warnings)] + +#![allow(dead_code)] +#![allow(non_snake_case)] +#![allow(unused_imports)] +#![allow(unused_mut)] +#![allow(unused_variables)] + +#[macro_use] extern crate serde_derive; + +pub mod job_state; +pub mod job_steps; +pub mod script_execution; + +use anyhow::anyhow; +use chrono::Utc; +use config::bad_urls::is_bad_tts_model_download_url; +use config::common_env::CommonEnv; +use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; +use config::shared_constants::DEFAULT_RUST_LOG; +use container_common::anyhow_result::AnyhowResult; +use container_common::filesystem::check_directory_exists::check_directory_exists; +use container_common::filesystem::check_file_exists::check_file_exists; +use container_common::filesystem::safe_delete_temp_directory::safe_delete_temp_directory; +use container_common::filesystem::safe_delete_temp_file::safe_delete_temp_file; +use container_common::hashing::hash_file_sha2::hash_file_sha2; +use crate::job_state::JobState; +use crate::job_steps::process_single_job::process_single_job; +use crate::script_execution::hifigan_model_check_command::HifiGanModelCheckCommand; +use database_queries::column_types::tts_model_type::TtsModelType; +use database_queries::mediators::badge_granter::BadgeGranter; +use database_queries::mediators::firehose_publisher::FirehosePublisher; +use database_queries::queries::generic_download::job::list_available_generic_download_jobs::{AvailableDownloadJob, list_available_generic_download_jobs}; +use database_queries::queries::generic_download::job::mark_generic_download_job_failure::mark_generic_download_job_failure; +use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use jobs_common::noop_logger::NoOpLogger; +use jobs_common::redis_job_status_logger::RedisJobStatusLogger; +use log::{warn, info}; +use r2d2_redis::RedisConnectionManager; +use r2d2_redis::r2d2; +use sqlx::MySqlPool; +use sqlx::mysql::MySqlPoolOptions; +use std::fs::File; +use std::fs; +use std::io::{BufReader, Read, Error}; +use std::ops::Deref; +use std::path::{PathBuf, Path}; +use std::process::Command; +use std::time::Duration; +use storage_buckets_common::bucket_client::BucketClient; +use storage_buckets_common::bucket_path_unifier::BucketPathUnifier; +use tempdir::TempDir; + +// Buckets +const ENV_ACCESS_KEY : &'static str = "ACCESS_KEY"; +const ENV_SECRET_KEY : &'static str = "SECRET_KEY"; +const ENV_REGION_NAME : &'static str = "REGION_NAME"; +const ENV_BUCKET_NAME : &'static str = "TTS_DOWNLOAD_BUCKET_NAME"; +const ENV_BUCKET_ROOT : &'static str = "TTS_DOWNLOAD_BUCKET_ROOT"; + +const DEFAULT_TEMP_DIR: &'static str = "/tmp"; + +#[tokio::main] +async fn main() -> AnyhowResult<()> { + easyenv::init_all_with_default_logging(Some(DEFAULT_RUST_LOG)); + + let _ = dotenv::from_filename(".env-secrets").ok(); + + info!("Obtaining hostname..."); + + let server_hostname = hostname::get() + .ok() + .and_then(|h| h.into_string().ok()) + .unwrap_or("generic-download-job".to_string()); + + info!("Hostname: {}", &server_hostname); + + // Bucket stuff + let access_key = easyenv::get_env_string_required(ENV_ACCESS_KEY)?; + let secret_key = easyenv::get_env_string_required(ENV_SECRET_KEY)?; + let region_name = easyenv::get_env_string_required(ENV_REGION_NAME)?; + let bucket_name = easyenv::get_env_string_required(ENV_BUCKET_NAME)?; + let bucket_root = easyenv::get_env_string_required(ENV_BUCKET_ROOT)?; + + let bucket_client = BucketClient::create( + &access_key, + &secret_key, + ®ion_name, + &bucket_name, + None, + )?; + + let temp_directory = easyenv::get_env_string_or_default( + "DOWNLOAD_TEMP_DIR", + DEFAULT_TEMP_DIR); + + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "./scripts/download_internet_file.py"); + + let google_drive_downloader = GoogleDriveDownloadCommand::new(&download_script); + + let temp_directory = PathBuf::from(temp_directory); + + check_directory_exists(&temp_directory)?; + + let db_connection_string = + easyenv::get_env_string_or_default( + "MYSQL_URL", + DEFAULT_MYSQL_CONNECTION_STRING); + + info!("Connecting to database..."); + + let mysql_pool = MySqlPoolOptions::new() + .max_connections(5) + .connect(&db_connection_string) + .await?; + + let common_env = CommonEnv::read_from_env()?; + + info!("Connecting to redis..."); + + let redis_manager = + RedisConnectionManager::new(common_env.redis_0_connection_string.deref())?; + + let redis_pool = r2d2::Pool::builder() + .build(redis_manager)?; + + let firehose_publisher = FirehosePublisher { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + }; + + let badge_granter = BadgeGranter { + mysql_pool: mysql_pool.clone(), // NB: Pool is sync/send/clone-safe + firehose_publisher: firehose_publisher.clone(), // NB: Also safe + }; + + let hifigan_model_check_command= HifiGanModelCheckCommand::new( + &easyenv::get_env_string_required("HIFIGAN_ROOT_CODE_DIRECTORY")?, + &easyenv::get_env_string_or_default( + "HIFIGAN_VIRTUAL_ENV_ACTIVATION_COMMAND", + "source python-tacotron/bin/activate"), + &easyenv::get_env_string_or_default( + "HIFIGAN_MODEL_CHECK_SCRIPT_NAME", + "vocodes_model_check_hifigan.py"), + ); + + let job_state = JobState { + download_temp_directory: temp_directory, + mysql_pool, + redis_pool, + bucket_client, + download_script, + google_drive_downloader, + bucket_path_unifier: BucketPathUnifier::default_paths(), + bucket_root_tts_model_uploads: bucket_root.to_string(), + firehose_publisher, + badge_granter, + hifigan_model_check_command, + job_batch_wait_millis: common_env.job_batch_wait_millis, + job_max_attempts: common_env.job_max_attempts as i32, + no_op_logger_millis: common_env.no_op_logger_millis, + }; + + main_loop(job_state).await; + + Ok(()) +} + +const START_TIMEOUT_MILLIS : u64 = 500; +const INCREASE_TIMEOUT_MILLIS : u64 = 1000; + +async fn main_loop(job_state: JobState) { + let mut error_timeout_millis = START_TIMEOUT_MILLIS; + + let mut noop_logger = NoOpLogger::new(job_state.no_op_logger_millis as i64); + + loop { + let num_records = 1; + let maybe_available_jobs = list_available_generic_download_jobs(&job_state.mysql_pool, num_records).await; + + let jobs = match maybe_available_jobs { + Ok(jobs) => jobs, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + }; + + if jobs.is_empty() { + noop_logger.log_after_awhile(); + + std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis)); + continue; + } + + let result = process_jobs(&job_state, jobs).await; + + match result { + Ok(_) => {}, + Err(e) => { + warn!("Error querying jobs: {:?}", e); + std::thread::sleep(Duration::from_millis(error_timeout_millis)); + error_timeout_millis += INCREASE_TIMEOUT_MILLIS; + continue; + } + } + + error_timeout_millis = START_TIMEOUT_MILLIS; // reset + + std::thread::sleep(Duration::from_millis(job_state.job_batch_wait_millis)); + } +} + +async fn process_jobs(job_state: &JobState, jobs: Vec) -> AnyhowResult<()> { + for job in jobs.into_iter() { + let result = process_single_job(job_state, &job).await; + match result { + Ok(_) => {}, + Err(e) => { + warn!("Failure to process job: {:?}", e); + let failure_reason = ""; + let _r = mark_generic_download_job_failure( + &job_state.mysql_pool, + &job, + failure_reason, + job_state.job_max_attempts + ).await; + } + } + } + + Ok(()) +} diff --git a/crates/service/generic_download_job/src/script_execution/hifigan_model_check_command.rs b/crates/service/generic_download_job/src/script_execution/hifigan_model_check_command.rs new file mode 100644 index 0000000000..1e48e9498e --- /dev/null +++ b/crates/service/generic_download_job/src/script_execution/hifigan_model_check_command.rs @@ -0,0 +1,71 @@ +use crate::AnyhowResult; +use log::{info,warn}; +use std::process::{Command, Stdio}; +use subprocess::{Popen, PopenConfig, Redirection}; +use std::fs::OpenOptions; +use std::path::Path; + +/// This command is used to check hifigan for being a real model +#[derive(Clone)] +pub struct HifiGanModelCheckCommand { + /// Where the HifiGan code lives + hifigan_root_code_directory: String, + + /// eg. `source python/bin/activate` + virtual_env_activation_command: String, + + hifigan_model_check_script_name: String, +} + +impl HifiGanModelCheckCommand { + pub fn new( + hifigan_root_code_directory: &str, + virtual_env_activation_command: &str, + hifigan_model_check_script_name: &str, + ) -> Self { + Self { + hifigan_root_code_directory: hifigan_root_code_directory.to_string(), + virtual_env_activation_command: virtual_env_activation_command.to_string(), + hifigan_model_check_script_name: hifigan_model_check_script_name.to_string(), + } + } + + pub fn execute>( + &self, + checkpoint_path: P, + output_metadata_filename: P, + ) -> AnyhowResult<()> { + + let mut command = String::new(); + command.push_str(&format!("cd {}", self.hifigan_root_code_directory)); + command.push_str(" && "); + command.push_str(&self.virtual_env_activation_command); + command.push_str(" && "); + command.push_str("python "); + command.push_str(&self.hifigan_model_check_script_name); + command.push_str(" --checkpoint_path "); + command.push_str(&checkpoint_path.as_ref().display().to_string()); + command.push_str(" --output_metadata_filename "); + command.push_str(&output_metadata_filename.as_ref().display().to_string()); + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut p = Popen::create(&command_parts, PopenConfig { + ..Default::default() + })?; + + info!("Pid : {:?}", p.pid()); + + let exit_status = p.wait()?; + + info!("Exit status: {:?}", exit_status); + + Ok(()) + } +} diff --git a/crates/service/generic_download_job/src/script_execution/mod.rs b/crates/service/generic_download_job/src/script_execution/mod.rs new file mode 100644 index 0000000000..a8fd2b8e3e --- /dev/null +++ b/crates/service/generic_download_job/src/script_execution/mod.rs @@ -0,0 +1 @@ +pub mod hifigan_model_check_command; \ No newline at end of file diff --git a/crates/service/storyteller_web/Cargo.toml b/crates/service/storyteller_web/Cargo.toml index 8db78ff6f9..3bb8eb3e6b 100644 --- a/crates/service/storyteller_web/Cargo.toml +++ b/crates/service/storyteller_web/Cargo.toml @@ -24,6 +24,7 @@ http_server_common = { path = "../../lib/http_server_common" } jobs_common = { path = "../../lib/jobs_common" } memory_caching = { path = "../../lib/memory_caching" } redis_common = { path = "../../lib/redis_common" } +reusable_types = { path = "../../lib/reusable_types" } storage_buckets_common = { path = "../../lib/storage_buckets_common" } tts_common = { path = "../../lib/tts_common" } twitch_common = { path = "../../lib/twitch_common" } @@ -45,6 +46,7 @@ lazy_static = "1.4.0" lexical-sort = "0.3.1" lfu = "0.2.5" log = "0.4.14" +once_cell = "1.9.0" rand = "0.8.3" regex = "1.5.4" time = "0.2.26" diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs new file mode 100644 index 0000000000..d825bf5658 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/enqueue_generic_download.rs @@ -0,0 +1,177 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; +use config::bad_urls::is_bad_tts_model_download_url; +use crate::http_server::web_utils::ip_address::get_request_ip; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use crate::validations::model_uploads::validate_model_title; +use crate::validations::passwords::validate_passwords; +use crate::validations::username::validate_username; +use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use database_queries::queries::generic_download::insert_generic_download_job::{Args, insert_generic_download_job}; +use database_queries::tokens::Tokens; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use regex::Regex; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::generic_download_type::GenericDownloadType; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +#[derive(Deserialize)] +pub struct EnqueueGenericDownloadRequest { + idempotency_token: String, + title: String, + download_url: String, + generic_download_type: GenericDownloadType, + + /// Not all upload types will have an associated visibility, so supplying this is optional. + /// If an upload type supports it and it is not supplied, assume a reasonable default or user + /// setting. + creator_set_visibility: Option, +} + +#[derive(Serialize)] +pub struct EnqueueGenericDownloadSuccessResponse { + pub success: bool, + /// This is how frontend clients can request the job execution status. + pub job_token: String, +} + +#[derive(Debug, Serialize)] +pub enum EnqueueGenericDownloadError { + BadInput(String), + MustBeLoggedIn, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueGenericDownloadError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueGenericDownloadError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueGenericDownloadError::MustBeLoggedIn => StatusCode::UNAUTHORIZED, + EnqueueGenericDownloadError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueGenericDownloadError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + serialize_as_json_error(self) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for EnqueueGenericDownloadError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn enqueue_generic_download_handler( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data>) -> Result +{ + if let Err(_err) = server_state.redis_rate_limiters.model_upload.rate_limit_request(&http_request) { + return Err(EnqueueGenericDownloadError::RateLimited); + } + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueGenericDownloadError::ServerError + })?; + + let user_session = match maybe_user_session { + Some(session) => session, + None => { + warn!("not logged in"); + return Err(EnqueueGenericDownloadError::MustBeLoggedIn); + } + }; + + let ip_address = get_request_ip(&http_request); + let uuid = request.idempotency_token.to_string(); + let title = request.title.trim().to_string(); + let download_url = request.download_url.trim().to_string(); + + let creator_set_visibility = request.creator_set_visibility + .unwrap_or(EntityVisibility::Public); + + if let Err(reason) = validate_idempotency_token_format(&uuid) { + return Err(EnqueueGenericDownloadError::BadInput(reason)); + } + + if let Err(reason) = validate_model_title(&title) { + return Err(EnqueueGenericDownloadError::BadInput(reason)); + } + + match is_bad_tts_model_download_url(&download_url) { + Ok(false) => {} // Ok case + Ok(true) => { + return Err(EnqueueGenericDownloadError::BadInput("Bad model download URL".to_string())); + } + Err(err) => { + warn!("Error parsing url: {:?}", err); + return Err(EnqueueGenericDownloadError::BadInput("Bad model download URL".to_string())); + } + } + + // This token is returned to the client. + let job_token = Tokens::new_generic_upload_job() + .map_err(|e| { + warn!("Error creating token"); + EnqueueGenericDownloadError::ServerError + })?; + + + let record_id = insert_generic_download_job(Args { + job_token: &job_token, + uuid_idempotency_token: &uuid, + download_type: request.generic_download_type, + download_url: &download_url, + title: &title, + creator_user_token: &user_session.user_token, + creator_ip_address: &ip_address, + creator_set_visibility, + mysql_pool: &server_state.mysql_pool, + }) + .await + .map_err(|err| { + warn!("New generic download creation DB error: {:?}", err); + EnqueueGenericDownloadError::ServerError + })?; + + info!("new generic download job id: {}", record_id); + + server_state.firehose_publisher.enqueue_generic_download(&user_session.user_token, &job_token) + .await + .map_err(|e| { + warn!("error publishing event: {:?}", e); + EnqueueGenericDownloadError::ServerError + })?; + + let response = EnqueueGenericDownloadSuccessResponse { + success: true, + job_token: job_token.to_string(), + }; + + let body = serde_json::to_string(&response) + .map_err(|e| EnqueueGenericDownloadError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs new file mode 100644 index 0000000000..88bc54400d --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/get_generic_upload_job_status.rs @@ -0,0 +1,159 @@ +use actix_http::Error; +use actix_http::http::header; +use actix_web::HttpResponseBuilder; +use actix_web::cookie::Cookie; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use crate::AnyhowResult; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::generic_download::get_generic_download_job_status::{GenericDownloadJobStatus, get_generic_download_job_status}; +use http_server_common::response::serialize_as_json_error::serialize_as_json_error; +use log::{info, warn, log}; +use r2d2_redis::redis::Commands; +use redis_common::redis_keys::RedisKeys; +use regex::Regex; +use sqlx::MySqlPool; +use sqlx::error::DatabaseError; +use sqlx::error::Error::Database; +use sqlx::mysql::MySqlDatabaseError; +use std::fmt; +use std::sync::Arc; + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetGenericDownloadJobStatusPathInfo { + token: String, +} + +#[derive(Serialize)] +pub struct GetGenericDownloadJobStatusForResponse { + pub job_token: String, + pub status: String, + + /// Extra, temporary status from Redis. + /// This can denote inference progress, and the Python code can write to it. + pub maybe_extra_status_description: Option, + + pub attempt_count: u8, + pub maybe_downloaded_entity_type: Option, + pub maybe_downloaded_entity_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +#[derive(Serialize)] +pub struct GetGenericDownloadJobStatusSuccessResponse { + pub success: bool, + pub state: GetGenericDownloadJobStatusForResponse, +} + +#[derive(Debug, Serialize)] +pub enum GetGenericDownloadJobStatusError { + NotFoundError, + ServerError, +} + +#[derive(Serialize)] +pub struct GetGenericDownloadJobStatusRecord { + pub job_token: String, + + pub status: String, + pub attempt_count: i32, + pub maybe_model_token: Option, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +impl ResponseError for GetGenericDownloadJobStatusError { + fn status_code(&self) -> StatusCode { + match *self { + GetGenericDownloadJobStatusError::NotFoundError => StatusCode::NOT_FOUND, + GetGenericDownloadJobStatusError::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 fmt::Display for GetGenericDownloadJobStatusError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + + +pub async fn get_generic_download_job_status_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data>) -> Result +{ + let maybe_job_status = get_generic_download_job_status( + &path.token, + &server_state.mysql_pool + ) + .await + .map_err(|err| { + warn!("error querying job record: {:?}", err); + return GetGenericDownloadJobStatusError::ServerError + })?; + + let job_status = match maybe_job_status { + Some(record) => record, + None => { + warn!("job record not found"); + return Err(GetGenericDownloadJobStatusError::NotFoundError); + } + }; + + let mut redis = server_state.redis_pool + .get() + .map_err(|e| { + warn!("redis error: {:?}", e); + GetGenericDownloadJobStatusError::ServerError + })?; + + let extra_status_key = RedisKeys::generic_download_extra_status_info(&path.token); + let maybe_extra_status_description : Option = match redis.get(&extra_status_key) { + Ok(Some(status)) => { + Some(status) + }, + Ok(None) => None, + Err(e) => { + warn!("redis error: {:?}", e); + None // Fail open + }, + }; + + let model_for_response = GetGenericDownloadJobStatusForResponse { + job_token: job_status.job_token, + status: job_status.status, + maybe_extra_status_description, + attempt_count: job_status.attempt_count as u8, + maybe_downloaded_entity_token: job_status.maybe_downloaded_entity_token, + maybe_downloaded_entity_type: job_status.maybe_downloaded_entity_type, + created_at: job_status.created_at.clone(), + updated_at: job_status.updated_at.clone(), + }; + + let response = GetGenericDownloadJobStatusSuccessResponse { + success: true, + state: model_for_response, + }; + + let body = serde_json::to_string(&response) + .map_err(|e| GetGenericDownloadJobStatusError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs new file mode 100644 index 0000000000..83ff514ad4 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/download_job/mod.rs @@ -0,0 +1,2 @@ +pub mod enqueue_generic_download; +pub mod get_generic_upload_job_status; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/mod.rs index 4c65b00eba..ffb2ac00e1 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/mod.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/mod.rs @@ -1,5 +1,6 @@ pub mod api_tokens; pub mod categories; +pub mod download_job; pub mod events; pub mod flags; pub mod investor_demo; @@ -10,5 +11,6 @@ pub mod service; pub mod tts; pub mod twitch; pub mod users; +pub mod vocoders; pub mod voice_clone_requests; pub mod w2l; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs index d7865dee60..5c400cb9ab 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/edit_tts_model.rs @@ -1,32 +1,28 @@ -use actix_http::Error; -use actix_http::http::header; -use actix_web::HttpResponseBuilder; -use actix_web::cookie::Cookie; +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use actix_web::error::ResponseError; use actix_web::http::StatusCode; use actix_web::web::{Path, Json}; -use actix_web::{Responder, web, HttpResponse, error, HttpRequest}; -use container_common::i18n::supported_languages_for_models::{is_valid_language_for_models, get_canonicalized_language_tag_for_model}; +use actix_web::{web, HttpResponse, HttpRequest}; +use container_common::i18n::supported_languages_for_models::get_canonicalized_language_tag_for_model; use crate::http_server::web_utils::ip_address::get_request_ip; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::http_server::web_utils::response_success_helpers::simple_json_success; use crate::server_state::ServerState; -use crate::util::email_to_gravatar::email_to_gravatar; use crate::util::markdown_to_html::markdown_to_html; use database_queries::column_types::record_visibility::RecordVisibility; use database_queries::column_types::vocoder_type::VocoderType; -use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; -use language_tags::{LanguageTag, ParseError}; -use log::{info, warn, log, error}; -use regex::Regex; -use sqlx::MySqlPool; -use sqlx::error::DatabaseError; -use sqlx::error::Error::Database; -use sqlx::mysql::MySqlDatabaseError; -use std::fmt; -use std::sync::Arc; use database_queries::queries::tts::tts_models::edit_tts_model_details::{edit_tts_model_details_as_author, edit_tts_model_details_as_mod}; use database_queries::queries::tts::tts_models::edit_tts_model_moderator_details::edit_tts_model_moderator_details; +use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; +use language_tags::LanguageTag; +use log::{info, warn, error}; +use sqlx::MySqlPool; +use std::fmt; +use std::sync::Arc; use tts_common::text_pipelines::text_pipeline_type::TextPipelineType; use user_input_common::check_for_slurs::contains_slurs; @@ -46,19 +42,31 @@ pub struct EditTtsModelRequest { pub title: Option, pub description_markdown: Option, pub creator_set_visibility: Option, - pub maybe_default_pretrained_vocoder: Option, - // Text pipeline won't be set by default on model upload (for now), - // and a lot of existing models will have null values. + // This controls how text will be pre-processed before it is handed off to the TTS + // synthesizer model. Different pipelines treat graphemes, acronyms, arpabet, language + // specifics, and integer encodings in their own way. + // Note that the text pipeline option won't be set by default on new model uploads (for + // now), and that most existing models will have null values here. pub text_pipeline_type: Option, + // [vocoders 1] + // This is the new type of vocoder configuration. Users can choose a custom trained + // vocoder to associate with their model. These tokens reference the `vocoder_models` + // table. + pub maybe_custom_vocoder_token: Option, + + // [vocoders 2] + // This is the old type of vocoder configuration, which leverages old pretrained + // vocoders that we manually uploaded. There aren't many good options for users to + // choose here, so this should be treated as a legacy option going forward. We'll + // likely be stuck with this configuration for some time, however, due to the large + // collection of legacy models we have. + pub maybe_default_pretrained_vocoder: Option, + // NB: We calculate 'ietf_primary_language_subtag' from this value. pub ietf_language_tag: Option, - //pub updatable_slug: Option, - //pub tts_model_type: Option, - //pub vocoder_token: Option, - // ========== Moderator options (protection) ========== pub is_public_listing_approved: Option, @@ -233,10 +241,10 @@ pub async fn edit_tts_model_handler( } } - let mut ietf_language_tag = ietf_language_tag + let ietf_language_tag = ietf_language_tag .unwrap_or(DEFAULT_IETF_LANGUAGE_TAG.to_string()); - let mut ietf_primary_language_subtag = ietf_primary_language_subtag + let ietf_primary_language_subtag = ietf_primary_language_subtag .unwrap_or(DEFAULT_IETF_PRIMARY_LANGUAGE_SUBTAG.to_string()); if let Some(visibility) = request.creator_set_visibility.as_deref() { @@ -244,14 +252,14 @@ pub async fn edit_tts_model_handler( .map_err(|_| EditTtsModelError::BadInput("bad record visibility".to_string()))?; } - if let Some(vocoder) = request.maybe_default_pretrained_vocoder { - maybe_default_pretrained_vocoder = Some(vocoder); - } - let text_pipeline_type = request.text_pipeline_type .map(|pipeline_type| pipeline_type.to_db_variant()) .map(|pipeline_type| pipeline_type.to_str()); + if let Some(vocoder) = request.maybe_default_pretrained_vocoder { + maybe_default_pretrained_vocoder = Some(vocoder); + } + let ip_address = get_request_ip(&http_request); let query_result = if is_author { @@ -266,6 +274,7 @@ pub async fn edit_tts_model_handler( &ietf_primary_language_subtag, creator_set_visibility, maybe_default_pretrained_vocoder, + request.maybe_custom_vocoder_token.as_deref(), text_pipeline_type, &ip_address, ).await @@ -281,6 +290,7 @@ pub async fn edit_tts_model_handler( &ietf_primary_language_subtag, creator_set_visibility, maybe_default_pretrained_vocoder, + request.maybe_custom_vocoder_token.as_deref(), text_pipeline_type, &user_session.user_token, ).await diff --git a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs index eb447e2d53..31e6272a3b 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/tts/get_tts_model.rs @@ -1,24 +1,19 @@ -use actix_http::http::header; -use actix_web::HttpResponseBuilder; -use actix_web::cookie::Cookie; +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use actix_web::error::ResponseError; use actix_web::http::StatusCode; use actix_web::web::Path; -use actix_web::{Responder, web, HttpResponse, error, HttpRequest, HttpMessage}; -use anyhow::Error; +use actix_web::{web, HttpResponse, HttpRequest}; use chrono::{DateTime, Utc}; -use crate::AnyhowResult; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; use crate::server_state::ServerState; use database_queries::column_types::record_visibility::RecordVisibility; use database_queries::column_types::vocoder_type::VocoderType; use database_queries::queries::tts::tts_models::get_tts_model::get_tts_model_by_token; -use log::{info, warn, log}; -use regex::Regex; -use sqlx::MySqlPool; -use sqlx::error::DatabaseError; -use sqlx::error::Error::Database; -use sqlx::mysql::MySqlDatabaseError; +use log::warn; use std::fmt; use std::sync::Arc; use tts_common::text_pipelines::guess_pipeline::guess_text_pipeline_heuristic; @@ -52,10 +47,22 @@ pub struct TtsModelInfo { /// inform the frontend (and inference pipeline) of our best guess according to a heuristic. pub text_pipeline_type: Option, pub text_pipeline_type_guess: TextPipelineType, - - pub maybe_default_pretrained_vocoder: Option, pub text_preprocessing_algorithm: String, + /// [vocoders 1] + /// This is the new type of vocoder configuration. Users can choose a custom trained + /// vocoder to associate with their model. The tokens reference the `vocoder_models` + /// table. + pub maybe_custom_vocoder: Option, + + /// [vocoders 2] + /// This is the old type of vocoder configuration, which leverages old pretrained + /// vocoders that we manually uploaded. There aren't many good options for users to + /// choose here, so this should be treated as a legacy option going forward. We'll + /// likely be stuck with this configuration for some time, however, due to the large + /// collection of legacy models we have. + pub maybe_default_pretrained_vocoder: Option, + pub creator_user_token: String, pub creator_username: String, pub creator_display_name: String, @@ -85,6 +92,17 @@ pub struct TtsModelInfo { pub maybe_moderator_fields: Option, } +/// New vocoder configuration options. +#[derive(Serialize)] +pub struct CustomVocoderInfo { + pub vocoder_token: String, + pub vocoder_title: String, + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, +} + /// "Moderator-only fields" that we wouldn't want to expose to ordinary users. /// It's the web endpoint controller's responsibility to clear these for non-mods. #[derive(Serialize)] @@ -189,7 +207,7 @@ pub async fn get_tts_model_handler( model.maybe_moderator_fields = None; } - let mut text_pipeline_type = model.text_pipeline_type + let text_pipeline_type = model.text_pipeline_type .as_deref() .and_then(|pipeline_type| { // If there's an error deserializing, turn it to None instead of 500ing. The column is @@ -219,6 +237,17 @@ pub async fn get_tts_model_handler( tts_model_type: model.tts_model_type, text_pipeline_type, text_pipeline_type_guess, + maybe_custom_vocoder: match model.maybe_custom_vocoder { + None => None, + Some(vocoder) => Some(CustomVocoderInfo { + vocoder_token: vocoder.vocoder_token, + vocoder_title: vocoder.vocoder_title, + creator_user_token: vocoder.creator_user_token, + creator_username: vocoder.creator_username, + creator_display_name: vocoder.creator_display_name, + creator_gravatar_hash: vocoder.creator_gravatar_hash, + }) + }, maybe_default_pretrained_vocoder: model.maybe_default_pretrained_vocoder, text_preprocessing_algorithm: model.text_preprocessing_algorithm, creator_user_token: model.creator_user_token, @@ -251,7 +280,7 @@ pub async fn get_tts_model_handler( }; let body = serde_json::to_string(&response) - .map_err(|e| GetTtsModelError::ServerError)?; + .map_err(|_e| GetTtsModelError::ServerError)?; Ok(HttpResponse::Ok() .content_type("application/json") diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_vocoders.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/list_user_vocoders.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs index 8d5324ab12..39f8c475ae 100644 --- a/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs +++ b/crates/service/storyteller_web/src/http_server/endpoints/users/mod.rs @@ -3,6 +3,7 @@ pub mod edit_profile; pub mod get_profile; pub mod list_user_tts_inference_results; pub mod list_user_tts_models; +pub mod list_user_vocoders; pub mod list_user_w2l_inference_results; pub mod list_user_w2l_templates; pub mod login; diff --git a/crates/service/storyteller_web/src/http_server/endpoints/vocoders/delete_vocoder.rs b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/delete_vocoder.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/service/storyteller_web/src/http_server/endpoints/vocoders/edit_vocoder.rs b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/edit_vocoder.rs new file mode 100644 index 0000000000..e69de29bb2 diff --git a/crates/service/storyteller_web/src/http_server/endpoints/vocoders/get_vocoder.rs b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/get_vocoder.rs new file mode 100644 index 0000000000..d71057def2 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/get_vocoder.rs @@ -0,0 +1,190 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::web::Path; +use actix_web::{web, HttpResponse, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::vocoder::get_vocoder_model::get_vocoder_model_by_token; +use log::warn; +use reusable_types::entity_visibility::EntityVisibility; +use reusable_types::vocoder_type::VocoderType; +use std::fmt; +use std::sync::Arc; + +// =============== Request =============== + +/// For the URL PathInfo +#[derive(Deserialize)] +pub struct GetVocoderPathInfo { + token: String, +} + +// =============== Success Response =============== + +#[derive(Serialize)] +pub struct GetVocoderSuccessResponse { + pub success: bool, + pub vocoder: Vocoder, +} + +#[derive(Serialize)] +pub struct Vocoder { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub description_markdown: String, + pub description_rendered_html: String, + + pub is_staff_recommended: bool, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub creator_set_visibility: EntityVisibility, + + pub created_at: DateTime, + pub updated_at: DateTime, + + /// Moderator fields are absent if not a moderator. + #[serde(skip_serializing_if = "Option::is_none")] + pub moderator_fields: Option, +} + +#[derive(Serialize)] +pub struct VocoderModFields { + pub creator_is_banned: bool, + pub creator_ip_address_creation: String, + pub creator_ip_address_last_update: String, + + pub is_mod_disabled_from_public_use: bool, + pub is_mod_disabled_from_author_use: bool, + pub is_mod_author_editing_locked: bool, + + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + +// =============== Error Response =============== + +#[derive(Debug)] +pub enum GetVocoderError { + NotFound, + ServerError, +} + +impl ResponseError for GetVocoderError { + fn status_code(&self) -> StatusCode { + match *self { + GetVocoderError::NotFound => StatusCode::NOT_FOUND, + GetVocoderError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + GetVocoderError::NotFound => "not found".to_string(), + GetVocoderError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for GetVocoderError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn get_vocoder_handler( + http_request: HttpRequest, + path: Path, + server_state: web::Data> +) -> Result { + + let is_moderator = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + GetVocoderError::ServerError + })? + .map(|session| { + // NB: Since we need to rip out and replace the permissions system, + // this is a proxy for being a moderator. + session.can_ban_users + }) + .unwrap_or(false); + + const NO_CREATOR_SCOPING_HERE : Option<&'static str> = None; + + let query_result = get_vocoder_model_by_token( + &path.token, + is_moderator, + &server_state.mysql_pool, + ).await; + + let vocoder = match query_result { + Ok(Some(model)) => model, + Ok(None) => return Err(GetVocoderError::NotFound), + Err(e) => { + warn!("query error: {:?}", e); + return Err(GetVocoderError::ServerError); + } + }; + + let mut vocoder = Vocoder { + vocoder_token: vocoder.vocoder_token, + vocoder_type: vocoder.vocoder_type, + title: vocoder.title, + description_markdown: vocoder.description_markdown, + description_rendered_html: vocoder.description_rendered_html, + is_staff_recommended: vocoder.is_staff_recommended, + creator_user_token: vocoder.creator_user_token, + creator_username: vocoder.creator_username, + creator_display_name: vocoder.creator_display_name, + creator_gravatar_hash: vocoder.creator_gravatar_hash, + creator_set_visibility: vocoder.creator_set_visibility, + created_at: vocoder.created_at, + updated_at: vocoder.updated_at, + moderator_fields: vocoder.maybe_moderator_fields.map(|mod_fields| { + VocoderModFields { + creator_is_banned: mod_fields.creator_is_banned, + creator_ip_address_creation: mod_fields.creator_ip_address_creation, + creator_ip_address_last_update: mod_fields.creator_ip_address_last_update, + is_mod_disabled_from_public_use: mod_fields.is_mod_disabled_from_public_use, + is_mod_disabled_from_author_use: mod_fields.is_mod_disabled_from_author_use, + is_mod_author_editing_locked: mod_fields.is_mod_author_editing_locked, + user_deleted_at: mod_fields.user_deleted_at, + mod_deleted_at: mod_fields.mod_deleted_at, + } + }) + }; + + if !is_moderator { + vocoder.moderator_fields = None; + } + + let response = GetVocoderSuccessResponse { + success: true, + vocoder, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| GetVocoderError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/vocoders/list_vocoders.rs b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/list_vocoders.rs new file mode 100644 index 0000000000..6cebb229c1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/list_vocoders.rs @@ -0,0 +1,161 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use actix_web::{web, HttpResponse, HttpRequest}; +use chrono::{DateTime, Utc}; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; +use database_queries::queries::vocoder::list_vocoder_models::{list_vocoder_models, VocoderModelListItem}; +use log::warn; +use reusable_types::vocoder_type::VocoderType; +use std::fmt; +use std::sync::Arc; + +#[derive(Serialize)] +pub struct ListVocodersSuccessResponse { + pub success: bool, + pub vocoders: Vec, +} + +#[derive(Serialize)] +pub struct VocoderListItem { + pub vocoder_token: String, + pub vocoder_type: VocoderType, + + pub title: String, + pub is_staff_recommended: bool, + + pub creator_user_token: String, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + pub created_at: DateTime, + pub updated_at: DateTime, + + /// Moderator fields are absent if not a moderator. + #[serde(skip_serializing_if = "Option::is_none")] + pub moderator_fields: Option, +} + +#[derive(Serialize)] +pub struct VocoderListItemModFields { + pub is_mod_disabled_from_public_use: bool, + pub is_mod_disabled_from_author_use: bool, + pub is_mod_author_editing_locked: bool, + pub user_deleted_at: Option>, + pub mod_deleted_at: Option>, +} + + +#[derive(Debug)] +pub enum ListVocodersError { + ServerError, +} + +impl ResponseError for ListVocodersError { + fn status_code(&self) -> StatusCode { + match *self { + ListVocodersError::ServerError=> StatusCode::INTERNAL_SERVER_ERROR, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + ListVocodersError::ServerError => "server error".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +// NB: Not using derive_more::Display since Clion doesn't understand it. +impl fmt::Display for ListVocodersError { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!(f, "{:?}", self) + } +} + +pub async fn list_vocoders_handler( + http_request: HttpRequest, + server_state: web::Data> +) -> Result { + + let is_moderator = server_state + .session_checker + .maybe_get_user_session(&http_request, &server_state.mysql_pool) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + ListVocodersError::ServerError + })? + .map(|session| { + // NB: Since we need to rip out and replace the permissions system, + // this is a proxy for being a moderator. + session.can_ban_users + }) + .unwrap_or(false); + + const NO_CREATOR_SCOPING_HERE : Option<&'static str> = None; + + let query_results = list_vocoder_models( + &server_state.mysql_pool, + NO_CREATOR_SCOPING_HERE, + false, + ).await; + + let vocoders = match query_results { + Ok(results) => results, + Err(e) => { + warn!("vocoder list query error: {:?}", e); + return Err(ListVocodersError::ServerError); + } + }; + + let vocoders = vocoders + .into_iter() + .map(|v: VocoderModelListItem| { + let mut vocoder = VocoderListItem { + vocoder_token: v.vocoder_token, + vocoder_type: v.vocoder_type, + title: v.title, + is_staff_recommended: v.is_staff_recommended, + creator_user_token: v.creator_user_token, + creator_username: v.creator_username, + creator_display_name: v.creator_display_name, + creator_gravatar_hash: v.creator_gravatar_hash, + created_at: v.created_at, + updated_at: v.updated_at, + moderator_fields: Some(VocoderListItemModFields { + is_mod_disabled_from_public_use: v.moderator_fields.is_mod_disabled_from_public_use, + is_mod_disabled_from_author_use: v.moderator_fields.is_mod_disabled_from_author_use, + is_mod_author_editing_locked: v.moderator_fields.is_mod_author_editing_locked, + user_deleted_at: v.moderator_fields.user_deleted_at, + mod_deleted_at:v.moderator_fields.mod_deleted_at, + }) + }; + + if !is_moderator { + vocoder.moderator_fields = None; + } + + vocoder + }) + .collect(); + + let response = ListVocodersSuccessResponse { + success: true, + vocoders, + }; + + let body = serde_json::to_string(&response) + .map_err(|_e| ListVocodersError::ServerError)?; + + Ok(HttpResponse::Ok() + .content_type("application/json") + .body(body)) +} diff --git a/crates/service/storyteller_web/src/http_server/endpoints/vocoders/mod.rs b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/mod.rs new file mode 100644 index 0000000000..ca2786f4d1 --- /dev/null +++ b/crates/service/storyteller_web/src/http_server/endpoints/vocoders/mod.rs @@ -0,0 +1,4 @@ +pub mod delete_vocoder; +pub mod edit_vocoder; +pub mod get_vocoder; +pub mod list_vocoders; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/routes.rs b/crates/service/storyteller_web/src/routes.rs index 11146bcbd0..b3f12e7e6b 100644 --- a/crates/service/storyteller_web/src/routes.rs +++ b/crates/service/storyteller_web/src/routes.rs @@ -12,6 +12,8 @@ use crate::http_server::endpoints::categories::create_category::create_category_ use crate::http_server::endpoints::categories::get_category::get_category_handler; use crate::http_server::endpoints::categories::list_tts_categories::list_tts_categories_handler; use crate::http_server::endpoints::categories::list_tts_model_assigned_categories::list_tts_model_assigned_categories_handler; +use crate::http_server::endpoints::download_job::enqueue_generic_download::enqueue_generic_download_handler; +use crate::http_server::endpoints::download_job::get_generic_upload_job_status::get_generic_download_job_status_handler; use crate::http_server::endpoints::events::list_events::list_events_handler; use crate::http_server::endpoints::flags::design_refresh_flag::disable_design_refresh_flag_handler::disable_design_refresh_flag_handler; use crate::http_server::endpoints::flags::design_refresh_flag::enable_design_refresh_flag_handler::enable_design_refresh_flag_handler; @@ -75,6 +77,8 @@ use crate::http_server::endpoints::users::list_user_w2l_templates::list_user_w2l use crate::http_server::endpoints::users::login::login_handler; use crate::http_server::endpoints::users::logout::logout_handler; use crate::http_server::endpoints::users::session_info::session_info_handler; +use crate::http_server::endpoints::vocoders::get_vocoder::get_vocoder_handler; +use crate::http_server::endpoints::vocoders::list_vocoders::list_vocoders_handler; use crate::http_server::endpoints::voice_clone_requests::check_if_voice_clone_request_submitted::check_if_voice_clone_request_submitted_handler; use crate::http_server::endpoints::voice_clone_requests::create_voice_clone_request::create_voice_clone_request_handler; use crate::http_server::endpoints::w2l::delete_w2l_result::delete_w2l_inference_result_handler; @@ -106,6 +110,8 @@ pub fn add_routes (app: App) -> App let mut app = add_moderator_routes(app); /* /moderation */ app = add_tts_routes(app); /* /tts */ app = add_w2l_routes(app); /* /w2l */ + app = add_vocoder_routes(app); /* /vocoder */ + app = add_retrieval_routes(app); /* /retrieval (aka. "generic_download_jobs") */ app = add_category_routes(app); /* /category */ app = add_user_profile_routes(app); /* /user */ app = add_api_token_routes(app); /* /api_tokens */ @@ -469,6 +475,72 @@ fn add_w2l_routes (app: App) -> App ) } +// ==================== VOCODER ROUTES ==================== + +fn add_vocoder_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/vocoder") + .service( + web::resource("/list") + .route(web::get().to(list_vocoders_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/model/{token}") + .route(web::get().to(get_vocoder_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + //.service( + // web::resource("/model/{token}/edit") + // .route(web::post().to(edit_w2l_template_handler)) + // .route(web::head().to(|| HttpResponse::Ok())) + //) + //.service( + // web::resource("/model/{token}/delete") + // .route(web::post().to(delete_w2l_template_handler)) + // .route(web::head().to(|| HttpResponse::Ok())) + //) + ) +} + +// ==================== RETRIEVAL ROUTES ("GENERIC_DOWNLOAD_JOBS") ==================== + +fn add_retrieval_routes (app: App) -> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + app.service( + web::scope("/retrieval") + .service( + web::resource("/enqueue") + .route(web::post().to(enqueue_generic_download_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service( + web::resource("/job_status/{token}") + .route(web::get().to(get_generic_download_job_status_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) +} + // ==================== CATEGORY ROUTES ==================== fn add_category_routes (app: App) -> App diff --git a/crates/service/storyteller_web/src/validations/mod.rs b/crates/service/storyteller_web/src/validations/mod.rs index e881b761a1..073fcfcec1 100644 --- a/crates/service/storyteller_web/src/validations/mod.rs +++ b/crates/service/storyteller_web/src/validations/mod.rs @@ -8,4 +8,5 @@ pub mod twitch_username; pub mod twitter_username; pub mod username; pub mod username_reservations; +pub mod validate_idempotency_token_format; pub mod website_url; \ No newline at end of file diff --git a/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs b/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs new file mode 100644 index 0000000000..46e8e9adba --- /dev/null +++ b/crates/service/storyteller_web/src/validations/validate_idempotency_token_format.rs @@ -0,0 +1,43 @@ +use once_cell::sync::Lazy; +use regex::Regex; + +static EXPECTED_FORMAT : Lazy = Lazy::new(|| { + // Support UUID with or without the dashes. + Regex::new(r"^(([a-zA-Z0-9]{32})|([a-zA-Z0-9\-]{36}))$").expect("should be valid regex") +}); + +/// Validate the idempotency token +/// Idempotency tokens are supposed to be random entropy from the frontend to prevent duplicate +/// creation/insertion requests from being sent and acknowledged by the database. +/// We expect them to be in UUID format. +pub fn validate_idempotency_token_format(token: &str) -> Result<(), String> { + if !EXPECTED_FORMAT.is_match(token) { + return Err("idempotency token does not match the expected format".to_string()); + } + Ok(()) +} + +#[cfg(test)] +mod tests { + use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; + + #[test] + fn valid_idempotency_tokens() { + // UUID with the dashes + assert!(validate_idempotency_token_format("bf12126a-2de0-4527-bbbf-156c4c45ae21").is_ok()); + assert!(validate_idempotency_token_format("220ace1b-db32-42c4-bb87-49176c85a907").is_ok()); + + // UUID without the dashes + assert!(validate_idempotency_token_format("220ace1bdb3242c4bb8749176c85a907").is_ok()); + } + + #[test] + fn invalid_idempotency_tokens() { + assert!(!validate_idempotency_token_format("").is_ok()); + assert!(!validate_idempotency_token_format(" ").is_ok()); + assert!(!validate_idempotency_token_format("\n\n\n\n").is_ok()); + assert!(!validate_idempotency_token_format("bob").is_ok()); + assert!(!validate_idempotency_token_format("220ace1bdb324").is_ok()); + assert!(!validate_idempotency_token_format("220ace1bdb3242c4bb8749176c85a907ace0b0def0ff0").is_ok()); + } +} diff --git a/crates/service/tts_inference_job/src/http_clients/tts_inference_sidecar_client.rs b/crates/service/tts_inference_job/src/http_clients/tts_inference_sidecar_client.rs index 0906194bae..8cf86635d6 100644 --- a/crates/service/tts_inference_job/src/http_clients/tts_inference_sidecar_client.rs +++ b/crates/service/tts_inference_job/src/http_clients/tts_inference_sidecar_client.rs @@ -48,6 +48,7 @@ impl TtsInferenceSidecarClient { } } + /// NB: 'hifigan_vocoder_checkpoint_path' may be either a pretrained or custom vocoder pub async fn request_inference>( &self, raw_text: &str, diff --git a/crates/service/tts_inference_job/src/job_steps/process_single_job.rs b/crates/service/tts_inference_job/src/job_steps/process_single_job.rs index 83c12a0d41..935d67b5c0 100644 --- a/crates/service/tts_inference_job/src/job_steps/process_single_job.rs +++ b/crates/service/tts_inference_job/src/job_steps/process_single_job.rs @@ -1,3 +1,8 @@ +// NB: Incrementally getting rid of build warnings... +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + use anyhow::anyhow; use container_common::anyhow_result::AnyhowResult; use container_common::filesystem::check_file_exists::check_file_exists; @@ -121,7 +126,7 @@ pub async fn process_single_job( // ==================== CONFIRM OR DOWNLOAD HIFIGAN (NORMAL) VOCODER MODEL ==================== // - let hifigan_vocoder_model_fs_path = { + let pretrained_hifigan_vocoder_model_fs_path = { let hifigan_vocoder_model_filename = inferencer.hifigan_vocoder_model_filename.clone(); let hifigan_vocoder_model_fs_path = inferencer.semi_persistent_cache.tts_pretrained_vocoder_model_path(&hifigan_vocoder_model_filename); let hifigan_vocoder_model_object_path = inferencer.bucket_path_unifier.tts_pretrained_vocoders_path(&hifigan_vocoder_model_filename); @@ -161,27 +166,31 @@ pub async fn process_single_job( hifigan_superres_vocoder_model_fs_path }; -// // ==================== LOOK UP TTS SYNTHESIZER RECORD (WHICH CONTAINS ITS BUCKET PATH) ==================== // -// -// info!("Looking up TTS model by token: {}", &job.model_token); -// -// let query_result = get_tts_model_by_token( -// &inferencer.mysql_pool, -// &job.model_token).await?; -// -// let tts_model = match query_result { -// Some(model) => model, -// None => { -// warn!("TTS model not found: {}", &job.model_token); -// return Err(anyhow!("Model not found!")) -// }, -// }; +// // ==================== CONFIRM OR DOWNLOAD OPTIONAL CUSTOM VOCODER MODEL ==================== // + + let custom_vocoder_fs_path = match &model_record.maybe_custom_vocoder { + None => None, + Some(vocoder) => { + let custom_vocoder_fs_path = inferencer.semi_persistent_cache.custom_vocoder_model_path(&vocoder.vocoder_token); + let custom_vocoder_object_path = inferencer.bucket_path_unifier.vocoder_path(&vocoder.vocoder_private_bucket_hash); + + maybe_download_file_from_bucket( + "custom vocoder", + &custom_vocoder_fs_path, + &custom_vocoder_object_path, + &inferencer.private_bucket_client, + &mut redis_logger, + "downloading user vocoder", + job.id.0, + &inferencer.scoped_temp_dir_creator, + ).await?; + + Some(custom_vocoder_fs_path) + } + }; // ==================== CONFIRM OR DOWNLOAD TTS SYNTHESIZER MODEL ==================== // - // TODO: Let's just put paths in the db - // TODO: We'll probably need to LRU cache these. - let tts_synthesizer_fs_path = { let tts_synthesizer_fs_path = inferencer.semi_persistent_cache.tts_synthesizer_model_path(&model_record.model_token); let tts_synthesizer_object_path = inferencer.bucket_path_unifier.tts_synthesizer_path(&model_record.private_bucket_hash); @@ -248,16 +257,6 @@ pub async fn process_single_job( warn!("Unload model from sidecar: {:?}", &model_path); } - //inferencer.tts_inference_command.execute( - // &tts_synthesizer_fs_path, - // &tts_vocoder_model_fs_path, - // &text_input_fs_path, - // &output_audio_fs_path, - // &output_spectrogram_fs_path, - // &output_metadata_fs_path, - // false, - //)?; - let mut pretrained_vocoder = VocoderType::HifiGanSuperResolution; if let Some(default_vocoder) = model_record.maybe_default_pretrained_vocoder.as_deref() { pretrained_vocoder = VocoderType::from_str(default_vocoder) @@ -277,12 +276,23 @@ pub async fn process_single_job( &model_record.text_pipeline_type, &text_pipeline_type_or_guess); + let hifigan_vocoder_model_fs_path_to_use = match custom_vocoder_fs_path { + None => { + info!("using pretrained HiFi-GAN vocoder"); + pretrained_hifigan_vocoder_model_fs_path + }, + Some(custom_vocoder_fs_path) => { + info!("using custom user-trained HiFi-GAN vocoder: {:?}", custom_vocoder_fs_path); + custom_vocoder_fs_path + }, + }; + inferencer.http_clients.tts_inference_sidecar_client.request_inference( &cleaned_inference_text, &tts_synthesizer_fs_path, pretrained_vocoder, &text_pipeline_type_or_guess.to_str(), - &hifigan_vocoder_model_fs_path, + &hifigan_vocoder_model_fs_path_to_use, &hifigan_superres_vocoder_model_fs_path, &waveglow_vocoder_model_fs_path, &output_audio_fs_path, diff --git a/crates/service/tts_inference_job/src/main.rs b/crates/service/tts_inference_job/src/main.rs index 36958fb6b1..f7a118d365 100644 --- a/crates/service/tts_inference_job/src/main.rs +++ b/crates/service/tts_inference_job/src/main.rs @@ -202,6 +202,7 @@ async fn main() -> AnyhowResult<()> { info!("Creating pod semi-persistent cache dirs..."); semi_persistent_cache.create_tts_synthesizer_model_path()?; semi_persistent_cache.create_tts_pretrained_vocoder_model_path()?; + semi_persistent_cache.create_custom_vocoder_model_path()?; let waveglow_vocoder_model_filename = easyenv::get_env_string_or_default( "TTS_WAVEGLOW_VOCODER_MODEL_FILENAME", "waveglow.pth");