diff --git a/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json b/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json new file mode 100644 index 0000000000..8fc92594f2 --- /dev/null +++ b/.sqlx/query-1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nINSERT INTO model_weights_extension_tts_details\nSET\n model_weights_token = ?,\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n text_pipeline_type = ?,\n use_default_mel_multiply_factor = ?,\n maybe_custom_mel_multiply_factor = ?,\n maybe_default_pretrained_vocoder = ?,\n maybe_custom_vocoder_token = ?\nON DUPLICATE KEY UPDATE\n ietf_language_tag = ?,\n ietf_primary_language_subtag = ?,\n text_pipeline_type = ?,\n use_default_mel_multiply_factor = ?,\n maybe_custom_mel_multiply_factor = ?,\n maybe_default_pretrained_vocoder = ?,\n maybe_custom_vocoder_token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 15 + }, + "nullable": [] + }, + "hash": "1ce9219b8f05fe54adaa6841acb708e165e1c7855b441b9623a2151197578929" +} diff --git a/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json b/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json new file mode 100644 index 0000000000..0f795c35d7 --- /dev/null +++ b/.sqlx/query-4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5.json @@ -0,0 +1,146 @@ +{ + "db_name": "MySQL", + "query": "\nSELECT\n w.token as `token: tokens::tokens::model_weights::ModelWeightToken`,\n w.weights_type as `weight_type: enums::by_table::model_weights::weights_types::WeightsType`,\n w.title,\n w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`,\n users.username as creator_username,\n users.display_name as creator_display_name,\n users.email_gravatar_hash as creator_gravatar_hash,\n w_extension.ietf_language_tag,\n w_extension.ietf_primary_language_subtag,\n w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`,\n w.created_at,\n w.updated_at\nFROM model_weights as w\nLEFT OUTER JOIN model_weights_extension_tts_details as w_extension\n ON w_extension.model_weights_token = w.token\nJOIN users\n ON users.token = w.creator_user_token\nWHERE\n w.weights_type IN (\"tt2\")\n AND w.user_deleted_at IS NULL\n AND w.mod_deleted_at IS NULL\n ", + "describe": { + "columns": [ + { + "ordinal": 0, + "name": "token: tokens::tokens::model_weights::ModelWeightToken", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 1, + "name": "weight_type: enums::by_table::model_weights::weights_types::WeightsType", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 2, + "name": "title", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 1020 + } + }, + { + "ordinal": 3, + "name": "creator_user_token: tokens::tokens::users::UserToken", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 4, + "name": "creator_username", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | UNIQUE_KEY | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 5, + "name": "creator_display_name", + "type_info": { + "type": "VarString", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 80 + } + }, + { + "ordinal": 6, + "name": "creator_gravatar_hash", + "type_info": { + "type": "String", + "flags": "NOT_NULL | BINARY | NO_DEFAULT_VALUE", + "char_set": 224, + "max_size": 128 + } + }, + { + "ordinal": 7, + "name": "ietf_language_tag", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 256 + } + }, + { + "ordinal": 8, + "name": "ietf_primary_language_subtag", + "type_info": { + "type": "VarString", + "flags": "BINARY", + "char_set": 224, + "max_size": 48 + } + }, + { + "ordinal": 9, + "name": "creator_set_visibility: enums::common::visibility::Visibility", + "type_info": { + "type": "String", + "flags": "NOT_NULL | MULTIPLE_KEY | BINARY | ENUM", + "char_set": 224, + "max_size": 28 + } + }, + { + "ordinal": 10, + "name": "created_at", + "type_info": { + "type": "Timestamp", + "flags": "NOT_NULL | BINARY | TIMESTAMP", + "char_set": 63, + "max_size": 19 + } + }, + { + "ordinal": 11, + "name": "updated_at", + "type_info": { + "type": "Timestamp", + "flags": "NOT_NULL | BINARY | TIMESTAMP | ON_UPDATE_NOW", + "char_set": 63, + "max_size": 19 + } + } + ], + "parameters": { + "Right": 0 + }, + "nullable": [ + false, + false, + false, + false, + false, + false, + false, + true, + true, + false, + false, + false + ] + }, + "hash": "4d45f5ce213627465e991eaec9b509100e22c614367b67ab026b5edd70927ea5" +} diff --git a/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json b/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json new file mode 100644 index 0000000000..34731cd396 --- /dev/null +++ b/.sqlx/query-aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112.json @@ -0,0 +1,12 @@ +{ + "db_name": "MySQL", + "query": "\nUPDATE tts_models\nSET\n maybe_migration_new_model_weights_token = ?\nWHERE token = ?\n ", + "describe": { + "columns": [], + "parameters": { + "Right": 2 + }, + "nullable": [] + }, + "hash": "aa7ec1c1e32f773212b264f34173f99db1e515e370c47065f932aaec56f89112" +} diff --git a/Cargo.lock b/Cargo.lock index 5e29620c49..fcc378de3e 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1259,6 +1259,24 @@ dependencies = [ "my-workspace-hack", ] +[[package]] +name = "composite_identifiers" +version = "0.0.1" +dependencies = [ + "enums", + "my-workspace-hack", + "serde", + "serde_derive", + "serde_json", + "sqlx", + "sqlx-core", + "sqlx-mysql", + "strum", + "tokens", + "toml 0.5.11", + "utoipa", +] + [[package]] name = "concurrency" version = "0.0.1" @@ -3052,6 +3070,7 @@ dependencies = [ "chrono", "cloud_storage", "collections", + "composite_identifiers", "concurrency", "config", "dotenv", @@ -3059,6 +3078,7 @@ dependencies = [ "enums", "errors", "filesys", + "google_drive_common", "hashing", "hostname", "hyper", @@ -3857,6 +3877,7 @@ dependencies = [ "anyhow", "buckets", "chrono", + "composite_identifiers", "config", "container_common", "dotenv", @@ -6006,6 +6027,7 @@ dependencies = [ "buckets", "chrono", "cloud_storage", + "composite_identifiers", "config", "container_common", "cookies", @@ -6068,6 +6090,7 @@ dependencies = [ "sha2 0.9.9", "speculoos", "sqlx", + "strum", "symphonia", "time 0.3.30", "tokens", diff --git a/Cargo.toml b/Cargo.toml index 0bd2f6e824..4b04d5e36c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -5,6 +5,7 @@ members = [ # Types: common strongly-typed tokens, ids, payloads, etc. # Also common cloud buckets and other "strongly typed" configs. "crates/_types/buckets", + "crates/_types/composite_identifiers", "crates/_types/enums", "crates/_types/reusable_types", # TODO(bt, 2022-12-21): Deprecate "crates/_types/tokens", diff --git a/_sql/analytics/tts/tts_inference_jobs.sql b/_sql/analytics/tts/tts_inference_jobs.sql index 17af4d8bde..171de2b1a3 100644 --- a/_sql/analytics/tts/tts_inference_jobs.sql +++ b/_sql/analytics/tts/tts_inference_jobs.sql @@ -2,6 +2,20 @@ -- noinspection SqlNoDataSourceInspectionForFile -- noinspection SqlResolveForFile + +-- Find jobs that are broken for a given user +select id, last_assigned_worker, internal_debugging_failure_reason,created_at +from tts_inference_jobs +where maybe_creator_user_token IN ( + select token + from users + where username IN ('knifecat ', 'rewritten_code', 'johnloberger') +) + and attempt_count > 1 + and status IN ('complete_failure', 'dead', 'attempt_failed') +and created_at > ( CURDATE() - INTERVAL 5 DAY ); + + -- Total pending jobs select count(*) from tts_inference_jobs diff --git a/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql new file mode 100644 index 0000000000..b9a1398228 --- /dev/null +++ b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE model_weights_extension_tts_details DROP COLUMN text_pipeline_type; diff --git a/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql new file mode 100644 index 0000000000..30aefdf219 --- /dev/null +++ b/_sql/migrations/2024-01-14-030220_alter_model_weights_extension_tts_add_pipeline/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE model_weights_extension_tts_details + ADD COLUMN text_pipeline_type VARCHAR(64) DEFAULT NULL + AFTER ietf_primary_language_subtag; diff --git a/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql b/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql new file mode 100644 index 0000000000..ab73e06165 --- /dev/null +++ b/_sql/migrations/2024-01-16-231833_idempotency_tokens/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table idempotency_tokens; \ No newline at end of file diff --git a/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql b/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql new file mode 100644 index 0000000000..5667756d56 --- /dev/null +++ b/_sql/migrations/2024-01-16-231833_idempotency_tokens/up.sql @@ -0,0 +1,19 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- A form of optimistic locking +CREATE TABLE idempotency_tokens ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- This can be a UUID or any entropic value. + idempotency_value VARCHAR(36) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (idempotency_value) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations/2024-01-16-231841_batch_generations/down.sql b/_sql/migrations/2024-01-16-231841_batch_generations/down.sql new file mode 100644 index 0000000000..419c98f550 --- /dev/null +++ b/_sql/migrations/2024-01-16-231841_batch_generations/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +drop table batch_generations; diff --git a/_sql/migrations/2024-01-16-231841_batch_generations/up.sql b/_sql/migrations/2024-01-16-231841_batch_generations/up.sql new file mode 100644 index 0000000000..d738abb9f5 --- /dev/null +++ b/_sql/migrations/2024-01-16-231841_batch_generations/up.sql @@ -0,0 +1,42 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Batch generations +CREATE TABLE batch_generations ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- The batch token + -- This is **NOT** a unique token. + token VARCHAR(32) NOT NULL, + + -- ========== ENTITY LINKAGE ========== + + -- The type of record this is attached to + -- Currently supported values: + -- * 'model_weight' + -- * 'media_file' + -- * 'comment' + entity_type VARCHAR(32) NOT NULL, + + -- The token for the record. + entity_token VARCHAR(32) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token, entity_type, entity_token), + KEY index_token (token), + KEY index_entity_token (entity_token) + + -- NB: Realistically, we probably don't need this index until we have a lot of entity types, and the entity_token + -- entropy is probably sufficient enough as-is. + -- KEY index_entity_type_and_token (entity_type, entity_token), + + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql new file mode 100644 index 0000000000..924a36d70c --- /dev/null +++ b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/down.sql @@ -0,0 +1,5 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files DROP COLUMN is_batch_generated; diff --git a/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql new file mode 100644 index 0000000000..1cbc612de5 --- /dev/null +++ b/_sql/migrations/2024-01-16-231852_alter_media_files_add_batch_token/up.sql @@ -0,0 +1,7 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +ALTER TABLE media_files + ADD COLUMN is_batch_generated BOOLEAN NOT NULL DEFAULT FALSE + AFTER maybe_origin_filename; diff --git a/_sql/migrations_squashed/batch_generations.sql b/_sql/migrations_squashed/batch_generations.sql new file mode 100644 index 0000000000..d738abb9f5 --- /dev/null +++ b/_sql/migrations_squashed/batch_generations.sql @@ -0,0 +1,42 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- Batch generations +CREATE TABLE batch_generations ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- The batch token + -- This is **NOT** a unique token. + token VARCHAR(32) NOT NULL, + + -- ========== ENTITY LINKAGE ========== + + -- The type of record this is attached to + -- Currently supported values: + -- * 'model_weight' + -- * 'media_file' + -- * 'comment' + entity_type VARCHAR(32) NOT NULL, + + -- The token for the record. + entity_token VARCHAR(32) NOT NULL, + + -- ========== TIMESTAMPS ========== + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + updated_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP ON UPDATE CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (token, entity_type, entity_token), + KEY index_token (token), + KEY index_entity_token (entity_token) + + -- NB: Realistically, we probably don't need this index until we have a lot of entity types, and the entity_token + -- entropy is probably sufficient enough as-is. + -- KEY index_entity_type_and_token (entity_type, entity_token), + + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations_squashed/idempotency_tokens.sql b/_sql/migrations_squashed/idempotency_tokens.sql new file mode 100644 index 0000000000..5667756d56 --- /dev/null +++ b/_sql/migrations_squashed/idempotency_tokens.sql @@ -0,0 +1,19 @@ +-- noinspection SqlDialectInspectionForFile +-- noinspection SqlNoDataSourceInspectionForFile +-- noinspection SqlResolveForFile + +-- A form of optimistic locking +CREATE TABLE idempotency_tokens ( + -- Not used for anything except replication. + id BIGINT(20) NOT NULL AUTO_INCREMENT, + + -- This can be a UUID or any entropic value. + idempotency_value VARCHAR(36) NOT NULL, + + created_at TIMESTAMP NOT NULL DEFAULT CURRENT_TIMESTAMP, + + -- INDICES -- + PRIMARY KEY (id), + UNIQUE KEY (idempotency_value) + +) ENGINE=InnoDB DEFAULT CHARSET=utf8mb4 COLLATE=utf8mb4_bin; diff --git a/_sql/migrations_squashed/media_files.sql b/_sql/migrations_squashed/media_files.sql index 90fc13febb..35d7334ca9 100644 --- a/_sql/migrations_squashed/media_files.sql +++ b/_sql/migrations_squashed/media_files.sql @@ -26,7 +26,7 @@ CREATE TABLE media_files ( -- * 'device_api' for direct user uploads recorded using Browser/Device APIs. origin_category VARCHAR(16) NOT NULL, - -- TODO: Remove default value once records backfilled and queries updated. + -- TODO(bt,2024-01-12): Rename to origin_product. -- Product area where the media file originated. -- This is not the *model* that created the thing, this is the *product*. -- (The underlying models can change over time.) @@ -37,8 +37,13 @@ CREATE TABLE media_files ( -- "Voice Designer" dataset samples in a video generation flow.) -- -- Possible values: + -- * 'unknown' for legacy records without an associated product -- * 'face_animator' for uploads and outputs - -- * 'unknown' for legacy records without an associated product (TODO: This should be temporary) + -- * 'tts' for text to speech outputs + -- * 'voice_conversion' for uploads or outputs (RVC or SVC) + -- * 'zs_voice' for uploads or outputs for zero shot voice products + -- * 'video_filter' filters on videos + -- * 'mocap' for files uploaded or processed by motion capture origin_product_category VARCHAR(16) NOT NULL DEFAULT "unknown", -- For inference that can be tied back to a model, the type of model. @@ -53,6 +58,11 @@ CREATE TABLE media_files ( -- The original filename of the media (if uploaded by a user) maybe_origin_filename VARCHAR(255) DEFAULT NULL, + -- Whether this media file was generated as part of a batch + -- If so, we can look up the batch in a separate query/call to the `batch_generations` table. + -- We won't hold the batch token in this table since it'll be very sparsely populated. + is_batch_generated BOOLEAN NOT NULL DEFAULT FALSE, + -- ========== METADATA ========== -- We can store text transcripts, etc. in other tables and refer back via foreign key. diff --git a/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql b/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql index 01478a453f..edb2534d2b 100644 --- a/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql +++ b/_sql/migrations_squashed/model_weights/model_weights_extension_tts_details.sql @@ -28,7 +28,7 @@ CREATE TABLE model_weights_extension_tts_details ( -- live in the `storyteller-ml` repository. -- Older models and brand new models will have this field set to null, and we'll infer which value to -- lazily backfill based on the creation date and language setting. - -- tts_text_pipeline_type VARCHAR(64) DEFAULT NULL, + text_pipeline_type VARCHAR(64) DEFAULT NULL, -- If true, multiply the mel outputs before being vocoded by a globally default constant. use_default_mel_multiply_factor BOOLEAN NOT NULL DEFAULT FALSE, diff --git a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs index 9d4ecee304..0c95cf153a 100644 --- a/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs +++ b/crates/_types/buckets/src/private/zs_voices/bucket_directory.rs @@ -26,12 +26,14 @@ impl ModelCategory { #[derive(Debug, Clone, Copy, Eq, PartialEq)] pub enum ModelType { VallEx, + StyleTTS2, } impl ModelType { pub fn to_str(&self) -> &'static str { match self { ModelType::VallEx => "vallex", + ModelType::StyleTTS2 => "styletts2", } } } diff --git a/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs b/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs index e0e959ae0c..e1d7babedf 100644 --- a/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs +++ b/crates/_types/buckets/src/public/weight_files/bucket_file_path.rs @@ -25,6 +25,11 @@ impl PublicPath for WeightFileBucketPath {} impl WeightFileBucketPath { + // TODO(bt,2024-01-14): This is temporary standardization. Clean this up. + pub fn generate_for_tt2_model() -> Self { + Self::generate_new(Some("tt2_"), Some(".pt")) + } + // TODO(bt,2023-12-19): This is temporary standardization. Clean this up. pub fn generate_for_svc_model() -> Self { Self::generate_new(Some("svc_"), Some(".pt")) @@ -35,6 +40,11 @@ impl WeightFileBucketPath { Self::generate_new(Some("rvc_"), Some(".pt")) } + // TODO(bt,2024-01-15): This is temporary standardization. Clean this up. + pub fn tt2_model_file_from_object_hash(hash: &str) -> Self { + Self::from_object_hash(hash, Some("tt2_"), Some(".pt")) + } + // TODO(bt,2023-12-19): This is temporary standardization. Clean this up. pub fn svc_model_file_from_object_hash(hash: &str) -> Self { Self::from_object_hash(hash, Some("svc_"), Some(".pt")) diff --git a/crates/_types/composite_identifiers/Cargo.toml b/crates/_types/composite_identifiers/Cargo.toml new file mode 100644 index 0000000000..f9aa2156b2 --- /dev/null +++ b/crates/_types/composite_identifiers/Cargo.toml @@ -0,0 +1,40 @@ +[package] +name = "composite_identifiers" +edition = "2018" +version = "0.0.1" +authors = [ + "Brandon Thomas ", + "Brandon Thomas ", +] +publish = false + +[lib] +name = "composite_identifiers" +path = "src/lib.rs" + +[dependencies] +# Internal +enums = { path = "../enums" } +tokens = { path = "../tokens" } + +# Workspace hack (faster compile times) +my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" } + +# Serialization +serde = "1.0.125" +serde_derive = "1.0.125" + +sqlx = { version = "0.7.2", features = [ "mysql", "sqlite", "runtime-tokio-native-tls", "chrono" ] } +sqlx-core = { version = "0.7.2" } +sqlx-mysql= { version = "0.7.2" } + +utoipa = { version = "4", features = ["chrono","actix_extras"] } + +[dev-dependencies] + +# Enum utilities +strum = { version = "0.25", features = ["derive"] } + +# Test serialization +serde_json = "1.0.108" +toml = "0.5.8" diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs new file mode 100644 index 0000000000..0c9a261314 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/batch_generation_entity.rs @@ -0,0 +1,28 @@ +use enums::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; +use tokens::tokens::media_files::MediaFileToken; + +// TODO(bt, 2024-01-16): A ref type with <'a> lifetime of inner data instead of ownership? +pub enum BatchGenerationEntity { + MediaFile(MediaFileToken), +} + +// TODO(bt, 2024-01-16): Codegen? Make traits for these? Maybe overkill. +impl BatchGenerationEntity { + pub fn get_entity_type(&self) -> BatchGenerationEntityType { + match self { + Self::MediaFile(_) => BatchGenerationEntityType::MediaFile, + } + } + + pub fn get_entity_token_str(&self) -> &str { + match self { + Self::MediaFile(token) => token.as_str(), + } + } + + pub fn get_composite_keys(&self) -> (BatchGenerationEntityType, &str) { + match self { + Self::MediaFile(token) => (BatchGenerationEntityType::MediaFile, token.as_str()), + } + } +} diff --git a/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..0f4d8b5201 --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity; \ No newline at end of file diff --git a/crates/_types/composite_identifiers/src/by_table/mod.rs b/crates/_types/composite_identifiers/src/by_table/mod.rs new file mode 100644 index 0000000000..4bbf977cdb --- /dev/null +++ b/crates/_types/composite_identifiers/src/by_table/mod.rs @@ -0,0 +1 @@ +pub mod batch_generations; diff --git a/crates/_types/composite_identifiers/src/lib.rs b/crates/_types/composite_identifiers/src/lib.rs new file mode 100644 index 0000000000..08d90188da --- /dev/null +++ b/crates/_types/composite_identifiers/src/lib.rs @@ -0,0 +1,22 @@ +//! composite_identifiers +//! +//! The purpose of this library is to provide easy to construct, **typesafe**, composite identifiers. +//! This is mostly a combination of `tokens` and `enums`. +//! + +// Never allow these +#![forbid(private_bounds)] +#![forbid(private_interfaces)] +#![forbid(unused_must_use)] // NB: It's unsafe to not close/check some things + +// Okay to toggle +#![forbid(unreachable_patterns)] +#![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +// Always allow +#![allow(dead_code)] +#![allow(non_snake_case)] + +pub mod by_table; diff --git a/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs new file mode 100644 index 0000000000..41633225f9 --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/batch_generation_entity_type.rs @@ -0,0 +1,109 @@ +use std::collections::BTreeSet; + +use serde::Deserialize; +use serde::Serialize; +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + +/// Used in the `batch_generations` table in a `VARCHAR(32)` field named `entity_type`. +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Deserialize, Serialize, Ord, PartialOrd, ToSchema)] +pub enum BatchGenerationEntityType { + /// Media files + /// This will probably be the only type supported, but we'll leave the possibility of future types. + #[serde(rename = "media_file")] + MediaFile, +} + +// TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! +impl_enum_display_and_debug_using_to_str!(BatchGenerationEntityType); +impl_mysql_enum_coders!(BatchGenerationEntityType); +impl_mysql_from_row!(BatchGenerationEntityType); + +/// NB: Legacy API for older code. +impl BatchGenerationEntityType { + pub fn to_str(&self) -> &'static str { + match self { + Self::MediaFile => "media_file", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "media_file" => Ok(Self::MediaFile), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + // NB: BTreeSet is sorted + // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. + BTreeSet::from([ + Self::MediaFile, + ]) + } +} + +#[cfg(test)] +mod tests { + use crate::by_table::batch_generations::batch_generation_entity_type::BatchGenerationEntityType; + use crate::test_helpers::assert_serialization; + + mod serde { + use super::*; + + #[test] + fn test_serialization() { + assert_serialization(BatchGenerationEntityType::MediaFile, "media_file"); + } + } + + mod impl_methods { + use super::*; + + #[test] + fn to_str() { + assert_eq!(BatchGenerationEntityType::MediaFile.to_str(), "media_file"); + } + + #[test] + fn from_str() { + assert_eq!(BatchGenerationEntityType::from_str("media_file").unwrap(), BatchGenerationEntityType::MediaFile); + assert!(BatchGenerationEntityType::from_str("foo").is_err()); + } + } + + mod manual_variant_checks { + use super::*; + + #[test] + fn all_variants() { + let mut variants = BatchGenerationEntityType::all_variants(); + assert_eq!(variants.len(), 1); + assert_eq!(variants.pop_first(), Some(BatchGenerationEntityType::MediaFile)); + assert_eq!(variants.pop_first(), None); + } + } + + mod mechanical_checks { + use super::*; + + #[test] + fn variant_length() { + use strum::IntoEnumIterator; + assert_eq!(BatchGenerationEntityType::all_variants().len(), BatchGenerationEntityType::iter().len()); + } + + #[test] + fn round_trip() { + for variant in BatchGenerationEntityType::all_variants() { + assert_eq!(variant, BatchGenerationEntityType::from_str(variant.to_str()).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{}", variant)).unwrap()); + assert_eq!(variant, BatchGenerationEntityType::from_str(&format!("{:?}", variant)).unwrap()); + } + } + } +} diff --git a/crates/_types/enums/src/by_table/batch_generations/mod.rs b/crates/_types/enums/src/by_table/batch_generations/mod.rs new file mode 100644 index 0000000000..2babf4fabc --- /dev/null +++ b/crates/_types/enums/src/by_table/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod batch_generation_entity_type; \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs index 12405e4bf8..641dab68c2 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -29,6 +29,9 @@ pub enum InferenceCategory { #[serde(rename = "video_filter")] VideoFilter, + #[serde(rename = "image_generation")] + ImageGeneration, + #[serde(rename = "mocap")] Mocap, @@ -48,6 +51,7 @@ impl InferenceCategory { Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", Self::VideoFilter => "video_filter", + Self::ImageGeneration => "image_generation", Self::Mocap => "mocap", Self::Workflow => "workflow", } @@ -59,6 +63,7 @@ impl InferenceCategory { "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), "video_filter" => Ok(Self::VideoFilter), + "image_generation" => Ok(Self::ImageGeneration), "mocap" => Ok(Self::Mocap), _ => Err(format!("invalid value: {:?}", value)), } @@ -72,6 +77,7 @@ impl InferenceCategory { Self::TextToSpeech, Self::VoiceConversion, Self::VideoFilter, + Self::ImageGeneration, Self::Mocap, Self::Workflow, ]) @@ -89,6 +95,7 @@ mod tests { assert_serialization(InferenceCategory::TextToSpeech, "text_to_speech"); assert_serialization(InferenceCategory::VoiceConversion, "voice_conversion"); assert_serialization(InferenceCategory::VideoFilter, "video_filter"); + assert_serialization(InferenceCategory::ImageGeneration, "image_generation"); assert_serialization(InferenceCategory::Mocap, "mocap"); assert_serialization(InferenceCategory::Workflow, "workflow") } @@ -99,6 +106,7 @@ mod tests { assert_eq!(InferenceCategory::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(InferenceCategory::VideoFilter.to_str(), "video_filter"); + assert_eq!(InferenceCategory::ImageGeneration.to_str(), "image_generation"); assert_eq!(InferenceCategory::Mocap.to_str(), "mocap"); assert_eq!(InferenceCategory::Workflow.to_str(), "workflow"); } @@ -109,6 +117,7 @@ mod tests { assert_eq!(InferenceCategory::from_str("text_to_speech").unwrap(), InferenceCategory::TextToSpeech); assert_eq!(InferenceCategory::from_str("voice_conversion").unwrap(), InferenceCategory::VoiceConversion); assert_eq!(InferenceCategory::from_str("video_filter").unwrap(), InferenceCategory::VideoFilter); + assert_eq!(InferenceCategory::from_str("image_generation").unwrap(), InferenceCategory::ImageGeneration); assert_eq!(InferenceCategory::from_str("mocap").unwrap(), InferenceCategory::Mocap); assert_eq!(InferenceCategory::from_str("workflow").unwrap(), InferenceCategory::Workflow); } @@ -122,6 +131,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceCategory::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(InferenceCategory::VideoFilter)); + assert_eq!(variants.pop_first(), Some(InferenceCategory::ImageGeneration)); assert_eq!(variants.pop_first(), Some(InferenceCategory::Mocap)); assert_eq!(variants.pop_first(), Some(InferenceCategory::Workflow)); assert_eq!(variants.pop_first(), None); diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs index ca4e7bb38d..cd139eb497 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_model_type.rs @@ -33,10 +33,14 @@ pub enum InferenceModelType { VallEX, #[serde(rename = "rerender_a_video")] RerenderAVideo, + #[serde(rename = "stable_diffusion")] + StableDiffusion, #[serde(rename = "mocap_net")] MocapNet, #[serde(rename = "comfy_ui")] ComfyUi, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -54,7 +58,9 @@ impl InferenceModelType { Self::Vits => "vits", Self::VallEX => "vall_e_x", Self::RerenderAVideo => "rerender_a_video", + Self::StableDiffusion => "stable_diffusion", Self::MocapNet => "mocap_net", + Self::StyleTTS2 => "styletts2", Self::ComfyUi => "comfy_ui", } } @@ -68,7 +74,9 @@ impl InferenceModelType { "vits" => Ok(Self::Vits), "vall_e_x" => Ok(Self::VallEX), "rerender_a_video" => Ok(Self::RerenderAVideo), + "stable_diffusion" => Ok(Self::StableDiffusion), "mocap_net" => Ok(Self::MocapNet), + "styletts2" => Ok(Self::StyleTTS2), "comfy_ui" => Ok(Self::ComfyUi), _ => Err(format!("invalid value: {:?}", value)), } @@ -85,7 +93,9 @@ impl InferenceModelType { InferenceModelType::Vits, InferenceModelType::VallEX, InferenceModelType::RerenderAVideo, + InferenceModelType::StableDiffusion, InferenceModelType::MocapNet, + InferenceModelType::StyleTTS2, ]) } } @@ -104,8 +114,10 @@ mod tests { assert_serialization(InferenceModelType::Vits, "vits"); assert_serialization(InferenceModelType::VallEX, "vall_e_x"); assert_serialization(InferenceModelType::RerenderAVideo, "rerender_a_video"); + assert_serialization(InferenceModelType::StableDiffusion, "stable_diffusion"); assert_serialization(InferenceModelType::MocapNet, "mocap_net"); assert_serialization(InferenceModelType::ComfyUi, "comfy_ui"); + assert_serialization(InferenceModelType::StyleTTS2, "styletts2"); } #[test] @@ -117,7 +129,9 @@ mod tests { assert_eq!(InferenceModelType::Vits.to_str(), "vits"); assert_eq!(InferenceModelType::VallEX.to_str(), "vall_e_x"); assert_eq!(InferenceModelType::RerenderAVideo.to_str(), "rerender_a_video"); + assert_eq!(InferenceModelType::StableDiffusion.to_str(), "stable_diffusion"); assert_eq!(InferenceModelType::MocapNet.to_str(), "mocap_net"); + assert_eq!(InferenceModelType::StyleTTS2.to_str(), "styletts2"); assert_eq!(InferenceModelType::ComfyUi.to_str(), "comfy_ui"); } @@ -130,7 +144,9 @@ mod tests { assert_eq!(InferenceModelType::from_str("vits").unwrap(), InferenceModelType::Vits); assert_eq!(InferenceModelType::from_str("vall_e_x").unwrap(), InferenceModelType::VallEX); assert_eq!(InferenceModelType::from_str("rerender_a_video").unwrap(), InferenceModelType::RerenderAVideo); + assert_eq!(InferenceModelType::from_str("stable_diffusion").unwrap(), InferenceModelType::StableDiffusion); assert_eq!(InferenceModelType::from_str("mocap_net").unwrap(), InferenceModelType::MocapNet); + assert_eq!(InferenceModelType::from_str("styletts2").unwrap(), InferenceModelType::StyleTTS2); assert_eq!(InferenceModelType::from_str("comfy_ui").unwrap(), InferenceModelType::ComfyUi); } @@ -138,7 +154,7 @@ mod tests { fn all_variants() { // Static check let mut variants = InferenceModelType::all_variants(); - assert_eq!(variants.len(), 9); + assert_eq!(variants.len(), 10); assert_eq!(variants.pop_first(), Some(InferenceModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(InferenceModelType::SoVitsSvc)); @@ -146,7 +162,9 @@ mod tests { assert_eq!(variants.pop_first(), Some(InferenceModelType::Vits)); assert_eq!(variants.pop_first(), Some(InferenceModelType::VallEX)); assert_eq!(variants.pop_first(), Some(InferenceModelType::RerenderAVideo)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::StableDiffusion)); assert_eq!(variants.pop_first(), Some(InferenceModelType::MocapNet)); + assert_eq!(variants.pop_first(), Some(InferenceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), Some(InferenceModelType::ComfyUi)); assert_eq!(variants.pop_first(), None); diff --git a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs index 4a7d879c21..e411079f40 100644 --- a/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs +++ b/crates/_types/enums/src/by_table/generic_inference_jobs/inference_result_type.rs @@ -29,6 +29,9 @@ pub enum InferenceResultType { #[serde(rename = "zs_voice_embedding")] ZeroShotVoiceEmbedding, + + #[serde(rename = "upload_model")] + UploadModel } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -43,6 +46,7 @@ impl InferenceResultType { Self::TextToSpeech => "text_to_speech", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoiceEmbedding => "zs_voice_embedding", + Self::UploadModel => "upload_model", } } @@ -52,6 +56,7 @@ impl InferenceResultType { "text_to_speech" => Ok(Self::TextToSpeech), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice_embedding" => Ok(Self::ZeroShotVoiceEmbedding), + "upload_model" => Ok(Self::UploadModel), _ => Err(format!("invalid value: {:?}", value)), } } @@ -64,6 +69,7 @@ impl InferenceResultType { InferenceResultType::TextToSpeech, InferenceResultType::VoiceConversion, InferenceResultType::ZeroShotVoiceEmbedding, + InferenceResultType::UploadModel, ]) } } @@ -79,6 +85,7 @@ mod tests { assert_serialization(InferenceResultType::TextToSpeech, "text_to_speech"); assert_serialization(InferenceResultType::VoiceConversion, "voice_conversion"); assert_serialization(InferenceResultType::ZeroShotVoiceEmbedding, "zs_voice_embedding"); + assert_serialization(InferenceResultType::UploadModel, "upload_model"); } #[test] @@ -87,6 +94,7 @@ mod tests { assert_eq!(InferenceResultType::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(InferenceResultType::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(InferenceResultType::ZeroShotVoiceEmbedding.to_str(), "zs_voice_embedding"); + assert_eq!(InferenceResultType::UploadModel.to_str(), "upload_model"); } #[test] @@ -95,17 +103,19 @@ mod tests { assert_eq!(InferenceResultType::from_str("text_to_speech").unwrap(), InferenceResultType::TextToSpeech); assert_eq!(InferenceResultType::from_str("voice_conversion").unwrap(), InferenceResultType::VoiceConversion); assert_eq!(InferenceResultType::from_str("zs_voice_embedding").unwrap(), InferenceResultType::ZeroShotVoiceEmbedding); + assert_eq!(InferenceResultType::from_str("upload_model").unwrap(), InferenceResultType::UploadModel); } #[test] fn all_variants() { // Static check let mut variants = InferenceResultType::all_variants(); - assert_eq!(variants.len(), 4); + assert_eq!(variants.len(), 5); assert_eq!(variants.pop_first(), Some(InferenceResultType::MediaFile)); assert_eq!(variants.pop_first(), Some(InferenceResultType::TextToSpeech)); assert_eq!(variants.pop_first(), Some(InferenceResultType::VoiceConversion)); assert_eq!(variants.pop_first(), Some(InferenceResultType::ZeroShotVoiceEmbedding)); + assert_eq!(variants.pop_first(), Some(InferenceResultType::UploadModel)); assert_eq!(variants.pop_first(), None); // Generated check diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs index 2a2ef67aa4..f1fa83a365 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs @@ -41,6 +41,9 @@ pub enum MediaFileOriginModelType { #[serde(rename = "comfy_ui")] ComfyUi, + + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -60,6 +63,7 @@ impl MediaFileOriginModelType { Self::Rerender => "rerender", Self::MocapNet => "mocap_net", Self::ComfyUi => "comfy_ui", + Self::StyleTTS2 => "styletts2", } } @@ -73,6 +77,7 @@ impl MediaFileOriginModelType { "rerender" => Ok(Self::Rerender), "mocap_net" => Ok(Self::MocapNet), "comfy_ui" => Ok(Self::ComfyUi), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -89,6 +94,7 @@ impl MediaFileOriginModelType { Self::Rerender, Self::MocapNet, Self::ComfyUi, + Self::StyleTTS2, ]) } } @@ -116,6 +122,7 @@ mod tests { assert_eq!(MediaFileOriginModelType::Rerender.to_str(), "rerender"); assert_eq!(MediaFileOriginModelType::MocapNet.to_str(), "mocap_net"); assert_eq!(MediaFileOriginModelType::ComfyUi.to_str(), "comfy_ui"); + assert_eq!(MediaFileOriginModelType::StyleTTS2.to_str(), "styletts2"); } #[test] @@ -128,13 +135,14 @@ mod tests { assert_eq!(MediaFileOriginModelType::from_str("rerender").unwrap(), MediaFileOriginModelType::Rerender); assert_eq!(MediaFileOriginModelType::from_str("mocap_net").unwrap(), MediaFileOriginModelType::MocapNet); assert_eq!(MediaFileOriginModelType::from_str("comfy_ui").unwrap(), MediaFileOriginModelType::ComfyUi); + assert_eq!(MediaFileOriginModelType::from_str("styletts2").unwrap(), MediaFileOriginModelType::StyleTTS2); assert!(MediaFileOriginModelType::from_str("foo").is_err()); } #[test] fn all_variants() { let mut variants = MediaFileOriginModelType::all_variants(); - assert_eq!(variants.len(), 8); + assert_eq!(variants.len(), 10); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::RvcV2)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SadTalker)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::SoVitsSvc)); @@ -143,6 +151,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::Rerender)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::MocapNet)); assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::ComfyUi)); + assert_eq!(variants.pop_first(), Some(MediaFileOriginModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs index 6508c6f9df..0efb1294c4 100644 --- a/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs +++ b/crates/_types/enums/src/by_table/media_files/media_file_origin_product_category.rs @@ -17,6 +17,11 @@ use utoipa::ToSchema; #[cfg_attr(test, derive(EnumIter, EnumCount))] #[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Deserialize, Serialize, ToSchema)] pub enum MediaFileOriginProductCategory { + /// Unknown which product is attached to the file (generated the file, the file was + /// uploaded on behalf of, etc.) + #[serde(rename = "unknown")] + Unknown, + /// Media files created by (or uploaded for) the Face Animator product. /// The underlying model could be SadTalker, Wav2Lip, or some future model #[serde(rename = "face_animator")] @@ -26,12 +31,6 @@ pub enum MediaFileOriginProductCategory { #[serde(rename = "tts")] TextToSpeech, - // TODO: This should be a temporary category until we migrate the DB to remove this default value - /// Unknown which product is attached to the file (generated the file, the file was - /// uploaded on behalf of, etc.) - #[serde(rename = "unknown")] - Unknown, - /// Voice conversion (either RVC or SVC) #[serde(rename = "voice_conversion")] VoiceConversion, @@ -62,9 +61,9 @@ impl_mysql_from_row!(MediaFileOriginProductCategory); impl MediaFileOriginProductCategory { pub fn to_str(&self) -> &'static str { match self { + Self::Unknown => "unknown", Self::FaceAnimator => "face_animator", Self::TextToSpeech => "tts", - Self::Unknown => "unknown", Self::VoiceConversion => "voice_conversion", Self::ZeroShotVoice => "zs_voice", Self::VideoFilter => "video_filter", @@ -75,9 +74,9 @@ impl MediaFileOriginProductCategory { pub fn from_str(value: &str) -> Result { match value { + "unknown" => Ok(Self::Unknown), "face_animator" => Ok(Self::FaceAnimator), "tts" => Ok(Self::TextToSpeech), - "unknown" => Ok(Self::Unknown), "voice_conversion" => Ok(Self::VoiceConversion), "zs_voice" => Ok(Self::ZeroShotVoice), "video_filter" => Ok(Self::VideoFilter), @@ -91,10 +90,10 @@ impl MediaFileOriginProductCategory { // NB: BTreeSet is sorted // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ + Self::Unknown, Self::FaceAnimator, Self::TextToSpeech, Self::VideoFilter, - Self::Unknown, Self::VoiceConversion, Self::ZeroShotVoice, Self::Mocap, @@ -113,9 +112,9 @@ mod tests { #[test] fn test_serialization() { + assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::FaceAnimator, "face_animator"); assert_serialization(MediaFileOriginProductCategory::TextToSpeech, "tts"); - assert_serialization(MediaFileOriginProductCategory::Unknown, "unknown"); assert_serialization(MediaFileOriginProductCategory::VoiceConversion, "voice_conversion"); assert_serialization(MediaFileOriginProductCategory::ZeroShotVoice, "zs_voice"); assert_serialization(MediaFileOriginProductCategory::VideoFilter, "video_filter"); @@ -125,9 +124,9 @@ mod tests { #[test] fn to_str() { + assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::FaceAnimator.to_str(), "face_animator"); assert_eq!(MediaFileOriginProductCategory::TextToSpeech.to_str(), "tts"); - assert_eq!(MediaFileOriginProductCategory::Unknown.to_str(), "unknown"); assert_eq!(MediaFileOriginProductCategory::VoiceConversion.to_str(), "voice_conversion"); assert_eq!(MediaFileOriginProductCategory::ZeroShotVoice.to_str(), "zs_voice"); assert_eq!(MediaFileOriginProductCategory::VideoFilter.to_str(), "video_filter"); @@ -137,9 +136,9 @@ mod tests { #[test] fn from_str() { + assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("face_animator").unwrap(), MediaFileOriginProductCategory::FaceAnimator); assert_eq!(MediaFileOriginProductCategory::from_str("tts").unwrap(), MediaFileOriginProductCategory::TextToSpeech); - assert_eq!(MediaFileOriginProductCategory::from_str("unknown").unwrap(), MediaFileOriginProductCategory::Unknown); assert_eq!(MediaFileOriginProductCategory::from_str("voice_conversion").unwrap(), MediaFileOriginProductCategory::VoiceConversion); assert_eq!(MediaFileOriginProductCategory::from_str("zs_voice").unwrap(), MediaFileOriginProductCategory::ZeroShotVoice); assert_eq!(MediaFileOriginProductCategory::from_str("video_filter").unwrap(), MediaFileOriginProductCategory::VideoFilter); @@ -151,9 +150,9 @@ mod tests { fn all_variants() { let mut variants = MediaFileOriginProductCategory::all_variants(); assert_eq!(variants.len(), 8); + assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::FaceAnimator)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::TextToSpeech)); - assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::Unknown)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VoiceConversion)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::ZeroShotVoice)); assert_eq!(variants.pop_first(), Some(MediaFileOriginProductCategory::VideoFilter)); diff --git a/crates/_types/enums/src/by_table/mod.rs b/crates/_types/enums/src/by_table/mod.rs index b3134ad59d..0c0169b2ac 100644 --- a/crates/_types/enums/src/by_table/mod.rs +++ b/crates/_types/enums/src/by_table/mod.rs @@ -1,5 +1,6 @@ // ===== MySql ===== pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_jobs; pub mod entity_stats; @@ -17,6 +18,7 @@ pub mod user_ratings; pub mod voice_conversion_models; pub mod voice_conversion_results; pub mod zs_voices; + // ===== Sqlite ===== pub mod tts_render_tasks; pub mod web_rendition_targets; diff --git a/crates/_types/enums/src/by_table/model_weights/weights_category.rs b/crates/_types/enums/src/by_table/model_weights/weights_category.rs index f2b10e67d2..d281b9dfb6 100644 --- a/crates/_types/enums/src/by_table/model_weights/weights_category.rs +++ b/crates/_types/enums/src/by_table/model_weights/weights_category.rs @@ -17,6 +17,8 @@ pub enum WeightsCategory { Vocoder, #[serde(rename = "voice_conversion")] VoiceConversion, + #[serde(rename = "workflow_config")] + WorkflowConfig, } impl WeightsCategory { @@ -26,6 +28,7 @@ impl WeightsCategory { Self::TextToSpeech => "text_to_speech", Self::Vocoder => "vocoder", Self::VoiceConversion => "voice_conversion", + Self::WorkflowConfig => "workflow_config", } } @@ -35,6 +38,7 @@ impl WeightsCategory { "text_to_speech" => Ok(Self::TextToSpeech), "vocoder" => Ok(Self::Vocoder), "voice_conversion" => Ok(Self::VoiceConversion), + "workflow_config" => Ok(Self::WorkflowConfig), _ => Err(format!("invalid value: {:?}", value)), } } @@ -45,6 +49,7 @@ impl WeightsCategory { Self::TextToSpeech, Self::Vocoder, Self::VoiceConversion, + Self::WorkflowConfig, ]) } } @@ -62,6 +67,7 @@ mod tests { assert_eq!(WeightsCategory::TextToSpeech.to_str(), "text_to_speech"); assert_eq!(WeightsCategory::Vocoder.to_str(), "vocoder"); assert_eq!(WeightsCategory::VoiceConversion.to_str(), "voice_conversion"); + assert_eq!(WeightsCategory::WorkflowConfig.to_str(), "workflow_config"); } #[test] @@ -70,16 +76,18 @@ mod tests { assert_eq!(WeightsCategory::from_str("text_to_speech").unwrap(), WeightsCategory::TextToSpeech); assert_eq!(WeightsCategory::from_str("vocoder").unwrap(), WeightsCategory::Vocoder); assert_eq!(WeightsCategory::from_str("voice_conversion").unwrap(), WeightsCategory::VoiceConversion); + assert_eq!(WeightsCategory::from_str("workflow_config").unwrap(), WeightsCategory::WorkflowConfig); assert!(WeightsCategory::from_str("invalid").is_err()); } #[test] fn test_all_variants() { let variants = WeightsCategory::all_variants(); - assert_eq!(variants.len(), 4); + assert_eq!(variants.len(), 5); assert!(variants.contains(&WeightsCategory::ImageGeneration)); assert!(variants.contains(&WeightsCategory::TextToSpeech)); assert!(variants.contains(&WeightsCategory::Vocoder)); assert!(variants.contains(&WeightsCategory::VoiceConversion)); + assert!(variants.contains(&WeightsCategory::WorkflowConfig)); } } \ No newline at end of file diff --git a/crates/_types/enums/src/by_table/model_weights/weights_types.rs b/crates/_types/enums/src/by_table/model_weights/weights_types.rs index 12f00955df..c0ecb6fce3 100644 --- a/crates/_types/enums/src/by_table/model_weights/weights_types.rs +++ b/crates/_types/enums/src/by_table/model_weights/weights_types.rs @@ -25,6 +25,8 @@ pub enum WeightsType { LoRA, #[serde(rename = "vall_e")] VallE, + #[serde(rename = "comfy_ui")] + ComfyUi, } @@ -39,6 +41,7 @@ impl WeightsType { Self::Tacotron2 => "tt2", Self::LoRA => "loRA", Self::VallE => "vall_e", + Self::ComfyUi => "comfy_ui", } } @@ -52,6 +55,7 @@ impl WeightsType { "tt2" => Ok(Self::Tacotron2), "loRA" => Ok(Self::LoRA), "vall_e" => Ok(Self::VallE), + "comfy_ui" => Ok(Self::ComfyUi), _ => Err(format!("invalid value: {:?}", value)), } } @@ -65,7 +69,8 @@ impl WeightsType { Self::SoVitsSvc, Self::Tacotron2, Self::LoRA, - Self::VallE + Self::VallE, + Self::ComfyUi, ]) } } @@ -88,6 +93,7 @@ mod tests { assert_eq!(WeightsType::Tacotron2.to_str(), "tt2"); assert_eq!(WeightsType::LoRA.to_str(), "loRA"); assert_eq!(WeightsType::VallE.to_str(), "vall_e"); + assert_eq!(WeightsType::ComfyUi.to_str(), "comfy_ui"); } #[test] @@ -100,13 +106,14 @@ mod tests { assert_eq!(WeightsType::from_str("tt2").unwrap(), WeightsType::Tacotron2); assert_eq!(WeightsType::from_str("loRA").unwrap(), WeightsType::LoRA); assert_eq!(WeightsType::from_str("vall_e").unwrap(), WeightsType::VallE); + assert_eq!(WeightsType::from_str("comfy_ui").unwrap(), WeightsType::ComfyUi); assert!(WeightsType::from_str("invalid").is_err()); } #[test] fn test_all_variants() { let variants = WeightsType::all_variants(); - assert_eq!(variants.len(), 8); + assert_eq!(variants.len(), 9); assert!(variants.contains(&WeightsType::HifiganTacotron2)); assert!(variants.contains(&WeightsType::RvcV2)); assert!(variants.contains(&WeightsType::StableDiffusion15)); @@ -115,5 +122,6 @@ mod tests { assert!(variants.contains(&WeightsType::Tacotron2)); assert!(variants.contains(&WeightsType::LoRA)); assert!(variants.contains(&WeightsType::VallE)); + assert!(variants.contains(&WeightsType::ComfyUi)); } } diff --git a/crates/_types/enums/src/by_table/zs_voices/model_type.rs b/crates/_types/enums/src/by_table/zs_voices/model_type.rs index d8f022035d..a5e4dca813 100644 --- a/crates/_types/enums/src/by_table/zs_voices/model_type.rs +++ b/crates/_types/enums/src/by_table/zs_voices/model_type.rs @@ -16,6 +16,8 @@ pub enum ZsVoiceModelType { /// TTS-type zero shot models #[serde(rename = "vall-e-x")] VallEX, + #[serde(rename = "styletts2")] + StyleTTS2, } // TODO(bt, 2023-01-17): This desperately needs MySQL integration tests! @@ -27,12 +29,14 @@ impl ZsVoiceModelType { pub fn to_str(&self) -> &'static str { match self { Self::VallEX=> "vall-e-x", + Self::StyleTTS2 => "styletts2", } } pub fn from_str(value: &str) -> Result { match value { "vall-e-x" => Ok(Self::VallEX), + "styletts2" => Ok(Self::StyleTTS2), _ => Err(format!("invalid value: {:?}", value)), } } @@ -42,6 +46,7 @@ impl ZsVoiceModelType { // NB: BTreeSet::from() isn't const, but not worth using LazyStatic, etc. BTreeSet::from([ Self::VallEX, + Self::StyleTTS2, ]) } } @@ -66,11 +71,13 @@ mod tests { #[test] fn to_str() { assert_eq!(ZsVoiceModelType::VallEX.to_str(), "vall-e-x"); + assert_eq!(ZsVoiceModelType::StyleTTS2.to_str(), "styletts2"); } #[test] fn from_str() { assert_eq!(ZsVoiceModelType::from_str("vall-e-x").unwrap(), ZsVoiceModelType::VallEX); + assert_eq!(ZsVoiceModelType::from_str("styletts2").unwrap(), ZsVoiceModelType::StyleTTS2); assert!(ZsVoiceModelType::from_str("foo").is_err()); } } @@ -81,8 +88,9 @@ mod tests { #[test] fn all_variants() { let mut variants = ZsVoiceModelType::all_variants(); - assert_eq!(variants.len(), 1); + assert_eq!(variants.len(), 2); assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::VallEX)); + assert_eq!(variants.pop_first(), Some(ZsVoiceModelType::StyleTTS2)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/_types/tokens/src/prefixes.rs b/crates/_types/tokens/src/prefixes.rs index 382492bae4..e4263e21bb 100644 --- a/crates/_types/tokens/src/prefixes.rs +++ b/crates/_types/tokens/src/prefixes.rs @@ -24,6 +24,7 @@ use strum::EnumIter; pub(crate) enum TokenPrefix { AnonymousVisitorTracking, // AVTs are not stored as primary keys in any table, but an index in many tables. AuditLog, + BatchGeneration, Comment, DownloadJob, EmailSenderJob, @@ -90,6 +91,7 @@ impl PrefixGenerator for TokenPrefix { match self { Self::AnonymousVisitorTracking => "avt_", Self::AuditLog => "audit_", + Self::BatchGeneration => "batch_g_", Self::Comment => "comment_", Self::DownloadJob => "jdown_", // NB: Previously "JGUP:" Self::EmailSenderJob => "email_job_", @@ -297,7 +299,13 @@ mod tests { #[test] fn test_all_prefixes_end_with_separator_length_one() { for prefix in get_all_prefixes().iter().map(|s| *s) { - if prefix == "news_story_" || prefix == "tts_task_" || prefix == "INT_API:" || prefix == "pw_reset_" || prefix == "email_job_" { + if prefix == "news_story_" + || prefix == "tts_task_" + || prefix == "INT_API:" + || prefix == "pw_reset_" + || prefix == "email_job_" + || prefix == "batch_g_" + { // TODO/FIXME: I'm too tired at 5AM to replacen from the left. Make this test valid. // These tokens are from the AIChatBot sidecar, so asserting their validity is less important. continue; diff --git a/crates/_types/tokens/src/tokens/batch_generations.rs b/crates/_types/tokens/src/tokens/batch_generations.rs new file mode 100644 index 0000000000..052681f21a --- /dev/null +++ b/crates/_types/tokens/src/tokens/batch_generations.rs @@ -0,0 +1,15 @@ +use std::fmt::Debug; + +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use crate::prefixes::TokenPrefix; + +/// The primary key for Audit Logs. +#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, sqlx::Type, Debug, Serialize, Deserialize, ToSchema)] +#[sqlx(transparent)] +pub struct BatchGenerationToken(pub String); + +impl_string_token!(BatchGenerationToken); +impl_crockford_generator!(BatchGenerationToken, 32usize, TokenPrefix::BatchGeneration, CrockfordLower); diff --git a/crates/_types/tokens/src/tokens/mod.rs b/crates/_types/tokens/src/tokens/mod.rs index b7c4eed592..447cbedf43 100644 --- a/crates/_types/tokens/src/tokens/mod.rs +++ b/crates/_types/tokens/src/tokens/mod.rs @@ -2,6 +2,7 @@ pub mod anonymous_visitor_tracking; pub mod api_tokens_external; pub mod api_tokens_internal; pub mod audit_logs; +pub mod batch_generations; pub mod comments; pub mod email_sender_job_token; pub mod firehose_entries; @@ -20,6 +21,7 @@ pub mod tts_results; pub mod twitch_event_rule; pub mod twitch_oauth_tokens_grouping; pub mod twitch_oauth_tokens_internal; +pub mod user_bookmarks; pub mod user_sessions; pub mod user_subscriptions; pub mod users; @@ -34,4 +36,3 @@ pub mod w2l_templates; pub mod zs_voice_dataset_samples; pub mod zs_voice_datasets; pub mod zs_voices; -pub mod user_bookmarks; diff --git a/crates/cli/dev-database-seed/src/main.rs b/crates/cli/dev-database-seed/src/main.rs index 12d16ddd0b..e123dc3c24 100644 --- a/crates/cli/dev-database-seed/src/main.rs +++ b/crates/cli/dev-database-seed/src/main.rs @@ -7,13 +7,11 @@ use errors::AnyhowResult; use crate::bucket_clients::get_bucket_clients; use crate::cli_args::parse_cli_args; -use crate::seeding::model_weights::seed_weights; // V1 -use crate::seeding::seed_model_weights::seed_weights_files; // V2 v3 will come later on when everything is more conslidated and clear. +use crate::seeding::model_weights::seed_weights; use crate::seeding::seed_media_files::seed_media_files; use crate::seeding::tts_tacotron2::seed_tts_tacotron2; use crate::seeding::users::seed_user_accounts; use crate::seeding::voice_conversion::seed_voice_conversion; -use crate::seeding::zero_shot_tts::seed_zero_shot_tts; pub mod bucket_clients; @@ -61,9 +59,10 @@ pub async fn main() -> AnyhowResult<()> { seed_tts_tacotron2(&pool, maybe_bucket_clients.as_ref()).await?; // should seed the weights with a few files for hanashi - seed_weights_files(&pool, maybe_bucket_clients.as_ref()).await?; - - info!("Done!"); + //seed_weights_files(&pool, maybe_bucket_clients.as_ref()).await?; + //println!("TESTING DOWLOAD"); + //test_seed_weights_files().await?; + //info!("Done!"); Ok(()) } diff --git a/crates/cli/dev-database-seed/src/seeding/model_weights.rs b/crates/cli/dev-database-seed/src/seeding/model_weights.rs index 5234bd700a..497d996af0 100644 --- a/crates/cli/dev-database-seed/src/seeding/model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/model_weights.rs @@ -1,8 +1,12 @@ +use std::path::PathBuf; use std::vec; use log::info; use sqlx::{MySql, Pool}; +use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use cloud_storage::remote_file_manager::weights_descriptor::{WeightsLoRADescriptor, WeightsSD15Descriptor, WeightsWorkflowDescriptor}; use enums::by_table::model_weights::{ weights_category::WeightsCategory, weights_types::WeightsType, @@ -11,11 +15,48 @@ use enums::common::visibility::Visibility; use errors::{anyhow, AnyhowResult}; use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; use mysql_queries::queries::users::user::get_user_token_by_username::get_user_token_by_username; +use storyteller_root::get_seed_tool_data_root; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; use crate::seeding::users::HANASHI_USERNAME; +pub async fn test_seed_weights_files() -> AnyhowResult<()> { + + let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data"); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await; + let remote_cloud_file_client = match remote_cloud_file_client { + Ok(res) => { + res + } + Err(_) => { + return Err(anyhow!("failed to get remote cloud file client")); + } + }; + + let mut path_dl1 = seed_path.clone(); + path_dl1.push("downloads/loRA"); + let mut path_dl2 = seed_path.clone(); + path_dl2.push("downloads/checkpoint"); + + let bucket_details1 = RemoteCloudBucketDetails { + object_hash: String::from("apa0ej6es8d3ss2gwtf1cghge35qn9tn"), + prefix: String::from("sd15"), + suffix: String::from("safetensors"), + }; + + let bucket_details2 = RemoteCloudBucketDetails { + object_hash: String::from("27kz11et18fargyyxbj66ntfn621k9d3"), + prefix: String::from("loRA"), + suffix: String::from("safetensors"), + }; + + remote_cloud_file_client.download_file(bucket_details1, String::from("./checkpoint")).await?; + remote_cloud_file_client.download_file(bucket_details2, String::from("./loRA")).await?; + + Ok(()) +} + pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()> { let sd1_5_markdown_description = r#" # Dragonfruit AI Models Update and Workflow @@ -330,6 +371,351 @@ pub async fn seed_weights_for_paging(mysql_pool: &Pool, user_token: UserT } Ok(()) } + +pub async fn seed_weights_for_testing_inference(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()>{ + + + let miyhoyo_description = r#"一个面向米哈游角色的模型合集~A collection for MIHOYO Characters~ +对应人物tag记录在版本信息里，可在右侧“About this version”选项中查看。The corresponding character prompts are recorded in the version information and can be viewed in the About this version option on the right. +打开目录寻找最爱~（附带链接，点击直达）~Open the catalog, find your favorites~ (with links, click to go directly): +崩坏三 Honkai Impact 3rd： + +梅比乌斯- Mobius + +琪亚娜-Kiana：琪亚娜 合集-Collections for Kiana (Honkai Impact 3rd) + +布洛妮娅-bronya：布洛妮娅 合集-Collections for Bronya (Honkai Impact 3rd | Honkai Star Rail) + +芽衣：雷电芽衣 合集-Collections for Raiden Mei (Honkai Impact 3rd) + +西琳-Sirin + +魔法少女-Miracle ☆ Magical Girl + +赤鸢-Jingwei + +丹朱苍玄-Nuwa/Fuxi + +李素裳-Li Sushang + +玉骑士·月痕/锦岁玉团-Jade Knight/Auspicious Dazzle + +卡莲-kallen + +第六夜想曲-Sixth Serenade + +符华-Fu hua + +崩落 + +云墨丹心-Azure Empyrea + +八重樱-Yae Sakura + +逆神巫女-Gyakushinn Miko + +御神装·勿忘-Goushinnso Memento + +信花舞伎-Blooming Maiko + +希尔-Seele + +生死律者-Herrscher of Rebirth + +栖于永夜-Eternal Night's Embrace + +德莉莎-Theresa：德莉莎 合集-Collections for Theresa (Honkai Impact 3rd) + +月下誓约-Lunar Vow: Crimson Love + +月下初拥-Luna Kindred + +花裳月纱-Rosy Bridesmaid + +格蕾修-Griseo + +繁星·绘世之卷-Starry Impression + +天行·绘星之卷-Cosmic Expression + +画中的童话-Maroon Riding Hood + +原神 Genshin： + +凝光-Ningguang + +申鹤-Shenhe + +妮露 + +诺艾尔 + +神里绫华（包含礼装） + +甘雨(包含喜茶联动礼装)-Ganyu + +胡桃(包含森罗联动礼装)-Hutao + +琳妮特-Lynette + +绮良良-Kirara + +坎蒂丝-Candace + +芙宁娜-Furina + +芙宁娜/芙卡洛斯-Furina/Focalors + +克洛琳德-Clorinde + +娜维娅-Navia + +丝柯克-Skirk + +夏洛蒂-Charlotte + +归终-Guizhong + +闲云-Xianyun + +千织-Chiori + +夏沃蕾-Chevreuse + +崩坏：星穹铁道 Honkai Star Rand： + +驭空-Yukong + +银狼-Silver Wolf + +符玄-Fu Xuan + +李素裳-Sushang + +三月七(礼装)-March 7th (New Outfit) + +停云-Tingyun + +青雀-Qingque + +布洛妮娅·兰德-Bronya Rand + +镜流-Jingliu + +卡芙卡-Kafka + +托帕-Topaz + +桂乃芬-Guinaifen + +玲可-Lynx + +佩拉-Peia + +寒鸦-Hanya + +藿藿-Huohuo + +阮·梅-Ruan Mei + +雪衣-Xueyi + +希露瓦-Serval + +绝区零 Zenless-zone-zero： + +星见雅-Miyabi + +猫宫又奈-Nekomata + +艾莲·乔-Ellen + +妮可-Nicole + +苍角-Soukaku + +格莉丝-Grace + +珂蕾妲-koleda + +可琳-Corin + +其他 Others： + +鹿鸣-Luming + """#; + + let sd1_5_markdown_description = r#" +00025-2161235528-1girl,(orange_hair_1.1),(zentangle, mandala, tangle, entangle_0.6),(fractal art),the most beautiful form of chaos,brutalist desi.png + + +00010-1626070972-elf portrait,enchanting beauty,fantasy,ethereal glow,pointed ears,delicate facial features,long elegant hair,nature-themed attir.png + + +00001-301370162-girl in a snowy landscape,winter coat,fluffy hood,white fur trim,breath visible in cold air,soft falling snowflakes,glistening s.png + + +00002-3996269517-venomous sorcerer,ominous dark robes,glowing green eyes,holding a staff with a serpent design,mystical green flames,poisonous mi.png + + +00003-3817469732-Thunder Deity,immense power,hovering above the sea,stormy ocean waves below,crackling lightning in fingertips,dramatic stormy cl.png + +if you want to remake my front image, check this out: Lightflow | SD webUI Extension - workflow restore & save in one click - majicMIXfantasy example | Stable Diffusion Other | Civitai + +v3 is here. Shouldn't be difficult to use. Simple and complex prompts can be used, unleash your imagination. + +v3来了，简单复杂的提示词都可以用。 + +I developed the v3 version towards 2D in order to make more room for 2.5D models like majicmix lux. + +v3版本我往2D发展是为了给majicmix lux等2.5D模型让出更多空间。 + +merged from融了: +Jack Of All by GeneratedJudge + +AWPainting by DynamicWang + +majicMIX fantasy v2 = +Noosphere by skumerz + dalcefoPainting + 饭特稀V08 by zhazhahui345 + GhostMix +"#; + + + let sd_vae_description = r"This is an earlier version of a stable VAE. Compared to other VAEs, it has a higher level of stability. I am not the creator, but I could not find this VAE on the website, so I am sharing it here. +这是较早时候的稳定VAE，与其他VAE相比具有较高的稳定性，不容易坏图。我不是作者，但是站里没找着所以搬运。#"; + + + info!("Seeding weights for inference..."); + ModelWeightToken::reset_rng_for_testing_and_dev_seeding_never_use_in_production_seriously(54321); + let model_weight_token1 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token2 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let model_weight_token3 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + + let sd_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); + let lora_descriptor = Box::new(WeightsLoRADescriptor{}); + let sd_vae_15_weights_descriptor = Box::new(WeightsSD15Descriptor {}); + + let mut path_object_SD = get_seed_tool_data_root(); + path_object_SD.push("models/imagegen/sd15/majicmixFantasy_v30Vae.safetensors"); + + let mut path_object_loRA = get_seed_tool_data_root(); + path_object_loRA.push("models/imagegen/loRA/xiawolei-v100-000019.safetensors"); + + let mut path_to_VAE = get_seed_tool_data_root(); + path_to_VAE.push("models/imagegen/vae/zVae_v20.safetensors"); + + let metadata1 = remote_cloud_file_client.upload_file(sd_15_weights_descriptor,path_object_SD.as_path().to_str().unwrap()).await?; + let metadata2 = remote_cloud_file_client.upload_file(lora_descriptor,path_object_loRA.as_path().to_str().unwrap()).await?; + let metadata3 = remote_cloud_file_client.upload_file(sd_vae_15_weights_descriptor,path_to_VAE.as_path().to_str().unwrap()).await?; + + let weights1 = CreateModelWeightsArgs { + token: &model_weight_token1, // replace with actual ModelWeightToken + weights_type: WeightsType::StableDiffusion15, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "Dragonfruit".to_string(), + maybe_description_markdown: Some(sd1_5_markdown_description.to_string()), + maybe_description_rendered_html: Some("

Description

".to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("Last Update User Token".to_string()), + original_download_url: Some("https://civitai.com/models/41865/majicmix-fantasy".to_string()), + original_filename: Some("majicmixFantasy_v30Vae.safetensors".to_string()), + file_size_bytes: metadata1.file_size_bytes as i32, + file_checksum_sha2: metadata1.sha256_checksum.to_string(), + public_bucket_hash: metadata1.bucket_details.clone().unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata1.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata1.bucket_details.clone().unwrap().suffix), + version: 1, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + let weights2 = CreateModelWeightsArgs { + token: &model_weight_token2, // replace with actual ModelWeightToken + weights_type: WeightsType::LoRA, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "MIHOYO Collection 米家全家桶 (Honkai Impact 3rd | Honkai Star Rail | Genshin Impact | Zenless Zone Zero)".to_string(), + maybe_description_markdown: Some(miyhoyo_description.to_string()), + maybe_description_rendered_html: Some(miyhoyo_description.to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "292.268.2.2", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("

Honkai

".to_string()), + original_download_url: Some("https://civitai.com/models/95243/mihoyo-collection-honkai-impact-3rd-or-honkai-star-rail-or-genshin-impact-or-zenless-zone-zero".to_string()), + original_filename: Some("xiawolei-v100-000019.safetensors".to_string()), + file_size_bytes: metadata2.file_size_bytes as i32, + file_checksum_sha2: metadata2.sha256_checksum.to_string(), + public_bucket_hash: metadata2.bucket_details.clone().unwrap().object_hash.clone(), + maybe_public_bucket_prefix: Some(metadata2.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata2.bucket_details.clone().unwrap().suffix), + version: 2, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + let weights3 = CreateModelWeightsArgs { + token: &model_weight_token3, // replace with actual ModelWeightToken + weights_type: WeightsType::StableDiffusion15, // replace with actual WeightsType + weights_category: WeightsCategory::ImageGeneration, // replace with actual WeightsCategory + title: "z-vae".to_string(), + maybe_description_markdown: Some(sd_vae_description.to_string()), + maybe_description_rendered_html: Some("

Description

".to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("Last Update User Token".to_string()), + original_download_url: Some("https://civitai.com/models/97653/z-vae".to_string()), + original_filename: Some("zVae_v20.safetensors".to_string()), + file_size_bytes: metadata3.file_size_bytes as i32, + file_checksum_sha2: metadata3.sha256_checksum.to_string(), + public_bucket_hash: metadata3.bucket_details.clone().unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata3.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata3.bucket_details.clone().unwrap().suffix), + version: 1, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + + create_weight(weights1).await?; + create_weight(weights2).await?; + create_weight(weights3).await?; + + Ok(()) +} + +pub async fn seed_workflows_for_testing_inference(mysql_pool: &Pool, user_token: UserToken) -> AnyhowResult<()>{ + let model_weight_token4 = ModelWeightToken::generate_for_testing_and_dev_seeding_never_use_in_production_seriously(); + let mut path_to_comfy = get_seed_tool_data_root(); + path_to_comfy.push("models/workflows/comfyui/test-workflow.json"); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + let comfy_weights_descriptor = Box::new(WeightsWorkflowDescriptor {}); + let metadata4 = remote_cloud_file_client.upload_file(comfy_weights_descriptor,path_to_comfy.as_path().to_str().unwrap()).await?; + let weights4 = CreateModelWeightsArgs { + token: &model_weight_token4, // replace with actual ModelWeightToken + weights_type: WeightsType::ComfyUi, // replace with actual WeightsType + weights_category: WeightsCategory::WorkflowConfig, // replace with actual WeightsCategory + title: "comfy-test-workflow".to_string(), + maybe_description_markdown: Some("Test workflow for ComfyUI".to_string()), + maybe_description_rendered_html: Some("

Description

".to_string()), + creator_user_token: Some(&user_token), // replace with actual UserToken + creator_ip_address: "192.168.1.1", + creator_set_visibility: Visibility::Public, + maybe_last_update_user_token: Some("Last Update User Token".to_string()), + original_download_url: Some("https://github.com/comfyanonymous/ComfyUI".to_string()), + original_filename: Some("test-workflow.json".to_string()), + file_size_bytes: metadata4.file_size_bytes as i32, + file_checksum_sha2: metadata4.sha256_checksum.to_string(), + public_bucket_hash: metadata4.bucket_details.clone().unwrap().object_hash, + maybe_public_bucket_prefix: Some(metadata4.bucket_details.clone().unwrap().prefix), + maybe_public_bucket_extension: Some(metadata4.bucket_details.clone().unwrap().suffix), + version: 1, + mysql_pool: &mysql_pool, // replace with actual MySqlPool + }; + create_weight(weights4).await?; + + Ok(()) +} + pub async fn seed_weights_for_user_token( mysql_pool: &Pool, user_token: UserToken @@ -630,6 +1016,7 @@ pub async fn original_seed_weights(mysql_pool: &Pool, _user_token: UserTo Ok(()) } + pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { info!("Seeding weights..."); @@ -642,7 +1029,10 @@ pub async fn seed_weights(mysql_pool: &Pool) -> AnyhowResult<()> { //original_seed_weights(mysql_pool,user_token).await?; //seed_weights_for_user_token(mysql_pool, user_token).await?; - seed_weights_for_paging(mysql_pool,user_token).await?; - + //seed_weights_for_paging(mysql_pool,user_token).await?; + // seed_weights_for_testing_inference(mysql_pool,user_token).await?; + //println!("TESTING DOWLOAD"); + //test_seed_weights_files().await?; + seed_workflows_for_testing_inference(mysql_pool,user_token).await?; Ok(()) } diff --git a/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs b/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs index 474f08503d..a0cc6f133d 100644 --- a/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs +++ b/crates/cli/dev-database-seed/src/seeding/seed_model_weights.rs @@ -226,3 +226,5 @@ async fn seed_file_to_bucket( maybe_bucket_extension: maybe_bucket_extension.map(|s| s.to_string()), }) } + + diff --git a/crates/cli/migration-tool/src/cli_args.rs b/crates/cli/migration-tool/src/cli_args.rs new file mode 100644 index 0000000000..0d39c15587 --- /dev/null +++ b/crates/cli/migration-tool/src/cli_args.rs @@ -0,0 +1,47 @@ +use std::str::FromStr; + +use clap::Parser; +use strum::{EnumCount, EnumString, IntoEnumIterator}; +use strum::EnumIter; + +use errors::{anyhow, AnyhowResult}; + +pub struct ParsedArgs { + pub action: Action, +} + +#[derive(Clone, Copy, Debug, EnumIter, EnumCount, EnumString, strum::Display)] +#[strum(serialize_all = "snake_case")] +pub enum Action { + /// Migrate VC + MigrateVoiceConversion, + /// Migrate TTS + MigrateTts, +} + +#[derive(Parser, Debug)] +#[command(name="migration-tool")] +pub struct Args { + #[arg(name="action", long="action", help="action to take", required=true)] + action: String, +} + +pub fn parse_cli_args() -> AnyhowResult { + let args = Args::parse(); + + Ok(ParsedArgs { + action: action_from_str(&args.action)?, + }) +} + + +fn action_from_str(value: &str) -> AnyhowResult { + let action = Action::from_str(value) + .map_err(|err| { + let choices = Action::iter() + .map(|e| e.to_string()) + .collect::>(); + anyhow!("parse error: {:?}, provided: \"{}\" choices: {:?}", err, value, choices) + })?; + Ok(action) +} diff --git a/crates/cli/migration-tool/src/main.rs b/crates/cli/migration-tool/src/main.rs index 9fdf2a3452..71a4270e68 100644 --- a/crates/cli/migration-tool/src/main.rs +++ b/crates/cli/migration-tool/src/main.rs @@ -5,17 +5,21 @@ use std::time::Duration; + use log::info; use sqlx::{MySql, Pool}; use sqlx::mysql::MySqlPoolOptions; -use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_client::BucketClient; use config::shared_constants::DEFAULT_RUST_LOG; use errors::AnyhowResult; +use crate::cli_args::{Action, parse_cli_args}; use crate::deps::Deps; +use crate::migrations::tts_models_to_weights::migrate::migrate_tts_to_weights; use crate::migrations::voice_conversion_to_weights::migrate::migrate_voice_conversion_to_weights; +pub mod cli_args; pub mod deps; pub mod migrations; @@ -40,10 +44,20 @@ pub async fn main() -> AnyhowResult<()> { bucket_production_private: get_bucket_client("PRODUCTION_PRIVATE")?, }; - migrate_voice_conversion_to_weights(&deps).await?; + let args = parse_cli_args()?; + + match args.action { + Action::MigrateVoiceConversion => { + migrate_voice_conversion_to_weights(&deps).await?; + } + Action::MigrateTts => { + migrate_tts_to_weights(&deps).await?; + } + } Ok(()) } + async fn get_mysql(env_var_name: &str) -> AnyhowResult> { info!("Connecting to MySQL {env_var_name}..."); diff --git a/crates/cli/migration-tool/src/migrations/mod.rs b/crates/cli/migration-tool/src/migrations/mod.rs index c0183ac4c9..63f3ec6607 100644 --- a/crates/cli/migration-tool/src/migrations/mod.rs +++ b/crates/cli/migration-tool/src/migrations/mod.rs @@ -1 +1,2 @@ +pub mod tts_models_to_weights; pub mod voice_conversion_to_weights; \ No newline at end of file diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs new file mode 100644 index 0000000000..41880f1e24 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/copy_cloud_files.rs @@ -0,0 +1,47 @@ +use tempdir::TempDir; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use errors::AnyhowResult; +use filesys::safe_delete_temp_directory::safe_delete_temp_directory; +use filesys::safe_delete_temp_file::safe_delete_temp_file; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_tts_model::CopiedTtsFileData; +use mysql_queries::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::WholeTtsModelRecord; + +use crate::deps::Deps; + +pub async fn copy_cloud_files(model: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult { + let copied_file_data = copy_model(model, deps).await?; + Ok(copied_file_data) +} + +async fn copy_model(model: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult { + let bucket_path_unifier = BucketPathUnifier::default_paths(); + + let old_model_bucket_path = bucket_path_unifier.tts_synthesizer_path(&model.private_bucket_hash); + + // TODO(bt,2023-12-19): Probably faster to stream between buckets, but whatever. + let temp_dir = TempDir::new("model_transfer")?; + let model_temp_fs_path = temp_dir.path().join("model.bin"); + + deps.bucket_production_private.download_file_to_disk(&old_model_bucket_path, &model_temp_fs_path).await?; + + let file_checksum = sha256_hash_file(&model_temp_fs_path)?; + + let new_model_bucket_path = WeightFileBucketPath::generate_for_tt2_model(); + + deps.bucket_production_public.upload_filename_with_content_type( + &new_model_bucket_path.get_full_object_path_str(), + &model_temp_fs_path, + "application/octet-stream").await?; + + safe_delete_temp_file(&model_temp_fs_path); + safe_delete_temp_directory(&temp_dir); + + Ok(CopiedTtsFileData { + bucket_path: new_model_bucket_path, + file_sha_hash: file_checksum, + }) +} + diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs new file mode 100644 index 0000000000..50f10e04c9 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/migrate.rs @@ -0,0 +1,93 @@ +use std::thread; +use std::time::Duration; + +use log::{info, warn}; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; +use mysql_queries::queries::model_weights::migration::upsert_model_weight_from_tts_model::{CopiedTtsFileData, upsert_model_weight_from_tts_model}; +use mysql_queries::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::{list_whole_tts_models_using_cursor, WholeTtsModelRecord}; +use tokens::tokens::model_weights::ModelWeightToken; + +use crate::deps::Deps; +use crate::migrations::tts_models_to_weights::copy_cloud_files::copy_cloud_files; + +const PAGE_SIZE: u64 = 10; + +// TODO(bt,2024-01-15): This hasn't been fully vetted / tested. Do not run until it's combed over. + +pub async fn migrate_tts_to_weights(deps: &Deps) -> AnyhowResult<()> { + + let mut cursor = 0; + + loop { + info!("Querying {PAGE_SIZE} models at cursor = {cursor}"); + + let results + = list_whole_tts_models_using_cursor(&deps.mysql_production, PAGE_SIZE, cursor).await?; + + if results.is_empty() { + warn!("No more results found; exiting."); + break; + } + + for result in results.iter() { + println!("\n\nmigrating over result {:?} : {:?} ...", result.token, result); + copy_model_record_and_bucket_files(result, &deps).await?; + } + + if let Some(last_id) = results.last().map(|result| result.id) { + cursor = last_id as u64; + } + + thread::sleep(Duration::from_secs(2)); + } + + Ok(()) +} + +async fn copy_model_record_and_bucket_files(record: &WholeTtsModelRecord, deps: &Deps) -> AnyhowResult<()> { + let mut maybe_copied_data = None; + + match &record.maybe_migration_new_model_weights_token { + None => { + info!("Copying bucket files for record {:?} ...", record.token); + let result = copy_cloud_files(record, &deps).await?; + maybe_copied_data = Some(result); + } + Some(token) => { + info!("Existing migrated record; skipping cloud bucket copy..."); + maybe_copied_data = lookup_bucket_details_from_record(token, &deps).await?; + } + } + + if let Some(copied_data) = &maybe_copied_data { + info!("Upserting records..."); + upsert_model_weight_from_tts_model(record, &deps.mysql_production, copied_data).await?; + } + + Ok(()) +} + +/// Look up the bucket storage details from a previously-migrated record +async fn lookup_bucket_details_from_record(model_weight_token: &ModelWeightToken, deps: &Deps) -> AnyhowResult> { + let maybe_model_weight_record = + get_weight_by_token(model_weight_token, true, &deps.mysql_production).await?; + + let model_weight_record = match maybe_model_weight_record { + None => return Ok(None), + Some(record) => record, + }; + + let bucket_path = match model_weight_record.weights_type { + WeightsType::Tacotron2 => WeightFileBucketPath::tt2_model_file_from_object_hash(&model_weight_record.public_bucket_hash), + _ => return Ok(None), // NB: SVC, RVC, Stable Diffusion, etc. aren't valid for TTS migration + }; + + Ok(Some(CopiedTtsFileData { + bucket_path, + file_sha_hash: model_weight_record.file_checksum_sha2, + })) +} \ No newline at end of file diff --git a/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs new file mode 100644 index 0000000000..a2f14892d9 --- /dev/null +++ b/crates/cli/migration-tool/src/migrations/tts_models_to_weights/mod.rs @@ -0,0 +1,2 @@ +pub mod copy_cloud_files; +pub mod migrate; \ No newline at end of file diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs index 7b790fc432..dc7ef315df 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/bucket_orchestration.rs @@ -19,7 +19,6 @@ pub trait BucketOrchestrationCore { bytes: &[u8], content_type: &str, is_public: bool) -> AnyhowResult<()>; - } diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs index 8066d28659..b7f80df895 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/file_meta_data.rs @@ -1,6 +1,9 @@ +use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; + #[derive(Debug, Clone)] pub struct FileMetaData { pub file_size_bytes: u64, pub sha256_checksum: String, - pub mimetype: String + pub mimetype: String, + pub bucket_details: Option } diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs index e69de29bb2..d969193271 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/media_descriptor.rs @@ -0,0 +1,22 @@ +use crate::remote_file_manager::file_descriptor::FileDescriptor; +// TODO ASK BRANDON WHAT path should this be? +const MEDIA_FILE_DIRECTORY: &str = "/media"; + +pub struct MediaImagePngDescriptor; + +impl FileDescriptor for MediaImagePngDescriptor { + fn remote_directory_path(&self) -> &str { + MEDIA_FILE_DIRECTORY + } + fn get_suffix(&self)->String { + "png".to_string() + } + + fn get_prefix(&self)->String { + "image".to_string() + } + + fn is_public(&self) -> bool { + true + } +} diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs index 21b3a3aee1..68627337ff 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_bucket_details.rs @@ -1,7 +1,7 @@ use crate::remote_file_manager::weights_descriptor::{self}; - use super::file_descriptor::FileDescriptor; +#[derive(Debug, Clone)] pub struct RemoteCloudBucketDetails { pub object_hash: String, pub prefix: String, diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs index a57b89f97f..61e1f14e19 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/remote_cloud_file_manager.rs @@ -12,17 +12,13 @@ use crate::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDe use super::file_descriptor::FileDescriptor; use super::file_meta_data::FileMetaData; + pub struct RemoteCloudFileClient { bucket_orchestration_client: Box } impl RemoteCloudFileClient { - pub fn new(bucket_orchestration_client: Box) -> Self { - Self { - bucket_orchestration_client - } - } - + pub async fn get_remote_cloud_file_client() -> AnyhowResult { let bucket_orchestration = match BucketOrchestration::new_bucket_client_from_existing_env() { Ok(client) => client, @@ -36,6 +32,13 @@ impl RemoteCloudFileClient { }) } + pub fn new(bucket_orchestration_client: Box) -> Self { + Self { + bucket_orchestration_client + } + } + + pub async fn download_file(&self, remote_cloud_bucket_details:RemoteCloudBucketDetails, to_system_file_path:String) -> AnyhowResult<()> { let file_descriptor = remote_cloud_bucket_details.file_descriptor_from_bucket_details(); let file_bucket_directory = FileBucketDirectory::from_existing_bucket_details(remote_cloud_bucket_details); @@ -53,17 +56,26 @@ impl RemoteCloudFileClient { println!("Reading media file: {:?}", from_system_file_path); // get meta data let bytes = file_read_bytes(from_system_file_path)?; - let result = Self::get_file_meta_data(from_system_file_path)?; + let mut result = Self::get_file_meta_data(from_system_file_path)?; let is_public = file_descriptor.is_public(); + let suffix = file_descriptor.get_suffix().clone(); + let prefix = file_descriptor.get_prefix().clone(); + let directory = FileBucketDirectory::generate_new( file_descriptor ); + + result.bucket_details = Some(RemoteCloudBucketDetails { + object_hash: directory.get_file_object_hash().to_string(), + suffix, + prefix, + }); println!("Uploading media file to bucket path: {:?}",directory.get_full_remote_cloud_file_path()); self.bucket_orchestration_client.upload_file_with_content_type_process( - &directory.get_remote_cloud_base_directory(), + &directory.get_full_remote_cloud_file_path(), bytes.as_ref(), result.mimetype.as_ref(), is_public @@ -83,7 +95,8 @@ impl RemoteCloudFileClient { Ok(FileMetaData { file_size_bytes, sha256_checksum, - mimetype: mimetype.to_string() + mimetype: mimetype.to_string(), + bucket_details: None }) } } @@ -125,7 +138,8 @@ mod tests { is_public: bool) -> AnyhowResult<()> { println!("Upload File to Disk"); println!("{}",object_name); - assert_eq!(object_name,String::from("/weights/2/y/q/m/2/2yqm2f1bamh88seyd690h9v24apgezhr/loRA_2yqm2f1bamh88seyd690h9v24apgezhr.safetensors")); + // this is random you have to just check the outputs + //assert_eq!(object_name,String::from("/weights/2/y/q/m/2/2yqm2f1bamh88seyd690h9v24apgezhr/loRA_2yqm2f1bamh88seyd690h9v24apgezhr.safetensors")); println!("ContentType:{}",content_type); println!("{}",is_public); assert_eq!(is_public,true); diff --git a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs index 6c59c0f30b..0bbbfac9ce 100644 --- a/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs +++ b/crates/lib/api_clients/cloud_storage/src/remote_file_manager/weights_descriptor.rs @@ -3,7 +3,6 @@ use super::file_descriptor::FileDescriptor; pub struct WeightsLoRADescriptor; const REMOTE_FILE_DIRECTORY: &str = "/weights"; - impl FileDescriptor for WeightsLoRADescriptor { fn remote_directory_path(&self) -> &str { REMOTE_FILE_DIRECTORY @@ -26,6 +25,7 @@ impl FileDescriptor for WeightsLoRADescriptor { } } + pub struct WeightsSD15Descriptor; impl FileDescriptor for WeightsSD15Descriptor { @@ -166,6 +166,26 @@ impl FileDescriptor for crate::remote_file_manager::weights_descriptor::WeightsV } // This will be ensure that the right bucket is picked + fn is_public(&self) -> bool { + true + } +} + +pub struct WeightsWorkflowDescriptor; + +impl FileDescriptor for WeightsWorkflowDescriptor { + fn remote_directory_path(&self) -> &str { + REMOTE_FILE_DIRECTORY + } + + fn get_suffix(&self) -> String { + "json".to_string() + } + + fn get_prefix(&self) -> String { + "workflow".to_string() + } + fn is_public(&self) -> bool { true } diff --git a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs index 292de11bc0..56ab8d4931 100644 --- a/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs +++ b/crates/lib/api_clients/google_drive_common/src/google_drive_download_command.rs @@ -57,6 +57,59 @@ impl GoogleDriveDownloadCommand { } } + pub async fn download_file_with_file_name(&self, + download_url: &str, + temp_dir: &TempDir, + file_name:&str) -> AnyhowResult + { + let temp_dir_path = temp_dir.path() + .to_str() + .unwrap_or("/tmp") + .to_string(); + + let temp_filename = format!("{}/{}", temp_dir_path, file_name); + + info!("Downloading {} to: {}", download_url, temp_filename); + + let mut command = format!("{} --url \"{}\" --output_filename {}", + &self.download_script, + download_url, + &temp_filename); + + if let Some(venv_activation_script) = self.maybe_venv_activation_script.as_deref() { + // NB: "." is source for non-bash shells + command = format!(". {} && {}", + venv_activation_script, + &command); + } + + if let Some(change_to_directory) = self.maybe_change_to_directory.as_deref() { + command = format!("cd {} && {}", + change_to_directory, + &command); + } + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Running command: {}", command); + + let result = Command::new("sh") + .arg("-c") + .arg(command) + .output()?; + + info!("Downloader Result: {:?}", result); + + if !result.status.success() { + let reason = String::from_utf8(result.stderr).unwrap_or("UNKNOWN".to_string()); + return Err(anyhow!("Failure to execute command: {:?}", reason)) + } + + Ok(temp_filename) + } + /// Download file from Google Drive into the `TempDir`. /// Return the local filesystem filename. pub async fn download_file(&self, diff --git a/crates/lib/deprecated/http_server_common/src/cors.rs b/crates/lib/deprecated/http_server_common/src/cors.rs index 95f51ff0e9..a72c174710 100644 --- a/crates/lib/deprecated/http_server_common/src/cors.rs +++ b/crates/lib/deprecated/http_server_common/src/cors.rs @@ -59,6 +59,8 @@ pub fn add_fakeyou(cors: Cors, is_production: bool) -> Cors { // FakeYou (Staging) .allowed_origin("http://staging.fakeyou.com") .allowed_origin("https://staging.fakeyou.com") + // FakeYou (Netlify Staging / Production) + .allowed_origin("https://feature-mvp--fakeyou.netlify.app") } else { cors // FakeYou (Development) diff --git a/crates/lib/migration/src/lib.rs b/crates/lib/migration/src/lib.rs index 89f29d3e6b..ff860f3c6f 100644 --- a/crates/lib/migration/src/lib.rs +++ b/crates/lib/migration/src/lib.rs @@ -18,4 +18,5 @@ #![allow(dead_code)] #![allow(non_snake_case)] +pub mod text_to_speech; pub mod voice_conversion; diff --git a/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs b/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs new file mode 100644 index 0000000000..d2a7439949 --- /dev/null +++ b/crates/lib/migration/src/text_to_speech/list_tts_models_for_migration.rs @@ -0,0 +1,155 @@ +use chrono::{DateTime, Utc}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; + +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use mysql_queries::queries::model_weights::list::list_model_weights_for_text_to_speech::{list_model_weights_for_text_to_speech, ModelWeightForTts}; +use mysql_queries::queries::tts::tts_models::list_tts_models::{list_tts_models_with_connection, TtsModelRecordForList}; + +/// List TTS models +/// This is for the tts model list page. +/// Since we're listing, we have to use a flag to determine which query to perform. +pub async fn list_tts_models_for_migration( + mysql_connection: &mut PoolConnection, + use_weights_table: bool, +) -> AnyhowResult> { + // NB: This is temporary migration code as we switch from the `tts_models` table to the `model_weights` table. + if use_weights_table { + let models = list_model_weights_for_text_to_speech( + mysql_connection).await?; + + Ok(models.into_iter() + .map(|model| TtsModelForList::ModelWeight(model)) + .collect()) + + } else { + let models = list_tts_models_with_connection( + mysql_connection, None, false).await?; + + Ok(models.into_iter() + .map(|model| TtsModelForList::LegacyTts(model)) + .collect()) + } +} + +/// Union over the legacy table and the new table to support an easier migration. +/// This enum can hold a record of either type and present a unified accessor interface. +#[derive(Clone)] +pub enum TtsModelForList { + /// Old type from the `tts_models` table, on the way out + LegacyTts(TtsModelRecordForList), + /// New type, replacing the `tts_models` table. + ModelWeight(ModelWeightForTts), +} + +impl TtsModelForList { + pub fn token(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => model.model_token.as_str(), + TtsModelForList::ModelWeight(ref model) => model.token.as_str(), + } + } + + pub fn title(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.title, + TtsModelForList::ModelWeight(ref model) => &model.title, + } + } + + pub fn ietf_language_tag(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.ietf_language_tag, + TtsModelForList::ModelWeight(ref model) => &model.ietf_language_tag, + } + } + + pub fn ietf_primary_language_subtag(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.ietf_primary_language_subtag, + TtsModelForList::ModelWeight(ref model) => &model.ietf_primary_language_subtag, + } + } + + pub fn creator_user_token(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_user_token, + TtsModelForList::ModelWeight(ref model) => model.creator_user_token.as_str(), + } + } + + pub fn creator_username(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_username, + TtsModelForList::ModelWeight(ref model) => &model.creator_username, + } + } + + pub fn creator_display_name(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_display_name, + TtsModelForList::ModelWeight(ref model) => &model.creator_display_name, + } + } + + pub fn creator_gravatar_hash(&self) -> &str { + match self { + TtsModelForList::LegacyTts(ref model) => &model.creator_gravatar_hash, + TtsModelForList::ModelWeight(ref model) => &model.creator_gravatar_hash, + } + } + + pub fn creator_set_visibility(&self) -> Visibility{ + match self { + TtsModelForList::LegacyTts(ref model) => model.creator_set_visibility, + TtsModelForList::ModelWeight(ref model) => model.creator_set_visibility, + } + } + + pub fn created_at(&self) -> &DateTime { + match self { + TtsModelForList::LegacyTts(ref model) => &model.created_at, + TtsModelForList::ModelWeight(ref model) => &model.created_at, + } + } + + pub fn updated_at(&self) -> &DateTime { + match self { + TtsModelForList::LegacyTts(ref model) => &model.updated_at, + TtsModelForList::ModelWeight(ref model) => &model.updated_at, + } + } + + // Not supported in the new table + pub fn is_front_page_featured(&self) -> bool { + match self { + TtsModelForList::LegacyTts(ref model) => model.is_front_page_featured, + TtsModelForList::ModelWeight(ref _model) => false, + } + } + + // Not supported in the new table + pub fn user_ratings_positive_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_positive_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } + + // Not supported in the new table + pub fn user_ratings_negative_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_negative_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } + + // Not supported in the new table + pub fn user_ratings_total_count(&self) -> u32 { + match self { + TtsModelForList::LegacyTts(ref model) => model.user_ratings_total_count, + TtsModelForList::ModelWeight(ref _model) => 0, + } + } +} diff --git a/crates/lib/migration/src/text_to_speech/mod.rs b/crates/lib/migration/src/text_to_speech/mod.rs new file mode 100644 index 0000000000..2e243cfce7 --- /dev/null +++ b/crates/lib/migration/src/text_to_speech/mod.rs @@ -0,0 +1 @@ +pub mod list_tts_models_for_migration; \ No newline at end of file diff --git a/crates/lib/migration/src/voice_conversion/list_vc_models_for_migration.rs b/crates/lib/migration/src/voice_conversion/list_vc_models_for_migration.rs index 9028b510f8..613d823557 100644 --- a/crates/lib/migration/src/voice_conversion/list_vc_models_for_migration.rs +++ b/crates/lib/migration/src/voice_conversion/list_vc_models_for_migration.rs @@ -67,6 +67,7 @@ impl VcModelForList { WeightsType::Tacotron2 => None, WeightsType::LoRA => None, WeightsType::VallE => None, + WeightsType::ComfyUi => None, }, } } diff --git a/crates/lib/migration/src/voice_conversion/query_vc_model_for_migration.rs b/crates/lib/migration/src/voice_conversion/query_vc_model_for_migration.rs index 40f714ead7..c282a7d631 100644 --- a/crates/lib/migration/src/voice_conversion/query_vc_model_for_migration.rs +++ b/crates/lib/migration/src/voice_conversion/query_vc_model_for_migration.rs @@ -106,6 +106,7 @@ impl VcModel { WeightsType::Tacotron2 => VcModelType::Invalid, WeightsType::LoRA => VcModelType::Invalid, WeightsType::VallE => VcModelType::Invalid, + WeightsType::ComfyUi => VcModelType::Invalid, } } } diff --git a/crates/lib/migration/src/voice_conversion/query_vc_model_info_lite_for_migration.rs b/crates/lib/migration/src/voice_conversion/query_vc_model_info_lite_for_migration.rs index 8912dfd5c7..98c4fbceb9 100644 --- a/crates/lib/migration/src/voice_conversion/query_vc_model_info_lite_for_migration.rs +++ b/crates/lib/migration/src/voice_conversion/query_vc_model_info_lite_for_migration.rs @@ -77,6 +77,7 @@ impl VcModelLite { WeightsType::HifiganTacotron2 => None, // NB: Not used directly in inference. Included by other models. WeightsType::LoRA => None, // NB: not used directly in inference. Included in other models. WeightsType::VallE => None, // TODO(bt,2023-12-18): Does this mean speaker weights? + WeightsType::ComfyUi => None, // NB: not used directly in inference. Included in other models. } } } @@ -97,6 +98,7 @@ impl VcModelLite { WeightsType::HifiganTacotron2 => None, // NB: Not used directly in inference. Included by other models. WeightsType::LoRA => None, // NB: not used directly in inference. Included in other models. WeightsType::VallE => None, // TODO(bt,2023-12-18): Does this mean speaker weights? + WeightsType::ComfyUi => None, // NB: not used directly in inference. Included in other models. } } } diff --git a/crates/lib/mysql_queries/Cargo.toml b/crates/lib/mysql_queries/Cargo.toml index 228c472322..2dcf665e29 100644 --- a/crates/lib/mysql_queries/Cargo.toml +++ b/crates/lib/mysql_queries/Cargo.toml @@ -16,6 +16,7 @@ path = "src/lib.rs" # Internal buckets = { path = "../../_types/buckets" } +composite_identifiers = { path = "../../_types/composite_identifiers" } config = { path = "../deprecated/config" } container_common = { path = "../deprecated/container_common" } enums = { path = "../../_types/enums" } diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs index 5ec0c07c54..10547fca35 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/generic_inference_args.rs @@ -1,13 +1,13 @@ use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; use errors::AnyhowResult; +use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use crate::payloads::generic_inference_args::lipsync_payload::LipsyncArgs; use crate::payloads::generic_inference_args::mocap_payload::MocapArgs; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; -use crate::payloads::generic_inference_args::videofilter_payload::{RerenderArgs}; +use crate::payloads::generic_inference_args::videofilter_payload::RerenderArgs; use crate::payloads::generic_inference_args::workflow_payload::WorkflowArgs; - /// Used to encode extra state for the `generic_inference_jobs` table. /// This should act somewhat like a serialized protobuf stored inside a record. #[derive(Clone, Debug, Serialize, Deserialize)] @@ -43,6 +43,10 @@ pub enum InferenceCategoryAbbreviated { #[serde(alias = "video_filter")] VideoFilter, + #[serde(rename = "ig")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + #[serde(alias = "image_generation")] + ImageGeneration, + #[serde(rename = "mc")] // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. #[serde(alias = "mocap")] Mocap, @@ -103,6 +107,9 @@ pub enum PolymorphicInferenceArgs { /// Rerender a video. (Short name to save space when serializing.) Rr(RerenderArgs), + /// Image generation. (Short name to save space when serializing.) + Ig(StableDiffusionArgs), + // Mocap (Short name to save space when serializing.) Mc(MocapArgs), @@ -110,7 +117,6 @@ pub enum PolymorphicInferenceArgs { Wf(WorkflowArgs) } - impl GenericInferenceArgs { pub fn from_json(json: &str) -> AnyhowResult { @@ -129,6 +135,7 @@ impl InferenceCategoryAbbreviated { InferenceCategory::TextToSpeech => Self::TextToSpeech, InferenceCategory::VoiceConversion => Self::VoiceConversion, InferenceCategory::VideoFilter => Self::VideoFilter, + InferenceCategory::ImageGeneration => Self::ImageGeneration, InferenceCategory::Mocap => Self::Mocap, InferenceCategory::Workflow => Self::Workflow, } @@ -140,6 +147,7 @@ impl InferenceCategoryAbbreviated { Self::TextToSpeech => InferenceCategory::TextToSpeech, Self::VoiceConversion => InferenceCategory::VoiceConversion, Self::VideoFilter => InferenceCategory::VideoFilter, + Self::ImageGeneration =>InferenceCategory::ImageGeneration, Self::Mocap => InferenceCategory::Mocap, Self::Workflow => InferenceCategory::Workflow, } @@ -148,7 +156,10 @@ impl InferenceCategoryAbbreviated { #[cfg(test)] mod tests { + use tokens::tokens::{media_files::MediaFileToken, model_weights::ModelWeightToken}; + use crate::payloads::generic_inference_args::generic_inference_args::{FundamentalFrequencyMethodForJob, GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; + use crate::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; use crate::payloads::generic_inference_args::lipsync_payload::{LipsyncAnimationAudioSource, LipsyncAnimationImageSource, LipsyncArgs}; use crate::payloads::generic_inference_args::tts_payload::TTSArgs; @@ -197,6 +208,48 @@ mod tests { assert!(json.len() < 1000); } + #[test] + fn typical_image_gen_args_serialize() { + + let video_media_token = MediaFileToken("video_media_token".to_string()); + let image_media_token = MediaFileToken("image_media_token".to_string()); + let sd_model_token = ModelWeightToken("sd_model_token".to_string()); + let lora_model_token = ModelWeightToken("lora_model_token".to_string()); + + let prompt = "prompt".to_string(); + let a_prompt = "a_prompt".to_string(); + let n_prompt = "n_prompt".to_string(); + let seed = 1; + let upload_path = "upload_path".to_string(); + let lora_upload_path = "lora_upload_path".to_string(); + let checkpoint = "checkpoint".to_string(); + let type_of_inference = "type_of_inference".to_string(); + let args = GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), + args: Some(PolymorphicInferenceArgs::Ig(StableDiffusionArgs { + maybe_sd_model_token: Some(sd_model_token), + maybe_lora_model_token: Some(lora_model_token), + maybe_sampler: Some("sampler".to_string()), + maybe_height: Some(512), + maybe_width: Some(512), + maybe_cfg_scale: Some(7), + maybe_prompt: Some(prompt), + maybe_n_prompt: Some(n_prompt), + maybe_batch_count: Some(1), + maybe_number_of_samples: Some(25), + maybe_seed: Some(seed), + maybe_upload_path: Some(upload_path), + maybe_lora_upload_path: Some(lora_upload_path), + type_of_inference, + description: Some("Option".to_string()), + name: Some("Model Name".to_string()) + })), + }; + + let json = serde_json::ser::to_string(&args).unwrap(); + assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"vs":"video_media_token","is":"image_media_token","sd":"sd_model_token","lm":"lora_model_token","p":"prompt","ap":"a_prompt","np":"n_prompt","se":1,"mu":"upload_path","lu":"lora_upload_path"}}}"#.to_string()); + } + #[test] fn typical_voice_conversion_args_serialize() { let args = GenericInferenceArgs { diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs new file mode 100644 index 0000000000..cdae13bdd2 --- /dev/null +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/image_generation_payload.rs @@ -0,0 +1,88 @@ +use tokens::tokens::model_weights::ModelWeightToken; + +/// Video sources can be one of several: +/// - F: media_files (todo) +/// - U: media_uploads (legacy) + +/// For image to image we probably want images +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] +pub enum StableDiffusionSource { + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media File Token (media_files table) + /// Serde cannot yet rename enum variants. + F(String), + + // NB: DO NOT CHANGE. It could break live jobs. Renamed to be fewer bytes. + /// Media Upload Token (media_uploads table) + /// Serde cannot yet rename enum variants. + U(String), +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +pub struct StableDiffusionArgs { + + #[serde(rename = "sd")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sd_model_token: Option, + + #[serde(rename = "lm")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_model_token: Option, + + #[serde(rename = "w")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_width: Option, + + #[serde(rename = "h")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_height: Option, + + #[serde(rename = "s")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_sampler: Option, + + #[serde(rename = "p")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_prompt: Option, + + #[serde(rename = "np")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_n_prompt: Option, + + #[serde(rename = "se")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_seed: Option, + + #[serde(rename = "mu")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_upload_path: Option, + + #[serde(rename = "cf")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_cfg_scale: Option, + + #[serde(rename = "lu")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_lora_upload_path: Option, + + #[serde(rename = "sa")] + #[serde(skip_serializing_if = "Option::is_none")] + pub maybe_number_of_samples: Option, + + #[serde(rename = "bc")] + pub maybe_batch_count: Option, + + #[serde(rename = "t")] + pub type_of_inference: String, + + // loRA / checkpoint description + #[serde(rename = "de")] + #[serde(skip_serializing_if = "Option::is_none")] + pub description: Option, + // loRA / checkpoint name and description + #[serde(rename = "na")] + #[serde(skip_serializing_if = "Option::is_none")] + pub name: Option +} + + diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs index 3ce1305f47..defc7a833b 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/mod.rs @@ -2,5 +2,6 @@ pub mod generic_inference_args; pub mod lipsync_payload; pub mod tts_payload; pub mod videofilter_payload; +pub mod image_generation_payload; pub mod mocap_payload; -pub mod workflow_payload; \ No newline at end of file +pub mod workflow_payload; diff --git a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs index 9cf68aaea4..2aadc215df 100644 --- a/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs +++ b/crates/lib/mysql_queries/src/payloads/generic_inference_args/workflow_payload.rs @@ -1,30 +1,16 @@ use std::collections::HashMap; use tokens::tokens::model_weights::ModelWeightToken; -#[derive(Clone, Debug, Serialize, Deserialize, PartialEq, Eq)] -pub enum FileSource { - F(String), - U(String), -} - -impl FileSource { - pub fn media_file_token(token: &str) -> Self { - FileSource::F(token.to_string()) - } - pub fn media_upload_token(token: &str) -> Self { - FileSource::U(token.to_string()) - } -} - -#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] +#[derive(Clone, Debug, Serialize, Deserialize, PartialEq)] #[serde(untagged)] -enum NewValue { +pub enum NewValue { String(String), Float(f32), Int(i32), + Bool(bool), } -#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq, Eq)] +#[derive(Clone, Debug, Default, Serialize, Deserialize, PartialEq)] pub struct WorkflowArgs { #[serde(rename = "sd")] #[serde(skip_serializing_if = "Option::is_none")] diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs new file mode 100644 index 0000000000..4b9c6f4d56 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs @@ -0,0 +1,99 @@ +use anyhow::anyhow; +use log::error; +use sqlx::{MySql, QueryBuilder, Transaction}; + +use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; +use errors::AnyhowResult; +use tokens::tokens::batch_generations::BatchGenerationToken; + +pub struct BatchEntry { + pub batch_token: BatchGenerationToken, + pub entity: BatchGenerationEntity, +} + +pub struct InsertBatchArgs<'a, 'b> { + pub entries: Vec, + pub transaction: &'a mut Transaction<'b, MySql>, +} + +/// Insert a list of entities into a "batch" together for grouping; returns the new batch token identifier to return +/// to the HTTP caller. +/// +/// NB: Calling code is responsible for rolling back the transaction if this fails. +pub async fn insert_batch_generation_records<'a, 'b>(args: InsertBatchArgs<'a, 'b>) -> AnyhowResult { + let batch_token = BatchGenerationToken::generate(); + + let mut query_builder = QueryBuilder::new(r#" + INSERT INTO batch_generations (token, entity_type, entity_token) VALUES + "#); + + for (i, entry) in args.entries.iter().enumerate() { + let (entity_type, entity_token) = entry.get_composite_keys(); + + query_builder.push("("); + query_builder.push_bind(&batch_token); + query_builder.push(","); + + query_builder.push_bind(entity_type); + query_builder.push(","); + + query_builder.push_bind(entity_token); + query_builder.push(")"); + + if i < args.entries.len() - 1 { + query_builder.push(","); + } + } + + let query = query_builder.build(); + + let query_result = query.execute(&mut **args.transaction).await; + + match query_result { + Ok(_) => Ok(batch_token), + Err(err) => { + error!("Error with batch generation query: {:?}", &err); + Err(anyhow!("model category creation error: {:?}", &err)) + }, + } +} + +#[ignore] +#[cfg(test)] +mod tests { + use sqlx::mysql::MySqlPoolOptions; + + use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity; + use config::shared_constants::DEFAULT_MYSQL_CONNECTION_STRING; + use tokens::tokens::media_files::MediaFileToken; + + use crate::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs}; + + async fn setup() -> sqlx::Pool { + MySqlPoolOptions::new() + .max_connections(3) + .connect(&DEFAULT_MYSQL_CONNECTION_STRING).await + .unwrap() + } + + #[ignore] + #[tokio::test] + async fn test() { + let pool = setup().await; + + let mut transaction = pool.begin().await.unwrap(); + + let entries = vec![ + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_foo")), + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_bar")), + BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_baz")), + ]; + + let r = insert_batch_generation_records(InsertBatchArgs { + entries, + transaction: &mut transaction, + }).await; + + transaction.commit().await.unwrap(); + } +} \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs new file mode 100644 index 0000000000..5f34676f55 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/batch_generations/mod.rs @@ -0,0 +1 @@ +pub mod insert_batch_generation_records; diff --git a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs index fddd32ab3d..dd13a620c0 100644 --- a/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs +++ b/crates/lib/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs @@ -179,6 +179,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::TextToSpeech => record.maybe_tts_model_title.as_deref(), InferenceCategory::VoiceConversion => record.maybe_voice_conversion_model_title.as_deref(), InferenceCategory::VideoFilter => Some("Video Filter"), + InferenceCategory::ImageGeneration => Some("Image Generation"), InferenceCategory::Mocap => Some("Mocap"), InferenceCategory::Workflow => Some("Workflow"), }; @@ -190,6 +191,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::TextToSpeech => (false, record.maybe_tts_public_bucket_path.as_deref()), InferenceCategory::VoiceConversion => (true, record.maybe_voice_conversion_public_bucket_hash.as_deref()), InferenceCategory::VideoFilter => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), + InferenceCategory::ImageGeneration => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::Mocap => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::Workflow => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), }; diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs index e250919752..3be8b2c70f 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_from_zero_shot_tts.rs @@ -18,6 +18,7 @@ pub struct InsertArgs<'a> { pub pool: &'a MySqlPool, pub job: &'a AvailableInferenceJob, + pub origin_model_type: MediaFileOriginModelType, pub maybe_mime_type: Option<&'a str>, pub file_size_bytes: u64, pub sha256_checksum: &'a str, @@ -65,11 +66,10 @@ pub async fn insert_media_file_from_zero_shot( maybe_creator_category_synthetic_id = Some(next_lipsync_id); } - // TODO: Ask if it make sense to use ONE function to do the insert + // TODO: Ask if it make sense to use ONE function to do the insert * I agree it does .... use insert media file generic and eventually remove all this. // Since fill this out apart of the funtion this seems like the only thing to change. const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::ZeroShotVoice; - const ORIGIN_MODEL_TYPE : MediaFileOriginModelType = MediaFileOriginModelType::VallEX; const MEDIA_TYPE : MediaFileType = MediaFileType::Audio; let record_id = { @@ -111,8 +111,8 @@ pub async fn insert_media_file_from_zero_shot( ORIGIN_CATEGORY.to_str(), ORIGIN_PRODUCT_CATEGORY.to_str(), - ORIGIN_MODEL_TYPE.to_str(), - + args.origin_model_type.to_str(), + MEDIA_TYPE.to_str(), args.maybe_mime_type, args.file_size_bytes, diff --git a/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs new file mode 100644 index 0000000000..ec51b21e76 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/media_files/create/insert_media_file_generic.rs @@ -0,0 +1,172 @@ +use anyhow::anyhow; +use sqlx; +use sqlx::MySqlPool; + +use enums::by_table::generic_synthetic_ids::id_category::IdCategory; +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use errors::AnyhowResult; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::weight_model_weight; +use tokens::tokens::users::UserToken; + +use crate::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use crate::queries::generic_synthetic_ids::transactional_increment_generic_synthetic_id::transactional_increment_generic_synthetic_id; +// thought about this it seems like this can be a bit more geneneric instead of having this we can ... +pub struct InsertArgs<'a> { + pub pool: &'a MySqlPool, + pub job: &'a AvailableInferenceJob, + + pub maybe_mime_type: Option<&'a str>, + pub file_size_bytes: u64, + pub sha256_checksum: &'a str, + + pub public_bucket_directory_hash: &'a str, + pub maybe_public_bucket_prefix: Option<&'a str>, + pub maybe_public_bucket_extension: Option<&'a str>, + + pub is_on_prem: bool, + + pub worker_hostname: &'a str, + pub worker_cluster: &'a str, + + pub maybe_creator_user_token: Option<&'a UserToken>, + pub maybe_creator_anonymous_visitor_token: Option<&'a AnonymousVisitorTrackingToken>, + pub creator_ip_address: &'a str, + pub creator_set_visibility: Visibility, + pub upload_type: UploadType, + pub media_file_type: MediaFileType, + pub duration_millis: u64, + + pub media_file_origin: MediaFileOriginProductCategory, + pub media_file_origin_product_category: MediaFileOriginProductCategory, + pub media_file_origin_model_type: Mediafile, + pub maybe_origin_filename: String, + pub maybe_origin_model_token:WeightsModelToken +} +pub async fn insert_media_file_generic( + args: InsertArgs<'_> +) -> AnyhowResult<(MediaFileToken, u64)> +{ + + let result_token = MediaFileToken::generate(); + + let mut maybe_creator_file_synthetic_id : Option = None; + let mut maybe_creator_category_synthetic_id : Option = None; + + let mut transaction = args.pool.begin().await?; + + // create a user token + if let Some(creator_user_token) = args.job.maybe_creator_user_token.as_deref() { + let user_token = UserToken::new_from_str(creator_user_token); + + let next_media_file_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::MediaFile, + &mut transaction + ).await?; + + let next_lipsync_id = transactional_increment_generic_synthetic_id( + &user_token, + IdCategory::ZeroShotTtsResult, + &mut transaction + ).await?; + + maybe_creator_file_synthetic_id = Some(next_media_file_id); + maybe_creator_category_synthetic_id = Some(next_lipsync_id); + } + + // TODO: Ask if it make sense to use ONE function to do the insert * I agree it does .... use insert media file generic and eventually remove all this. + // Since fill this out apart of the funtion this seems like the only thing to change. + const ORIGIN_CATEGORY : MediaFileOriginCategory = MediaFileOriginCategory::Inference; + const ORIGIN_PRODUCT_CATEGORY : MediaFileOriginProductCategory = MediaFileOriginProductCategory::ZeroShotVoice; + const ORIGIN_MODEL_TYPE : MediaFileOriginModelType = MediaFileOriginModelType::VallEX; + const MEDIA_TYPE : MediaFileType = MediaFileType::Audio; + + let record_id = { + let query_result = sqlx::query!( + r#" + INSERT INTO media_files + SET + token = ?, + + origin_category = ?, + origin_product_category = ?, + maybe_origin_model_type = ?, + + media_type = ?, + maybe_mime_type = ?, + file_size_bytes = ?, + + checksum_sha2 = ?, + + public_bucket_directory_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + + maybe_creator_file_synthetic_id = ?, + maybe_creator_category_synthetic_id = ?, + + is_generated_on_prem = ?, + generated_by_worker = ?, + generated_by_cluster = ? + + "#, + result_token.as_str(), + + ORIGIN_CATEGORY.to_str(), + ORIGIN_PRODUCT_CATEGORY.to_str(), + ORIGIN_MODEL_TYPE.to_str(), + + MEDIA_TYPE.to_str(), + args.maybe_mime_type, + args.file_size_bytes, + + args.sha256_checksum, + + args.public_bucket_directory_hash, + args.maybe_public_bucket_prefix, + args.maybe_public_bucket_extension, + + args.job.maybe_creator_user_token, + args.job.maybe_creator_anonymous_visitor_token, + args.job.creator_ip_address, + + args.job.creator_set_visibility.to_str(), + + maybe_creator_file_synthetic_id, + maybe_creator_category_synthetic_id, + + args.is_on_prem, + args.worker_hostname, + args.worker_cluster + ) + .execute(&mut *transaction) + .await; + + let record_id = match query_result { + Ok(res) => { + res.last_insert_id() + }, + Err(err) => { + // TODO: handle better + //transaction.rollback().await?; + return Err(anyhow!("Mysql error: {:?}", err)); + } + }; + + record_id + }; + + transaction.commit().await?; + + Ok((result_token, record_id)) +} \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs index 447a106c90..0d7a14faef 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files.rs @@ -30,9 +30,13 @@ pub struct MediaFileListItem { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, + pub media_type: MediaFileType, pub public_bucket_directory_hash: String, pub maybe_public_bucket_prefix: Option, @@ -98,6 +102,7 @@ pub async fn list_media_files(args: ListMediaFilesArgs<'_>) -> AnyhowResult, maybe_origin_model_token: Option, + maybe_origin_model_title: Option, media_type: MediaFileType, public_bucket_directory_hash: String, @@ -317,6 +327,7 @@ impl FromRow<'_, MySqlRow> for MediaFileListItemInternal { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, media_type: MediaFileType::try_from_mysql_row(row, "media_type")?, public_bucket_directory_hash: row.try_get("public_bucket_directory_hash")?, maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?, diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs index 7adb769287..76a7ab510b 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_by_tokens.rs @@ -19,9 +19,13 @@ pub struct MediaFilesByTokensRecord { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, + pub public_bucket_directory_hash: String, pub maybe_public_bucket_prefix: Option, pub maybe_public_bucket_extension: Option, @@ -68,9 +72,12 @@ async fn get_raw_media_files_by_tokens( m.origin_category, m.origin_product_category, + m.maybe_origin_model_type, m.maybe_origin_model_token, + w.title as maybe_origin_model_title, + users.token as maybe_creator_user_token, users.username as maybe_creator_username, users.display_name as maybe_creator_display_name, @@ -90,6 +97,8 @@ async fn get_raw_media_files_by_tokens( FROM media_files as m LEFT OUTER JOIN users ON users.token = m.maybe_creator_user_token + LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -107,9 +116,12 @@ async fn get_raw_media_files_by_tokens( m.origin_category, m.origin_product_category, + m.maybe_origin_model_type, m.maybe_origin_model_token, + w.title as maybe_origin_model_title, + users.token as maybe_creator_user_token, users.username as maybe_creator_username, users.display_name as maybe_creator_display_name, @@ -129,6 +141,8 @@ async fn get_raw_media_files_by_tokens( FROM media_files as m LEFT OUTER JOIN users ON users.token = m.maybe_creator_user_token + LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -173,6 +187,7 @@ fn map_to_media_files(dataset:Vec) -> Vec) -> Vec, pub maybe_origin_model_token: Option, + pub maybe_origin_model_title: Option, + pub maybe_creator_user_token: Option, pub maybe_creator_username: Option, pub maybe_creator_display_name: Option, @@ -246,6 +264,7 @@ impl FromRow<'_, MySqlRow> for RawMediaFileJoinUser { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, maybe_creator_user_token: row.try_get::, _>("maybe_creator_user_token")? .and_then(|token| Some(UserToken::new_from_str(&token))), maybe_creator_username: row.try_get("maybe_creator_username")?, diff --git a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs index 748cb46d98..9c10eb9e61 100644 --- a/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs +++ b/crates/lib/mysql_queries/src/queries/media_files/list/list_media_files_for_user.rs @@ -26,8 +26,12 @@ pub struct MediaFileListItem { pub origin_category: MediaFileOriginCategory, pub origin_product_category: MediaFileOriginProductCategory, + pub maybe_origin_model_type: Option, pub maybe_origin_model_token: Option, + + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + pub maybe_origin_model_title: Option, pub media_type: MediaFileType, pub public_bucket_directory_hash: String, @@ -97,6 +101,7 @@ pub async fn list_media_files_for_user(args: ListMediaFileForUserArgs<'_>) -> An origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, + maybe_origin_model_title: record.maybe_origin_model_title, media_type: record.media_type, public_bucket_directory_hash: record.public_bucket_directory_hash, maybe_public_bucket_prefix: record.maybe_public_bucket_prefix, @@ -125,13 +130,15 @@ fn select_result_fields() -> String { m.id, m.token, + m.media_type, + m.origin_category, m.origin_product_category, m.maybe_origin_model_type, m.maybe_origin_model_token, - - m.media_type, + + w.title as maybe_origin_model_title, m.public_bucket_directory_hash, m.maybe_public_bucket_prefix, @@ -175,6 +182,8 @@ SELECT FROM media_files AS m LEFT OUTER JOIN users AS u ON m.maybe_creator_user_token = u.token +LEFT OUTER JOIN model_weights as w + ON m.maybe_origin_model_token = w.token LEFT OUTER JOIN entity_stats ON entity_stats.entity_type = "media_file" AND entity_stats.entity_token = m.token @@ -227,8 +236,12 @@ struct MediaFileListItemInternal { origin_category: MediaFileOriginCategory, origin_product_category: MediaFileOriginProductCategory, + maybe_origin_model_type: Option, maybe_origin_model_token: Option, + + // NB: The title won't be populated for `tts_models` records or non-`model_weights` records. + maybe_origin_model_title: Option, media_type: MediaFileType, public_bucket_directory_hash: String, @@ -266,6 +279,7 @@ impl FromRow<'_, MySqlRow> for MediaFileListItemInternal { origin_product_category: MediaFileOriginProductCategory::try_from_mysql_row(row, "origin_product_category")?, maybe_origin_model_type: MediaFileOriginModelType::try_from_mysql_row_nullable(row, "maybe_origin_model_type")?, maybe_origin_model_token: row.try_get("maybe_origin_model_token")?, + maybe_origin_model_title: row.try_get("maybe_origin_model_title")?, media_type: MediaFileType::try_from_mysql_row(row, "media_type")?, public_bucket_directory_hash: row.try_get("public_bucket_directory_hash")?, maybe_public_bucket_prefix: row.try_get("maybe_public_bucket_prefix")?, diff --git a/crates/lib/mysql_queries/src/queries/mod.rs b/crates/lib/mysql_queries/src/queries/mod.rs index a02c0a3a01..86079b7494 100644 --- a/crates/lib/mysql_queries/src/queries/mod.rs +++ b/crates/lib/mysql_queries/src/queries/mod.rs @@ -1,5 +1,6 @@ pub mod api_tokens; pub mod audit_logs; +pub mod batch_generations; pub mod billing; pub mod comments; pub mod email_sender_jobs; @@ -18,12 +19,9 @@ pub mod stats; pub mod trending_model_analytics; pub mod tts; pub mod twitch; -pub mod user_bookmarks; pub mod users; pub mod vocoder; pub mod voice_clone_requests; pub mod voice_conversion; pub mod voice_designer; pub mod w2l; - - diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs new file mode 100644 index 0000000000..014c2c4ba0 --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/list_model_weights_for_text_to_speech.rs @@ -0,0 +1,153 @@ +use anyhow::anyhow; +use chrono::{DateTime, Utc}; +use sqlx::MySql; +use sqlx::pool::PoolConnection; + +use enums::by_table::model_weights::weights_types::WeightsType; +use enums::common::visibility::Visibility; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +#[derive(Serialize, Clone)] +pub struct ModelWeightForTts { + pub token: ModelWeightToken, + pub weight_type: WeightsType, + + pub title: String, + + // We only supported tacotron2 in the legacy case + //pub tts_model_type: String, + + pub ietf_language_tag: String, + pub ietf_primary_language_subtag: String, + + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // NB: We won't port these... + //pub is_locked_from_use: bool, + //pub is_front_page_featured: bool, + //pub is_twitch_featured: bool, + //pub maybe_suggested_unique_bot_command: Option, + //pub user_ratings_positive_count: u32, + //pub user_ratings_negative_count: u32, + //pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + + pub creator_set_visibility: Visibility, + + pub created_at: DateTime, + pub updated_at: DateTime, +} + +// TODO(bt,2024-01-13): This is written to support the migration to the `model_weights` +// table from the `tts_models` table. We should write a more cross-cutting +// query to treat this not as a special cased thing. All weights should be listable. + +/// This is to support the tts model list page. +/// Later this will be migrated or replaced with a more generic query. +pub async fn list_model_weights_for_text_to_speech( + mysql_connection: &mut PoolConnection, +) -> AnyhowResult> { + + let models = + list_tts_models_for_all_creators(mysql_connection).await?; + + Ok(models.into_iter() + .map(|model| { + ModelWeightForTts { + token: model.token, + weight_type: model.weight_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, + ietf_language_tag: model.ietf_language_tag.unwrap_or("en".to_string()), + ietf_primary_language_subtag: model.ietf_primary_language_subtag.unwrap_or("en".to_string()), + creator_set_visibility: model.creator_set_visibility, + created_at: model.created_at, + updated_at: model.updated_at, + } + }) + .collect::>()) +} + +async fn list_tts_models_for_all_creators( + mysql_connection: &mut PoolConnection, +) -> AnyhowResult> { + // NB: Scoped to only tt2 weights + let maybe_results = sqlx::query_as!( + RawModelWeightForTts, + r#" +SELECT + w.token as `token: tokens::tokens::model_weights::ModelWeightToken`, + w.weights_type as `weight_type: enums::by_table::model_weights::weights_types::WeightsType`, + w.title, + w.creator_user_token as `creator_user_token: tokens::tokens::users::UserToken`, + users.username as creator_username, + users.display_name as creator_display_name, + users.email_gravatar_hash as creator_gravatar_hash, + w_extension.ietf_language_tag, + w_extension.ietf_primary_language_subtag, + w.creator_set_visibility as `creator_set_visibility: enums::common::visibility::Visibility`, + w.created_at, + w.updated_at +FROM model_weights as w +LEFT OUTER JOIN model_weights_extension_tts_details as w_extension + ON w_extension.model_weights_token = w.token +JOIN users + ON users.token = w.creator_user_token +WHERE + w.weights_type IN ("tt2") + AND w.user_deleted_at IS NULL + AND w.mod_deleted_at IS NULL + "# + ) + .fetch_all(&mut **mysql_connection) + .await; + + match maybe_results { + Ok(results) => Ok(results), + Err(err) => match err { + sqlx::Error::RowNotFound => Ok(Vec::new()), + _ => { + Err(anyhow!("error querying : {:?}", err)) + } + } + } +} + +struct RawModelWeightForTts { + pub token: ModelWeightToken, + pub weight_type: WeightsType, + + pub title: String, + + // We only supported tacotron2 in the legacy case + //pub tts_model_type: String, + + pub ietf_language_tag: Option, + pub ietf_primary_language_subtag: Option, + + pub creator_user_token: UserToken, + pub creator_username: String, + pub creator_display_name: String, + pub creator_gravatar_hash: String, + + // NB: We won't port these... + //pub is_locked_from_use: bool, + //pub is_front_page_featured: bool, + //pub is_twitch_featured: bool, + //pub maybe_suggested_unique_bot_command: Option, + //pub user_ratings_positive_count: u32, + //pub user_ratings_negative_count: u32, + //pub user_ratings_total_count: u32, // NB: Does not include "neutral" ratings. + + pub creator_set_visibility: Visibility, + + pub created_at: DateTime, + pub updated_at: DateTime, +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs index 3b5b9b35b2..f56b2896be 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/list/mod.rs @@ -1,5 +1,6 @@ +pub mod list_model_weights_for_elastic_search_backfill_using_cursor; +pub mod list_model_weights_for_text_to_speech; pub mod list_model_weights_for_voice_conversion; pub mod list_weights_by_tokens; pub mod list_weights_by_user; -pub mod list_model_weights_for_elastic_search_backfill_using_cursor; pub mod list_weights_query_builder; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs b/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs index 2789fecc14..939e70a942 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/migration/mod.rs @@ -1 +1,2 @@ +pub mod upsert_model_weight_from_tts_model; pub mod upsert_model_weight_from_voice_conversion_model; \ No newline at end of file diff --git a/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs b/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs new file mode 100644 index 0000000000..ec9f0fd24e --- /dev/null +++ b/crates/lib/mysql_queries/src/queries/model_weights/migration/upsert_model_weight_from_tts_model.rs @@ -0,0 +1,249 @@ +use sqlx::{MySql, MySqlPool, Transaction}; + +use buckets::public::weight_files::bucket_file_path::WeightFileBucketPath; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +use errors::AnyhowResult; +use tokens::tokens::model_weights::ModelWeightToken; + +use crate::queries::tts::tts_models::migration::list_whole_tts_models_using_cursor::WholeTtsModelRecord; + +pub struct CopiedTtsFileData { + pub bucket_path: WeightFileBucketPath, + pub file_sha_hash: String, +} + +/// Migrate `tts_models` records to `model_weights` + `model_weights_extension_tts_details` records. +/// This is designed to be idempotent and re-runnable. Any time we re-run the query, we should get the same result. +/// This will enable us to perfect the query and get the write flows online and switched over. +/// +pub async fn upsert_model_weight_from_tts_model( + record: &WholeTtsModelRecord, + mysql_pool: &MySqlPool, + copied_data: &CopiedTtsFileData, +) -> AnyhowResult<()> { + + let mut transaction = mysql_pool.begin().await?; + + let model_weight_token = create_or_generate_token(record); + + upsert_model_weights_record(record, &model_weight_token, copied_data, &mut transaction).await?; + upsert_model_weights_extension_record(record, &model_weight_token, &mut transaction).await?; + + update_original_record(record, &model_weight_token, &mut transaction).await?; + + transaction.commit().await?; + + Ok(()) +} + +pub fn create_or_generate_token(record: &WholeTtsModelRecord) -> ModelWeightToken { + match record.maybe_migration_new_model_weights_token { + Some(ref token) => token.clone(), + None => ModelWeightToken::generate(), + } +} + +pub async fn upsert_model_weights_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + copied_data: &CopiedTtsFileData, + mut transaction: &mut Transaction<'_, MySql>, +) -> AnyhowResult<()> { + + // NB: We never supported other TTS models in the tts_models table. + const WEIGHTS_TYPE : WeightsType = WeightsType::Tacotron2; + const WEIGHTS_CATEGORY : WeightsCategory = WeightsCategory::TextToSpeech; + + // NB: Not setting a few fields (for now) + // maybe_last_update_user_token - seems like bad design + // TODO(bt): file checksum + // TODO(bt): rename maybe_public_bucket_extension to maybe_public_bucket_suffix (!!!) + // TODO(bt): do we need model_weights.ip_address_last_update without audit logs? + // TODO(bt): rename creator_ip_address to ip_address_creation (and add ip_address_last_update) + // TODO(bt): Check model_weights column integer types - signed vs unsigned + let query = sqlx::query!( + r#" +INSERT INTO model_weights +SET + token = ?, + weights_type = ?, + weights_category = ?, + title = ?, + maybe_cover_image_media_file_token = NULL, + maybe_description_markdown = ?, + maybe_description_rendered_html = ?, + creator_user_token = ?, + creator_ip_address = ?, + + creator_set_visibility = ?, + maybe_last_update_user_token = NULL, + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + file_checksum_sha2 = ?, + + public_bucket_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + + maybe_migration_old_model_token = ?, + version = ?, + created_at = ?, + updated_at = ?, + user_deleted_at = ?, + mod_deleted_at = ? + +ON DUPLICATE KEY UPDATE + weights_type = ?, + weights_category = ?, + title = ?, + maybe_cover_image_media_file_token = NULL, + maybe_description_markdown = ?, + maybe_description_rendered_html = ?, + creator_user_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + maybe_last_update_user_token = NULL, + original_download_url = ?, + original_filename = ?, + file_size_bytes = ?, + file_checksum_sha2 = ?, + public_bucket_hash = ?, + maybe_public_bucket_prefix = ?, + maybe_public_bucket_extension = ?, + maybe_migration_old_model_token = ?, + version = ?, + created_at = ?, + updated_at = ?, + user_deleted_at = ?, + mod_deleted_at = ? + "#, + // Insert + model_weight_token, + WEIGHTS_TYPE, + WEIGHTS_CATEGORY, + record.title, + record.description_markdown, + record.description_rendered_html, + record.creator_user_token, + record.creator_ip_address_creation, + record.creator_set_visibility.to_str(), + record.original_download_url, + record.original_filename, + record.file_size_bytes, + + copied_data.file_sha_hash, + copied_data.bucket_path.get_object_hash(), + copied_data.bucket_path.get_optional_prefix(), + copied_data.bucket_path.get_optional_extension(), + + record.token.as_str(), + record.version, + record.created_at, + record.updated_at, + record.user_deleted_at, + record.mod_deleted_at, + + // Update + WEIGHTS_TYPE, + WEIGHTS_CATEGORY, + record.title, + record.description_markdown, + record.description_rendered_html, + record.creator_user_token, + record.creator_ip_address_creation, + record.creator_set_visibility.to_str(), + record.original_download_url, + record.original_filename, + record.file_size_bytes, + + copied_data.file_sha_hash, + copied_data.bucket_path.get_object_hash(), + copied_data.bucket_path.get_optional_prefix(), + copied_data.bucket_path.get_optional_extension(), + + record.token.as_str(), + record.version, + record.created_at, + record.updated_at, + record.user_deleted_at, + record.mod_deleted_at, + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} + +pub async fn upsert_model_weights_extension_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + transaction: &mut Transaction<'_, MySql> +) -> AnyhowResult<()> { + let query = sqlx::query!( + r#" +INSERT INTO model_weights_extension_tts_details +SET + model_weights_token = ?, + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + text_pipeline_type = ?, + use_default_mel_multiply_factor = ?, + maybe_custom_mel_multiply_factor = ?, + maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ? +ON DUPLICATE KEY UPDATE + ietf_language_tag = ?, + ietf_primary_language_subtag = ?, + text_pipeline_type = ?, + use_default_mel_multiply_factor = ?, + maybe_custom_mel_multiply_factor = ?, + maybe_default_pretrained_vocoder = ?, + maybe_custom_vocoder_token = ? + "#, + // Insert + &model_weight_token, + record.ietf_language_tag, + record.ietf_primary_language_subtag, + record.text_pipeline_type, + record.use_default_mel_multiply_factor, + record.maybe_custom_mel_multiply_factor, + record.maybe_default_pretrained_vocoder, + record.maybe_custom_vocoder_token, + + // Update + record.ietf_language_tag, + record.ietf_primary_language_subtag, + record.text_pipeline_type, + record.use_default_mel_multiply_factor, + record.maybe_custom_mel_multiply_factor, + record.maybe_default_pretrained_vocoder, + record.maybe_custom_vocoder_token + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} + +pub async fn update_original_record( + record: &WholeTtsModelRecord, + model_weight_token: &ModelWeightToken, + transaction: &mut Transaction<'_, MySql> +) -> AnyhowResult<()> { + let query = sqlx::query!( + r#" +UPDATE tts_models +SET + maybe_migration_new_model_weights_token = ? +WHERE token = ? + "#, + model_weight_token, + record.token, + ); + + let _r = query.execute(&mut **transaction).await?; + + Ok(()) +} diff --git a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs index 7f72eb2fd5..13a70c3072 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/tests.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/tests.rs @@ -293,6 +293,8 @@ mod tests { maybe_scoped_weight_type: None, maybe_scoped_weight_category: None, mysql_pool: &pool, + maybe_scoped_weight_type: None, + maybe_scoped_weight_category: None, }).await?; for weight in weights_by_username.records.iter() { // print weight diff --git a/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs b/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs index 3c0c393b82..c9ee839d9a 100644 --- a/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs +++ b/crates/lib/mysql_queries/src/queries/tts/tts_models/list_tts_models.rs @@ -12,7 +12,7 @@ use crate::helpers::boolean_converters::i8_to_bool; // FIXME: This is the old style of query scoping and shouldn't be copied. // TODO/FIXME : This struct is returned publicly in some endpoints! -#[derive(Serialize)] +#[derive(Serialize, Clone)] pub struct TtsModelRecordForList { pub model_token: String, pub tts_model_type: String, diff --git a/crates/lib/mysql_queries/src/queries/users/mod.rs b/crates/lib/mysql_queries/src/queries/users/mod.rs index 1b5e04bc70..e64f9df9c5 100644 --- a/crates/lib/mysql_queries/src/queries/users/mod.rs +++ b/crates/lib/mysql_queries/src/queries/users/mod.rs @@ -1,5 +1,6 @@ pub mod user; pub mod user_badges; +pub mod user_bookmarks; pub mod user_password_resets; pub mod user_profiles; pub mod user_ratings; diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/batch_get_user_bookmarks.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/batch_get_user_bookmarks.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/batch_get_user_bookmarks.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/batch_get_user_bookmarks.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/delete_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/delete_user_bookmark.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/delete_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/delete_user_bookmark.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs similarity index 91% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs index c366eab630..6e153a5219 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_total_bookmark_count_for_entity.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_total_bookmark_count_for_entity.rs @@ -1,11 +1,11 @@ use anyhow::anyhow; use log::warn; -use sqlx::{MySql}; +use sqlx::MySql; use sqlx::pool::PoolConnection; use errors::AnyhowResult; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct BookmarkCount { pub total_count: usize, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs similarity index 98% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs index b849cc655d..bbc1490f27 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/get_user_bookmark_transactional_locking.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/get_user_bookmark_transactional_locking.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct UserBookmark { pub token: UserBookmarkToken, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs similarity index 98% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs index 8dbe73c1dc..4d8a1b71ef 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks.rs @@ -6,8 +6,8 @@ use enums::by_table::model_weights::weights_types::WeightsType; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use errors::AnyhowResult; -use crate::queries::user_bookmarks::list_user_bookmarks_result::RawUserBookmarkRecord; -use crate::queries::user_bookmarks::list_user_bookmarks_result::UserBookmark; +use crate::queries::users::user_bookmarks::list_user_bookmarks_result::RawUserBookmarkRecord; +use crate::queries::users::user_bookmarks::list_user_bookmarks_result::UserBookmark; pub struct UserBookmarkListPage { pub results: Vec, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs similarity index 96% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs index 3df0f0812c..785f1ac802 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_for_entity.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_for_entity.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct UserBookmark { pub token: UserBookmarkToken, diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs similarity index 82% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs index 88c6954d45..2184350686 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/list_user_bookmarks_result.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/list_user_bookmarks_result.rs @@ -70,56 +70,56 @@ pub struct UserBookmark { } pub struct RawUserBookmarkRecord { - pub (crate) token: UserBookmarkToken, + pub(crate) token: UserBookmarkToken, - pub (crate) entity_type: UserBookmarkEntityType, - pub (crate) entity_token: String, + pub(crate) entity_type: UserBookmarkEntityType, + pub(crate) entity_token: String, // TODO(bt,2023-12-30): I don't think these user fields are necessary. // We already know the user we're querying on behalf of, so presumably we // already have their info (ie. profile page use). - pub (crate) user_token: UserToken, - pub (crate) username: String, - pub (crate) user_display_name: String, - pub (crate) user_gravatar_hash: String, + pub(crate) user_token: UserToken, + pub(crate) username: String, + pub(crate) user_display_name: String, + pub(crate) user_gravatar_hash: String, - pub (crate) created_at: DateTime, - pub (crate) updated_at: DateTime, - pub (crate) deleted_at: Option>, + pub(crate) created_at: DateTime, + pub(crate) updated_at: DateTime, + pub(crate) deleted_at: Option>, - pub (crate) maybe_ratings_positive_count: Option, - pub (crate) maybe_ratings_negative_count: Option, - pub (crate) maybe_bookmark_count: Option, + pub(crate) maybe_ratings_positive_count: Option, + pub(crate) maybe_ratings_negative_count: Option, + pub(crate) maybe_bookmark_count: Option, - pub (crate) maybe_media_file_type: Option, - pub (crate) maybe_media_file_origin_category: Option, - pub (crate) maybe_media_file_public_bucket_hash: Option, - pub (crate) maybe_media_file_public_bucket_prefix: Option, - pub (crate) maybe_media_file_public_bucket_extension: Option, + pub(crate) maybe_media_file_type: Option, + pub(crate) maybe_media_file_origin_category: Option, + pub(crate) maybe_media_file_public_bucket_hash: Option, + pub(crate) maybe_media_file_public_bucket_prefix: Option, + pub(crate) maybe_media_file_public_bucket_extension: Option, - pub (crate) maybe_media_file_creator_user_token: Option, - pub (crate) maybe_media_file_creator_username: Option, - pub (crate) maybe_media_file_creator_display_name: Option, - pub (crate) maybe_media_file_creator_gravatar_hash: Option, + pub(crate) maybe_media_file_creator_user_token: Option, + pub(crate) maybe_media_file_creator_username: Option, + pub(crate) maybe_media_file_creator_display_name: Option, + pub(crate) maybe_media_file_creator_gravatar_hash: Option, - pub (crate) maybe_model_weight_type: Option, - pub (crate) maybe_model_weight_category: Option, + pub(crate) maybe_model_weight_type: Option, + pub(crate) maybe_model_weight_category: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_hash: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_prefix: Option, - pub (crate) maybe_model_weight_cover_image_public_bucket_extension: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_hash: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_prefix: Option, + pub(crate) maybe_model_weight_cover_image_public_bucket_extension: Option, - pub (crate) maybe_model_weight_creator_user_token: Option, - pub (crate) maybe_model_weight_creator_username: Option, - pub (crate) maybe_model_weight_creator_display_name: Option, - pub (crate) maybe_model_weight_creator_gravatar_hash: Option, + pub(crate) maybe_model_weight_creator_user_token: Option, + pub(crate) maybe_model_weight_creator_username: Option, + pub(crate) maybe_model_weight_creator_display_name: Option, + pub(crate) maybe_model_weight_creator_gravatar_hash: Option, - pub (crate) maybe_descriptive_text_model_weight_title: Option, - pub (crate) maybe_descriptive_text_tts_model_title: Option, - pub (crate) maybe_descriptive_text_tts_result_inference_text: Option, - pub (crate) maybe_descriptive_text_user_display_name: Option, - pub (crate) maybe_descriptive_text_voice_conversion_model_title: Option, - pub (crate) maybe_descriptive_text_zs_voice_title: Option, + pub(crate) maybe_descriptive_text_model_weight_title: Option, + pub(crate) maybe_descriptive_text_tts_model_title: Option, + pub(crate) maybe_descriptive_text_tts_result_inference_text: Option, + pub(crate) maybe_descriptive_text_user_display_name: Option, + pub(crate) maybe_descriptive_text_voice_conversion_model_title: Option, + pub(crate) maybe_descriptive_text_zs_voice_title: Option, } impl RawUserBookmarkRecord { diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/mod.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/mod.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/mod.rs diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs similarity index 94% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs index ef361360a6..34062a5594 100644 --- a/crates/lib/mysql_queries/src/queries/user_bookmarks/upsert_user_bookmark.rs +++ b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/upsert_user_bookmark.rs @@ -7,7 +7,7 @@ use errors::AnyhowResult; use tokens::tokens::user_bookmarks::UserBookmarkToken; use tokens::tokens::users::UserToken; -use crate::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use crate::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; pub struct CreateUserBookmarkArgs<'e, 'c, E> where E: 'e + Executor<'c, Database = MySql> diff --git a/crates/lib/mysql_queries/src/queries/user_bookmarks/user_bookmark_entity_token.rs b/crates/lib/mysql_queries/src/queries/users/user_bookmarks/user_bookmark_entity_token.rs similarity index 100% rename from crates/lib/mysql_queries/src/queries/user_bookmarks/user_bookmark_entity_token.rs rename to crates/lib/mysql_queries/src/queries/users/user_bookmarks/user_bookmark_entity_token.rs diff --git a/crates/service/job/download_job/src/job_loop/process_single_job.rs b/crates/service/job/download_job/src/job_loop/process_single_job.rs index ba69531e70..4a133216d9 100644 --- a/crates/service/job/download_job/src/job_loop/process_single_job.rs +++ b/crates/service/job/download_job/src/job_loop/process_single_job.rs @@ -76,6 +76,8 @@ pub async fn do_process_single_job(job_state: &JobState, job: &AvailableDownload return Err(anyhow!("Bad download URL: `{}`", &job.download_url)); } + // TODO USE THIS DOWNLOADER + let download_filename = match job_state.sidecar_configs.google_drive_downloader.download_file(&job.download_url, &temp_dir).await { Ok(filename) => filename, Err(e) => { diff --git a/crates/service/job/download_job/src/job_state.rs b/crates/service/job/download_job/src/job_state.rs index 14b476a1f3..26b3e90a94 100644 --- a/crates/service/job/download_job/src/job_state.rs +++ b/crates/service/job/download_job/src/job_state.rs @@ -70,7 +70,7 @@ pub struct PretrainedModels { /// Configurations and interfaces to code deployed as sidecars or container mounts. pub struct SidecarConfigs { - pub google_drive_downloader: GoogleDriveDownloadCommand, + pub(crate) google_drive_downloader: GoogleDriveDownloadCommand, pub rvc_v2_model_check_command: RvcV2ModelCheckCommand, pub softvc_model_check_command: SoftVcModelCheckCommand, pub so_vits_svc_model_check_command: SoVitsSvcModelCheckCommand, diff --git a/crates/service/job/inference_job/Cargo.toml b/crates/service/job/inference_job/Cargo.toml index 19e17e422b..80897d083f 100644 --- a/crates/service/job/inference_job/Cargo.toml +++ b/crates/service/job/inference_job/Cargo.toml @@ -19,7 +19,9 @@ actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } +google_drive_common = { path = "../../../lib/api_clients/google_drive_common" } collections = { path = "../../../lib/collections" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } concurrency = { path = "../../../lib/concurrency" } config = { path = "../../../lib/deprecated/config" } easyenv = { path = "../../../lib/easyenv" } diff --git a/crates/service/job/inference_job/config/inference-job.development.env b/crates/service/job/inference_job/config/inference-job.development.env index a8b9b0ec47..34dd797375 100644 --- a/crates/service/job/inference_job/config/inference-job.development.env +++ b/crates/service/job/inference_job/config/inference-job.development.env @@ -103,6 +103,14 @@ VITS_INFERENCE_MAYBE_DOCKER_IMAGE_SHA="03017896367b" RERENDER_INFERENCE_ROOT_DIRECTORY="/home/salt/clone/storyteller-ml/animation/rerender-lora/" RERENDER_INFERENCE_EXECUTABLE_OR_COMMAND="venv/bin/python rerender.py" +# ================================= # +# # +# STABLE DIFFUSION # +# # +# ================================= # +STABLE_DIFFUSION_INFERENCE_ROOT_DIRECTORY="/home/tensor/code/storyteller/storyteller-ml/image_generation/Storyteller-Studio-Filter/" +STABLE_DIFFUSION_INFERENCE_COMMAND=".venv/bin/python main.py" + # ================================= # # # # MocapNET # diff --git a/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs b/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs index 9a79c77746..57e99ae5b5 100644 --- a/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs +++ b/crates/service/job/inference_job/src/job/job_loop/determine_dependency_status.rs @@ -188,7 +188,7 @@ async fn get_model_record_from_cacheable_query(job_dependencies: &JobDependencie } } _ => { - warn!("Job does not have a category or a model token: {:?} - category: {:?}, token: {:?}", + warn!("Only applies to older job workflows. Job does not have a category or a model token: {:?} - category: {:?}, token: {:?}", job.id, job.inference_category, job.maybe_model_token); MaybeInferenceModel::None diff --git a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs index 0cc79420f1..0f84ea6100 100644 --- a/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs +++ b/crates/service/job/inference_job/src/job/job_loop/process_single_job.rs @@ -17,6 +17,7 @@ use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_l use crate::job::job_types::tts::process_single_tts_job::process_single_tts_job; use crate::job::job_types::vc::process_single_vc_job::process_single_vc_job; use crate::job::job_types::videofilter::process_single_vf_job::process_single_vf_job; +use crate::job::job_types::image_generation::process_single_ig_job::process_single_ig_job; use crate::job::job_types::mocap::process_single_mc_job::process_single_mc_job; use crate::job::job_types::workflow::process_single_wf_job::process_single_wf_job; use crate::job_dependencies::JobDependencies; @@ -188,6 +189,9 @@ async fn do_process_single_job( InferenceCategory::VideoFilter => { process_single_vf_job(job_dependencies, job).await? } + InferenceCategory::ImageGeneration => { + process_single_ig_job(job_dependencies, job).await? + } InferenceCategory::Mocap => { process_single_mc_job(job_dependencies, job).await? } diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs new file mode 100644 index 0000000000..fa1c0e36ec --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs @@ -0,0 +1,2 @@ +pub mod process_single_ig_job; +pub mod sd; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs new file mode 100644 index 0000000000..bb26c58933 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs @@ -0,0 +1,34 @@ +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use errors::anyhow; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::image_generation::sd::process_job::process_job_selection; +use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; +use crate::job_dependencies::JobDependencies; + +pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { + match job.maybe_model_type { + Some(InferenceModelType::StableDiffusion) => { + dispatch_sd_job( + job_dependencies, + job + ).await + }, + Some(other_model_type) => { + Err(ProcessSingleJobError::Other(anyhow!("Wrong model type for SD: {:?}", other_model_type))) + } + None => { + Err(ProcessSingleJobError::Other(anyhow!("SD model type not set"))) + } + } +} + +pub async fn dispatch_sd_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result { + let args = StableDiffusionProcessArgs { + job_dependencies, + job, + }; + process_job_selection(args).await +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs new file mode 100644 index 0000000000..bf85551a02 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/mod.rs @@ -0,0 +1,6 @@ +pub mod process_job; + +pub mod model_downloaders; +pub mod sd_inference_command; +pub mod stable_diffusion_dependencies; +pub mod validate_inputs; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs new file mode 100644 index 0000000000..0dd95e7cfb --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/model_downloaders.rs @@ -0,0 +1,20 @@ + +crate::impl_model_downloader!( + SDVAEEncoder, + "SD Variation Auto Encoder", + "SD_VARIATIONAL_AUTO_ENCODER_BUCKET_PATH", + "/animation_sadtalker/enhancer/parsing_parsenet.pth", + "SD_VARIATIONAL_AUTO_ENCODER_FILESYSTEM_PATH", + "/tmp/downloads/sadtalker/parsing_parsenet.pth" + ); + +pub struct StableDiffusionDownloaders { + sd_vae_encoder: SDVAEEncoder +} +impl StableDiffusionDownloaders { + pub fn build_all_from_env() -> Self { + Self { + sd_vae_encoder: SDVAEEncoder::from_env(), + } + } +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs new file mode 100644 index 0000000000..b7ea33cf16 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs @@ -0,0 +1,477 @@ +use std::thread; +use std::time::{Duration, Instant}; + +use anyhow::anyhow; +use log::info; + +use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor; +use cloud_storage::remote_file_manager::remote_cloud_bucket_details::RemoteCloudBucketDetails; +use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient; +use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; +use filesys::path_to_string::path_to_string; +use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType}; +use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs}; +use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs; +use crate::job_dependencies::JobDependencies; + +pub struct StableDiffusionProcessArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} + + +// run inference +// insert record into the db with the inference job token complete. + +pub async fn sd_args_from_job( + args: &StableDiffusionProcessArgs<'_> +) -> Result { + let inference_args = args.job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + let polymorphic_args = match inference_args { + Some(args) => args, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!")) + ); + } + }; + let sd_args = match polymorphic_args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); + } + }; + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); + Ok(stable_diffusion_args) +} + +// store the prompt and cluster them today. + +pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_> +) -> Result { + let sd_args = sd_args_from_job(&args).await?; + if sd_args.type_of_inference == "inference" { + process_job_inference(&args).await + } + else if sd_args.type_of_inference == "lora" { + process_job_lora(&args).await + } + else if sd_args.type_of_inference == "checkpoint"{ + process_job_inference(&args).await + } + else { + Err(ProcessSingleJobError::Other(anyhow!("inference type doesn't exist!"))) + } +} + +pub async fn process_job_checkpoint(_args: &StableDiffusionProcessArgs<'_>) -> Result { + + // insert model record in create model record + + // let maybe_result_entity = ResultEntity { + // entity_type: InferenceResultType::UploadModel, + // entity_token: + // }; + Ok(JobSuccessResult { + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) +} +pub async fn process_job_lora(_args: &StableDiffusionProcessArgs<'_>) -> Result { + Ok(JobSuccessResult { + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) +} + +pub async fn process_job_inference( + args: &StableDiffusionProcessArgs<'_> +) -> Result { + let job = args.job; + let deps = args.job_dependencies; + let mysql_pool = &deps.db.mysql_pool; + + + let sd_args = sd_args_from_job(&args).await?; + let sd_deps = match &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies { + None => { + return Err(ProcessSingleJobError::Other(anyhow!("Missing Job Specific Dependencies"))); + } + Some(val) => { + val + } + }; + + let creator_user_token:UserToken; + match &job.maybe_creator_user_token { + Some(token) => { + creator_user_token = UserToken::new_from_str(token); + }, + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + } + + + let _job_progress_reporter = args.job_dependencies.clients.job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + //==================== TEMP DIR ==================== // + + let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0); + + // thread::sleep(seconds) to check the directory + + //NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let mut work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + + let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors"); + let lora_path = work_temp_dir.path().join("lora.safetensors"); + let vae_path = work_temp_dir.path().join("vae.safetensors"); + let output_path = work_temp_dir.path().join("output"); + let g_drive_path = work_temp_dir.path().join("gdrive"); + + info!("Paths to download to:"); + info!("sd_checkpoint_path:{}",sd_checkpoint_path.display()); + info!("lora_path:{}",lora_path.display()); + info!("vae_path:{}",vae_path.display()); + info!("output_path:{}",output_path.display()); + info!("tmp_google_drive_path:{}",g_drive_path.display()); + + // If it is not inference then it must be loRA or checkpoint downloading + if sd_args.type_of_inference != "inference" { + let mut download_url= String::from(""); + + if sd_args.type_of_inference == "lora" { + download_url = match sd_args.maybe_lora_upload_path { + Some(val) => { val }, + None => { "".to_string() } + }; + } + else if sd_args.type_of_inference == "checkpoint" { + download_url = match sd_args.maybe_upload_path { + Some(val) => { val }, + None => { "".to_string() } + } + } + + if download_url.len() == 0 { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing"))); + } + + let file_name = "download.safetensors"; + let download_script = easyenv::get_env_string_or_default( + "DOWNLOAD_SCRIPT", + "./scripts/download_internet_file.py"); + let google_drive_downloader = + GoogleDriveDownloadCommand::new_production(&download_script); + info!("Downloading {}",download_url); + + let _download_filename = match google_drive_downloader.download_file_with_file_name(&download_url, &work_temp_dir,file_name).await { + Ok(filename) => filename, + Err(_e) => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download"))); + } + }; + let download_file_path = work_temp_dir.path().join(file_name); + info!("File Retrieved at {}",download_file_path.display()); + } + + // // Unpack loRA and Checkpoint + // // run inference by downloading from google drive. + let lora_token = sd_args.maybe_lora_model_token; + let weight_token = sd_args.maybe_sd_model_token; + + let retrieved_sd_record = match weight_token { + Some(token) => { + let retrieved_sd_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + match retrieved_sd_record { + Some(record) => record, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no record of model!")) + ); + } + } + } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no sd model token for job!")) + ); + } + }; + + // ignore if no lora token + let retrieved_loRA_record = match lora_token { + Some(token) => { + let retrieved_loRA_record = get_weight_by_token( + &token, + false, + &deps.db.mysql_pool + ).await?; + Some(retrieved_loRA_record) + } + None => None, + }; + + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?; + + // Details for SD checkpoint + let details = RemoteCloudBucketDetails::new( + retrieved_sd_record.public_bucket_hash.clone(), + retrieved_sd_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()), + retrieved_sd_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string()) + ); + remote_cloud_file_client.download_file(details, path_to_string(sd_checkpoint_path.clone())).await?; + + match retrieved_loRA_record { + Some(record) => { + match record { + Some(model_weight_record) => { + + let lora_details = RemoteCloudBucketDetails::new( + model_weight_record.public_bucket_hash.clone(), + model_weight_record.maybe_public_bucket_prefix + .clone() + .unwrap_or_else(|| "".to_string()), + model_weight_record.maybe_public_bucket_extension + .clone() + .unwrap_or_else(|| "".to_string()) + ); + remote_cloud_file_client.download_file( + lora_details, + path_to_string(lora_path.clone()) + ).await?; + + } + None => {} + } + } + None => {} + } + + // VAE token for now + let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw"); + let model_weight_vae = ModelWeightToken(vae_token); + + let vae_weight_record = get_weight_by_token( + &model_weight_vae, + false, + &deps.db.mysql_pool + ).await?; + + let vae_weight_record = match vae_weight_record { + Some(val) => val, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))) + } + }; + + + let vae_details = RemoteCloudBucketDetails::new( + vae_weight_record.public_bucket_hash.clone(), + vae_weight_record.maybe_public_bucket_prefix + .clone() + .unwrap_or_else(|| "".to_string()), + vae_weight_record.maybe_public_bucket_extension + .clone() + .unwrap_or_else(|| "".to_string()) + ); + + remote_cloud_file_client.download_file( + vae_details, + path_to_string(vae_path.clone()) + ).await?; + + let prompt = match sd_args.maybe_prompt { + Some(val) => val, + None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!"))) + }; + + let stderr_output_file = work_temp_dir.path().join("sd_err.txt"); + + let number_of_samples = match sd_args.maybe_number_of_samples { + Some(val) => val, + None => 20, + }; + + let inference_start_time = Instant::now(); + + sd_deps.inference_command.execute_inference(InferenceArgs { + work_dir: work_temp_dir.path().to_path_buf(), + output_file: output_path.clone(), + stderr_output_file, + prompt, + negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(), + number_of_samples, + samplers: sd_args.maybe_sampler.unwrap_or(String::from("Euler a")), + width: sd_args.maybe_width.unwrap_or(512), + height: sd_args.maybe_height.unwrap_or(512), + cfg_scale: sd_args.maybe_cfg_scale.unwrap_or(7), + seed: sd_args.maybe_seed.unwrap_or(1), + lora_path: lora_path.clone(), + checkpoint_path: sd_checkpoint_path.clone(), + vae: vae_path.clone(), + batch_count: sd_args.maybe_batch_count.unwrap_or(1), + }); + + let inference_duration = Instant::now().duration_since(inference_start_time); + let creator_ip_address = &job.creator_ip_address; + // run a for loop for output images output_0 in the folder then use upload media. + // pngs + + if sd_args.type_of_inference == "inference" { + // gather all the images into a single batch + + + for i in 0..sd_args.maybe_batch_count.unwrap_or(1) { + let path = output_path.clone(); + let file_path = format!("{}_{}.png", path_to_string(path), i); + println!("Upload File Path:{}", file_path); + let metadata = remote_cloud_file_client.upload_file(Box::new(MediaImagePngDescriptor {}), file_path.as_ref()).await?; + + let bucket_details = match metadata.bucket_details { + Some(val) => { val } + None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))) + }; + + // TODO: Generic Insert with batch id. + // insert media file generic should run here. + insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs { + pool: mysql_pool, + maybe_creator_user_token: Some(&creator_user_token), + maybe_creator_anonymous_visitor_token: None, + creator_ip_address, + creator_set_visibility: Default::default(), + upload_type: UploadType::Filesystem, + media_file_type: MediaFileType::Image, + maybe_mime_type: Some(metadata.mimetype.as_ref()), + file_size_bytes: metadata.file_size_bytes, + duration_millis: inference_duration.as_secs() * 1000, + sha256_checksum: metadata.sha256_checksum.as_str(), + public_bucket_directory_hash: bucket_details.object_hash.as_str(), + maybe_public_bucket_prefix: Some(bucket_details.prefix.as_str()), + maybe_public_bucket_extension: Some(bucket_details.suffix.as_str()), + }).await?; + } + + // hack to check the directory before clean up. + let thirtyMinutes = 1800; + thread::sleep(Duration::from_secs(thirtyMinutes)); + // upload media and create a record. + + Ok(JobSuccessResult {// TODO: batch token needs to go here + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) + } + else { // must be an upload and need upload token + + + // The parameters will be updated on another screen perhaps? + // so right now it will fill with the availible values. + create_weight(CreateModelWeightsArgs { + token: &ModelWeightToken::generate(), + weights_type: WeightsType::StableDiffusion15, + weights_category: WeightsCategory::ImageGeneration, + title: "".to_string(), + maybe_description_markdown: None, + maybe_description_rendered_html: None, + creator_user_token: Some(&creator_user_token), + creator_ip_address, + creator_set_visibility: Default::default(), + maybe_last_update_user_token: None, + original_download_url: None, + original_filename: None, + file_size_bytes: 0, + file_checksum_sha2: "".to_string(), + public_bucket_hash: "".to_string(), + maybe_public_bucket_prefix: None, + maybe_public_bucket_extension: None, + version: 0, + mysql_pool, + }).await?; + + Ok(JobSuccessResult {// TODO: batch token needs to go here + maybe_result_entity: None, + inference_duration: Duration::from_secs(0), + }) + } + + + +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use anyhow::anyhow; + + use cloud_storage::remote_file_manager::{ + remote_cloud_bucket_details::RemoteCloudBucketDetails, + remote_cloud_file_manager::RemoteCloudFileClient, + }; + use errors::AnyhowResult; + + #[tokio::test] + async fn test_seed_weights_files() -> AnyhowResult<()> { + let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data"); + let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await; + let remote_cloud_file_client = match remote_cloud_file_client { + Ok(res) => { + res + } + Err(_) => { + return Err(anyhow!("failed to get remote cloud file client")); + } + }; + + let mut path_dl1 = seed_path.clone(); + path_dl1.push("downloads/loRA"); + let mut path_dl2 = seed_path.clone(); + path_dl2.push("downloads/checkpoint"); + + let bucket_details1 = RemoteCloudBucketDetails { + object_hash: String::from("apa0ej6es8d3ss2gwtf1cghge35qn9tn"), + prefix: String::from("sd15"), + suffix: String::from("safetensors"), + }; + + let bucket_details2 = RemoteCloudBucketDetails { + object_hash: String::from("27kz11et18fargyyxbj66ntfn621k9d3"), + prefix: String::from("loRA"), + suffix: String::from("safetensors"), + }; + + remote_cloud_file_client.download_file(bucket_details1, String::from("./checkpoint")).await?; + remote_cloud_file_client.download_file(bucket_details2, String::from("./loRA")).await?; + + Ok(()) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs new file mode 100644 index 0000000000..265b35f8f7 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/sd_inference_command.rs @@ -0,0 +1,309 @@ + +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS : Lazy> = Lazy::new(|| { + let env_var_names= [ + "MYSQL_URL", + "ACCESS_KEY", + "SECRET_KEY", + "NEWRELIC_API_KEY", + ]; + + env_var_names.iter() + .map(|value| value.to_string()) + .collect::>() +}); + +#[derive(Clone)] +pub struct StableDiffusionInferenceCommand { + /// Where the code lives + stable_diffusion_root_code_directory: PathBuf, + + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + + /// Optional default config file to use + maybe_default_config_path: Option, + + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + + /// Eg. `python3 inference.py` + Command(String), +} + +pub struct InferenceArgs { + /// --source_image: path to the input image (or video) + /// --result_dir: path to directory work is performed + pub work_dir: PathBuf, + /// --result_file: path to final file output + pub output_file: PathBuf, + pub stderr_output_file: PathBuf, + pub prompt: String, + pub negative_prompt:String, + pub number_of_samples:i32, + pub samplers:String, + pub width:i32, + pub height:i32, + pub cfg_scale:i32, + pub seed:i64, + pub lora_path:PathBuf, + pub checkpoint_path:PathBuf, + pub vae:PathBuf, + pub batch_count:i32, +} + +impl StableDiffusionInferenceCommand { + pub fn new>( + stable_diffusion_root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + stable_diffusion_root_code_directory: stable_diffusion_root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| s.to_string()), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let stable_diffusion_root_code_directory = easyenv::get_env_pathbuf_required( + "STABLE_DIFFUSION_INFERENCE_ROOT_DIRECTORY")?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_COMMAND"); + + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_VENV_COMMAND"); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH"); + + let maybe_execution_timeout: Option = + easyenv::get_env_duration_seconds_optional("STABLE_DIFFUSION_TIMEOUT_SECONDS"); + + let maybe_docker_options = easyenv::get_env_string_optional( + "STABLE_DIFFUSION_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STABLE_DIFFUSION_ALTERNATE_CHECKPOINT_PATH"); + + Ok(Self { + stable_diffusion_root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference( + &self, + args: InferenceArgs, + ) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference( + &self, + args: InferenceArgs, + ) -> AnyhowResult { + + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.stable_diffusion_root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + command.push_str(format!("--prompt \"{}\"",&path_to_string(args.prompt)).as_ref()); + + command.push_str(" --output-dir "); + command.push_str(&path_to_string(args.work_dir)); + + command.push_str(" --output-name "); + command.push_str(&path_to_string(args.output_file)); + + // command.push_str(" --stderr-output-file "); + // command.push_str(&path_to_string(args.stderr_output_file.clone())); + + command.push_str(format!(" --negative-prompt \"{}\"",&path_to_string(args.negative_prompt)).as_ref()); + //command.push_str(&path_to_string(args.negative_prompt)); + + command.push_str(" --number-of-samples "); + command.push_str(args.number_of_samples.to_string().as_str()); + + command.push_str(format!(" --samplers \"{}\"",&path_to_string(args.samplers)).as_ref()); + //command.push_str(&path_to_string(args.samplers)); + + command.push_str(" --width "); + command.push_str(args.width.to_string().as_str()); + + command.push_str(" --height "); + command.push_str(args.height.to_string().as_str()); + + command.push_str(" --cfg-scale "); + command.push_str(args.cfg_scale.to_string().as_str()); + + command.push_str(" --seed "); + command.push_str(args.seed.to_string().as_str()); + + command.push_str(" --loRA-path "); + command.push_str(&path_to_string(args.lora_path)); + + command.push_str(" --check-point "); + command.push_str(&path_to_string(args.checkpoint_path)); + + command.push_str(" --vae "); + command.push_str(&path_to_string(args.vae)); + + command.push_str(" --batch-count "); + command.push_str(args.batch_count.to_string().as_str()); + + command.push_str(" --batch-size "); + command.push_str(1.to_string().as_str()); + + + info!("Command: {:?}", command); + + let command_parts = [ + "bash", + "-c", + &command + ]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push(( + OsString::from(env_key), + OsString::from(env_value), + )); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {}", path_to_string(args.stderr_output_file.clone())); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!("Subprocess didn't end after timeout: {:?}; terminating...", &timeout); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs new file mode 100644 index 0000000000..bf14dfbdc8 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/stable_diffusion_dependencies.rs @@ -0,0 +1,19 @@ +use anyhow::{Ok}; +use errors::AnyhowResult; + +use crate::job::job_types::image_generation::sd::model_downloaders::StableDiffusionDownloaders; +use crate::job::job_types::image_generation::sd::sd_inference_command::StableDiffusionInferenceCommand; + +pub struct StableDiffusionDependencies { + pub downloaders: StableDiffusionDownloaders, + pub inference_command: StableDiffusionInferenceCommand, +} + +impl StableDiffusionDependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + downloaders: StableDiffusionDownloaders::build_all_from_env(), + inference_command: StableDiffusionInferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs new file mode 100644 index 0000000000..75cf105f6f --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/image_generation/sd/validate_inputs.rs @@ -0,0 +1,42 @@ +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use anyhow::anyhow; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs; +use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs; + + +pub async fn validate_inputs(args: StableDiffusionProcessArgs<'_>) -> Result<(), ProcessSingleJobError> { + let inference_args = args.job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let polymorphic_args = match inference_args { + Some(args) => args, + None => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!"))); + } + }; + + let sd_args = match polymorphic_args { + PolymorphicInferenceArgs::Ig(args) => args, + _ => { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))); + } + }; + + let stable_diffusion_args: StableDiffusionArgs = StableDiffusionArgs::from(sd_args.clone()); + if stable_diffusion_args.type_of_inference == "checkpoint" { + + } else if stable_diffusion_args.type_of_inference == "lora" { + + } else if stable_diffusion_args.type_of_inference == "inference" { + + } else { + return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inference type for job!"))); + } + + + + Ok(()) +} diff --git a/crates/service/job/inference_job/src/job/job_types/mod.rs b/crates/service/job/inference_job/src/job/job_types/mod.rs index 88f6e673de..64b888e0d8 100644 --- a/crates/service/job/inference_job/src/job/job_types/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/mod.rs @@ -2,5 +2,6 @@ pub mod lipsync; pub mod tts; pub mod vc; pub(crate) mod videofilter; +pub mod image_generation; pub(crate) mod mocap; -pub(crate) mod workflow; \ No newline at end of file +pub(crate) mod workflow; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs index 26d581e703..137d21bf61 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/mod.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/mod.rs @@ -1,4 +1,5 @@ pub mod process_single_tts_job; +pub mod styletts2; pub mod tacotron2_v2_early_fakeyou; pub mod vits; pub mod vall_e_x; diff --git a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs index d9a6cf1f53..495901d1dc 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/process_single_tts_job.rs @@ -8,7 +8,8 @@ use mysql_queries::queries::tts::tts_models::get_tts_model_for_inference_improve use crate::job::job_loop::job_success_result::JobSuccessResult; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; -use crate::job::job_types::tts::{tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::{styletts2, tacotron2_v2_early_fakeyou, vits}; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; use crate::job::job_types::tts::tacotron2_v2_early_fakeyou::process_job::ProcessJobArgs; use crate::job::job_types::tts::vall_e_x::process_job::VALLEXProcessJobArgs; use crate::job::job_types::tts::vits::process_job::VitsProcessJobArgs; @@ -24,7 +25,7 @@ pub async fn process_single_tts_job( // TODO: Move common checks for slurs, etc. here. match job.maybe_model_type { - Some(InferenceModelType::VallEX) => { + Some(InferenceModelType::VallEX) | Some(InferenceModelType::StyleTTS2) => { // Zero-shot TTS does not need a fine-tuned model token. dispatch_zero_shot_model( job_dependencies, @@ -64,6 +65,12 @@ async fn dispatch_zero_shot_model( job, }).await } + Some(InferenceModelType::StyleTTS2) => { + styletts2::process_job::process_job(StyleTTS2ProcessJobArgs { + job_dependencies, + job, + }).await + } _ => { Err(ProcessSingleJobError::Other(anyhow!("not a zero-shot model"))) } diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs new file mode 100644 index 0000000000..5623a593ff --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/mod.rs @@ -0,0 +1,7 @@ +pub mod model_downloaders; +pub mod process_job; +pub mod process_job_create_voice; +pub mod process_job_inference_voice; +pub mod validate_job; +pub mod styletts2_dependencies; +pub mod styletts2_inference_command; \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs new file mode 100644 index 0000000000..fd32b2b2f8 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/model_downloaders.rs @@ -0,0 +1,19 @@ +use crate::util::model_downloader::ModelDownloader; + +pub struct StyleTTS2Downloaders; +//{ +// pub styletts2: StyleTTS2ModelMapping, +//} + +impl StyleTTS2Downloaders { + // TODO(KS): All checkpoints are in the container, so we don't need to download them here. + // Maybe they should be here instead + pub fn build_all_from_env() -> Self { + Self { + } + } + + pub fn all_downloaders(&self) -> Vec<&dyn ModelDownloader> { + vec![] + } +} \ No newline at end of file diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs new file mode 100644 index 0000000000..e98c004db0 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job.rs @@ -0,0 +1,45 @@ +use anyhow::anyhow; + +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job_create_voice::process_create_voice; +use crate::job::job_types::tts::styletts2::process_job_inference_voice::process_inference_voice; +use crate::job::job_types::tts::styletts2::validate_job::validate_job; +use crate::job_dependencies::JobDependencies; + +use super::validate_job::JobType; + +// This will download everything get into the root host OS then ... will invoke inference using the pathes from the files invoked +pub struct StyleTTS2ProcessJobArgs<'a> { + pub job_dependencies: &'a JobDependencies, + pub job: &'a AvailableInferenceJob, +} + +// query using the token then grab the bucket hash +pub async fn process_job( + args: StyleTTS2ProcessJobArgs<'_> +) -> Result { + let job = args.job; + + // get args token + let jobArgs = validate_job(&job)?; // bubbles error up + + match jobArgs.job_type { + JobType::Create => { + if let Some(voice_dataset_token) = jobArgs.voice_dataset_token { + process_create_voice(args, voice_dataset_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Dataset Token?"))) + } + } + JobType::Inference => { + if let Some(voice_token) = jobArgs.voice_token { + process_inference_voice(args, voice_token).await + } else { + Err(ProcessSingleJobError::Other(anyhow!("Missing Voice Token?"))) + } + } + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs new file mode 100644 index 0000000000..92b36939f8 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_create_voice.rs @@ -0,0 +1,349 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::zs_voices::encoding_type::ZsVoiceEncodingType; +use enums::by_table::zs_voices::model_category::ZsVoiceModelCategory; +use enums::by_table::zs_voices::model_type::ZsVoiceModelType; +use enums::common::visibility::Visibility; +use mysql_queries::queries::voice_designer::datasets::get_dataset::get_dataset_by_token; +use mysql_queries::queries::voice_designer::datasets::get_dataset::ZsDataset; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::DatasetSampleRecordForList; +use mysql_queries::queries::voice_designer::voice_samples::list_dataset_samples_for_dataset_token::list_dataset_samples_for_dataset_token; +use mysql_queries::queries::voice_designer::voices::create_voice::create_voice; +use mysql_queries::queries::voice_designer::voices::create_voice::CreateVoiceArgs; +use tokens::tokens::users::UserToken; + +use crate::job; +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; + +const BUCKET_FILE_PREFIX_CREATE: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION_CREATE: &str = ".bin"; +const MIME_TYPE_CREATE: &str = "application/x-binary"; + +/// We'll increment this if we change the underlying StyleTTS2 code in an incompatible way. +const CURRENT_STYLE_TTS2_MODEL_VERSION: u64 = 0; + +pub async fn process_create_voice( + args: StyleTTS2ProcessJobArgs<'_>, + dataset_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps.clients.job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + info!("token! {}", dataset_token); + let voice_dataset_token = tokens::tokens::zs_voice_datasets::ZsVoiceDatasetToken(dataset_token); + + let creator_ip_address = &job.creator_ip_address; + + let creator_user_token = match &job.maybe_creator_user_token { + Some(token) => { + UserToken::new_from_str(token) + }, + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token"))); + } + }; + + // STEP 1. SETUP A TEMP DIRECTORY + let work_temp_dir = format!("/tmp/temp_zeroshot_create_voice_{}", job.id.0); + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let workdir = work_temp_dir.path().to_path_buf(); + + // STEP 2. Get dataset for the title for the voice + let voice_dataset = get_dataset_by_token(&voice_dataset_token, false, &mysql_pool).await; + let single_dataset: ZsDataset; + match voice_dataset { + Ok(val) => { + match val { + Some(val) => { + single_dataset = val; + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Dataset"))); + } + } + } + Err(e) => { + return Err(ProcessSingleJobError::from_anyhow_error(e)); + } + } + info!("Title:{} Token:{}", &single_dataset.title, &single_dataset.token); + + // STEP 3. Download dataset each audio file + let result = list_dataset_samples_for_dataset_token( + &voice_dataset_token, + false, + &mysql_pool + ).await; + + let dataset: Vec = match result { + Ok(val) => val, + Err(e) => return Err(ProcessSingleJobError::from_anyhow_error(e)), + }; + + info!("Dataset length info: {}", dataset.len()); + + let mut downloaded_dataset: Vec = Vec::new(); + + for (index, record) in dataset.iter().enumerate() { + //https://storage.googleapis.com/dev-vocodes-public/media/5/3/3/w/8/533w8zs0fy11nv7gkcna7p7vt03h8nda/dev_zs_533w8zs0fy11nv7gkcna7p7vt03h8nda.bin <-- where the file actually is + let prefix: Option<&str> = record.maybe_public_bucket_prefix.as_ref().map(|s| s.as_str()); + let extension: Option<&str> = record.maybe_public_bucket_extension.as_ref().map(|s| s.as_str()); + + info!("Record=> hash:{} prefix:{:?} extension:{:?}", + record.public_bucket_directory_hash, + prefix, + extension + ); + + let audio_media_file = MediaFileBucketPath::from_object_hash( + &record.public_bucket_directory_hash, + prefix, + extension + ); + + info!("Download using audio_media_file_path: {:?}", audio_media_file.to_full_object_pathbuf()); + + let file_name_wav = format!("{}.wav", index); + + let mut file_path = PathBuf::new(); + file_path.push(workdir.clone()); + file_path.push(file_path.clone()); + file_path.push(file_name_wav); + + info!("Downloading to path: {:?}", file_path); + + // TODO: we might want to catch the error and not include the pathes into download dataset? + let result = deps.buckets.public_bucket_client.download_file_to_disk( + audio_media_file.to_full_object_pathbuf(), + &file_path + ).await; + + if let Err(err) = result { + error!("could not download sample: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + info!("FilePath to clone voice: {}", file_path.to_string_lossy()); + downloaded_dataset.push(file_path.clone()); + + + // TODO(KS): This is a hack to get only the first file to download + // Ideally we could have improved logic to pick the "best" media file from a dataset + if downloaded_dataset.len() >= 1 { + break; + } + } + + info!("Dataset Length downloaded: {}", downloaded_dataset.len()); + + // STEP 4 Download the models + info!("Download models (if not present)..."); + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // Command line arg for a list of paths to insert the container + let audio_files = join_paths(downloaded_dataset); + + info!("Files to process: {:?}", audio_files); + + // Name of the output file + let output_embedding_path = work_temp_dir.path().join("style.npz"); + let stderr_output_file = work_temp_dir.path().join("zero_shot_create_voice_err.txt"); + + let inference_start_time = Instant::now(); + + // Run Inference + let command_exit_status = + model_dependencies.create_embedding_command.execute_inference( + job::job_types::tts::styletts2::styletts2_inference_command::CreateVoiceInferenceArgs { + output_embedding_path: &output_embedding_path, + audio_file: audio_files, + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Create Embedding Failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + if let Ok(contents) = read_to_string(&stderr_output_file) { + warn!("Captured stderr output: {}", contents); + } + + //thread::sleep(Duration::from_secs(300)); + + // Clean up temp files + //safe_delete_temp_file(&audio_path.filesystem_path); + //safe_delete_temp_file(&image_path.filesystem_path); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + info!("Inference success!"); + + info!("Uploading Media ..."); + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::generate_new( + ModelCategory::Tts, + ModelType::StyleTTS2, + CURRENT_STYLE_TTS2_MODEL_VERSION + ); + + let embedding_bucket_object_pathbuf = embedding_bucket_location.to_full_object_pathbuf(); + + // Get Finished File + let mut finished_file = work_temp_dir.path().to_path_buf(); + + let output_bucket_file_name = String::from("style.npz"); // use extension for bucket upload. + finished_file.push(&output_bucket_file_name); + + info!("Upload File Path: {:?}", finished_file); + info!("Upload Bucket Path: {:?}", embedding_bucket_object_pathbuf); + + args.job_dependencies.buckets.private_bucket_client + .upload_filename_with_content_type( + &embedding_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE_CREATE + ) + .await // TODO: We should check the mimetype to make sure bad payloads can't get uploaded + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // CLEARIFY! these items + // 1.Should this be object hash? + // 2.As well as this what should the voice name be? + let voice_name = single_dataset.title; + + // insert record + let voice_token = create_voice(CreateVoiceArgs { + dataset_token: &voice_dataset_token, + model_category: ZsVoiceModelCategory::Tts, + model_type: ZsVoiceModelType::StyleTTS2, + model_version: CURRENT_STYLE_TTS2_MODEL_VERSION, + model_encoding_type: ZsVoiceEncodingType::Encodec, + voice_title: &voice_name, + bucket_hash: embedding_bucket_location.get_object_hash().clone(), + maybe_creator_user_token: Some(&creator_user_token), + creator_ip_address: &creator_ip_address, + creator_set_visibility: Visibility::Public, + mysql_pool, + }).await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::ZeroShotVoiceEmbedding, + entity_token: voice_token.to_string(), + }), + inference_duration, + }) +} + +pub struct AudioFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_audio_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download audio file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download audio file: {err}")) + })?; + + let audio_file = AudioFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(audio_file) +} + +fn join_paths(paths: Vec) -> String { + paths + .into_iter() + .map(|p| format!("\"{}\"", p.display())) + .collect::>() + .join(" ") +} + +#[cfg(test)] +mod tests { + use std::path::PathBuf; + + use crate::job::job_types::tts::styletts2::process_job_create_voice::join_paths; + + #[test] + fn test_path_build() { + let paths = vec![ + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav"), + PathBuf::from("/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav") + ]; + + let value = join_paths(paths); + let expected = + "\"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/20.wav\" \"/home/tensor/code/TTSDockerContainer/Vall-E-mount/input/21.wav\""; + assert_eq!(value, expected); + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs new file mode 100644 index 0000000000..199181d6f1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/process_job_inference_voice.rs @@ -0,0 +1,304 @@ +use std::fs::read_to_string; +use std::path::PathBuf; +use std::time::Instant; + +use anyhow::anyhow; +use log::{error, info, warn}; + +use buckets::private::zs_voices::bucket_directory::{ModelCategory, ModelType}; +use buckets::private::zs_voices::bucket_file_path::ZeroShotVoiceEmbeddingBucketPath; +use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; +use cloud_storage::bucket_client::BucketClient; +use cloud_storage::bucket_path_unifier::BucketPathUnifier; +use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use filesys::file_size::file_size; +use hashing::sha256::sha256_hash_file::sha256_hash_file; +use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; +use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::InsertArgs; +use mysql_queries::queries::voice_designer::voices::get_voice::{get_voice_by_token, ZsVoice}; +use tts_common::clean_symbols::clean_symbols; + +use crate::job::job_loop::job_success_result::JobSuccessResult; +use crate::job::job_loop::job_success_result::ResultEntity; +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; +use crate::job::job_types::tts::styletts2::process_job::StyleTTS2ProcessJobArgs; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::InferenceArgs; + +const BUCKET_FILE_PREFIX: &str = "fakeyou_"; +const BUCKET_FILE_EXTENSION: &str = ".wav"; +const MIME_TYPE: &str = "audio/wav"; + + +pub async fn process_inference_voice( + args: StyleTTS2ProcessJobArgs<'_>, + voice_token: String +) -> Result { + let deps = args.job_dependencies; + let job = args.job; + let mysql_pool = &deps.db.mysql_pool; + + let model_dependencies = deps + .job + .job_specific_dependencies + .maybe_styletts2_dependencies + .as_ref() + .ok_or_else(|| ProcessSingleJobError::JobSystemMisconfiguration(Some("missing StyleTTS2 dependencies".to_string())))?; + + // get some globals + let mut job_progress_reporter = deps + .clients + .job_progress_reporter + .new_generic_inference(job.inference_job_token.as_str()) + .map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?; + + // get job args + let text = match job.maybe_raw_inference_text.clone() { + Some(value) => { + clean_symbols(value.as_str()) + } + None => { + return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Text for Inference"))); + } + }; + + let voice_token = tokens::tokens::zs_voices::ZsVoiceToken(voice_token); + + // Get voice bucket hash - from voice token + let voice_lookup_result = get_voice_by_token(&voice_token, false, &mysql_pool).await; + + let voice = match voice_lookup_result { + Ok(Some(voice)) => voice, + Ok(None) => { + warn!("Voice not found: {:?}", voice_token); + return Err(ProcessSingleJobError::Other(anyhow!("Voice not found: {:?}", voice_token))); + } + Err(err) => { + warn!("Error looking up voice: {:?}", err); + return Err(ProcessSingleJobError::Other(anyhow!("Error looking up voice: {:?}", err))); + } + }; + + // Need to download the models + info!("Download models (if not present)..."); + + for downloader in model_dependencies.downloaders.all_downloaders() { + let result = downloader.download_if_not_on_filesystem( + &args.job_dependencies.buckets.private_bucket_client, + &args.job_dependencies.fs.scoped_temp_dir_creator_for_short_lived_downloads + ).await; + + if let Err(err) = result { + error!("could not download: {:?}", err); + return Err(err); + } + } + + // run inference + let work_temp_dir = format!("/tmp/temp_zeroshot_inference_{}", job.id.0); + + // NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem. + let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work + .new_tempdir(&work_temp_dir) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let _workdir = work_temp_dir.path().to_path_buf(); + let filename = "style.npz".to_string(); + + let mut downloaded_weights_path = work_temp_dir.path().to_path_buf(); + downloaded_weights_path.push(&filename); + + let voice_file = download_voice_embedding( + &voice, + &deps.buckets.private_bucket_client, + &downloaded_weights_path + ).await?; + + println!("voice file path! {}", voice_file.filesystem_path.to_string_lossy()); + + // Download embeddings file using embedding token + // Create a temp dir to download things to + job_progress_reporter + .log_status("running inference") + .map_err(|e| ProcessSingleJobError::Other(e))?; + + let text_input_fs_path = work_temp_dir.path().join("inference_input.txt"); + + std::fs::write(&text_input_fs_path, &text) + .map_err(|e| ProcessSingleJobError::from_io_error(e))?; + + let inference_start_time = Instant::now(); + + let output_file_name = work_temp_dir.path().join("output.wav"); + let stderr_output_file = work_temp_dir.path().join("zero_shot_inference.txt"); + + // Run Inference + let command_exit_status = + model_dependencies.inference_command.execute_inference( + InferenceArgs { + input_embedding_file_path: &downloaded_weights_path, + input_text_file_path: &text_input_fs_path, // text + output_file_name: &output_file_name, // output file name in the output folder + stderr_output_file: &stderr_output_file, + } + ); + + let inference_duration = Instant::now().duration_since(inference_start_time); + + if !command_exit_status.is_success() { + error!("Zero shot TTS failed: {:?}", command_exit_status); + + let error = ProcessSingleJobError::Other(anyhow!("CommandExitStatus: {:?}", command_exit_status)); + + if let Ok(contents) = read_to_string(&stderr_output_file) { + warn!("Captured stderr output: {}", contents); + } + + // Clean up temp files + //safe_delete_temp_file(&audio_path.filesystem_path); + //safe_delete_temp_file(&image_path.filesystem_path); + //safe_delete_temp_file(&usable_image_path); + //safe_delete_temp_file(&output_video_fs_path); + //safe_delete_temp_file(&stderr_output_file); + //safe_delete_temp_directory(&work_temp_dir); + + return Err(error); + } + + // upload audio to bucket + info!("Uploading media ..."); + + let result_bucket_location = MediaFileBucketPath::generate_new( + Some(BUCKET_FILE_PREFIX), + Some(BUCKET_FILE_EXTENSION) + ); + + let result_bucket_object_pathbuf = result_bucket_location.to_full_object_pathbuf(); + + // Finished file path + let mut finished_file = work_temp_dir.path().to_path_buf(); + finished_file.push(&output_file_name); + + info!("Upload Bucket Path: {:?}", result_bucket_object_pathbuf); + info!("Upload File Path: {:?}", finished_file); + + args.job_dependencies.buckets.public_bucket_client + .upload_filename_with_content_type( + &result_bucket_object_pathbuf, + &finished_file, + &MIME_TYPE + ) + .await + .map_err(|e| ProcessSingleJobError::Other(e))?; + + // ==================== UPLOAD AUDIO TO BUCKET ==================== + + info!("Calculating sha256..."); + + let file_checksum = sha256_hash_file(&finished_file).map_err(|err| { + ProcessSingleJobError::Other(anyhow!("Error hashing file: {:?}", err)) + })?; + + let file_size_bytes = file_size(&finished_file).map_err(|err| + ProcessSingleJobError::Other(err) + )?; + + job_progress_reporter.log_status("done").map_err(|e| ProcessSingleJobError::Other(e))?; + + // insert into db the record + let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { + pool: &args.job_dependencies.db.mysql_pool, + job: &job, + origin_model_type: MediaFileOriginModelType::StyleTTS2, + maybe_mime_type: Some(&MIME_TYPE), + file_size_bytes, + sha256_checksum: &file_checksum, + public_bucket_directory_hash: result_bucket_location.get_object_hash(), + maybe_public_bucket_prefix: Some(BUCKET_FILE_PREFIX), + maybe_public_bucket_extension: Some(BUCKET_FILE_EXTENSION), + is_on_prem: args.job_dependencies.job.info.container.is_on_prem, + worker_hostname: &args.job_dependencies.job.info.container.hostname, + worker_cluster: &args.job_dependencies.job.info.container.cluster_name, + }).await.map_err(|e| ProcessSingleJobError::Other(e))?; + + info!( + "Job {:?} complete success! Downloaded, ran inference, and uploaded. Saved model record: {}, Result Token: {}", + job.id, + id, + &media_file_token + ); + + Ok(JobSuccessResult { + maybe_result_entity: Some(ResultEntity { + entity_type: InferenceResultType::MediaFile, + entity_token: media_file_token.to_string(), + }), + inference_duration, + }) +} + +pub struct VoiceFile { + pub filesystem_path: PathBuf, +} + +pub async fn download_voice_embedding_from_hash( + bucket_hash: &str, + name: &str, + private_bucket_client: &BucketClient, + path: &PathBuf +) -> Result { + let unifer = BucketPathUnifier::default_paths(); + let object_path = unifer.zero_shot_tts_speaker_encoding(bucket_hash, 0); + + let mut path = path.clone(); + + let file_name = format!("{}", name); + path.push(&file_name); + + private_bucket_client.download_file_to_disk(object_path, &path) + .await + .map_err(|err| { + error!("Could not download embedding file: {err}"); + ProcessSingleJobError::from_anyhow_error(anyhow!("Could not download embedding file: {err}")) + })?; + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&path.clone()), + }; + + Ok(voice_file) +} + +pub async fn download_voice_embedding( + voice: &ZsVoice, + private_bucket_client: &BucketClient, + download_path: &PathBuf, +) -> Result { + + let embedding_bucket_location = ZeroShotVoiceEmbeddingBucketPath::from_object_hash( + &voice.bucket_hash, + ModelCategory::Tts, + ModelType::StyleTTS2, + voice.model_version + ); + + info!("Downloading embedding from: {:?}", &embedding_bucket_location.to_full_object_pathbuf()); + info!("Downloading to filesystem location: {:?}", download_path); + + let result = private_bucket_client.download_file_to_disk( + embedding_bucket_location.to_full_object_pathbuf(), + &download_path + ).await; + + if let Err(err) = result { + error!("could not download embedding file: {:?}", err); + return Err(ProcessSingleJobError::from_anyhow_error(err)); + } + + let voice_file = VoiceFile { + filesystem_path: PathBuf::from(&download_path.clone()), + }; + + Ok(voice_file) +} + diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs new file mode 100644 index 0000000000..eb96c87c7a --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_dependencies.rs @@ -0,0 +1,20 @@ +use errors::AnyhowResult; + +use crate::job::job_types::tts::styletts2::model_downloaders::StyleTTS2Downloaders; +use crate::job::job_types::tts::styletts2::styletts2_inference_command::{StyleTTS2CreateEmbeddingCommand, StyleTTS2InferenceCommand}; + +pub struct StyleTTS2Dependencies { + pub create_embedding_command: StyleTTS2CreateEmbeddingCommand, + pub downloaders: StyleTTS2Downloaders, + pub inference_command: StyleTTS2InferenceCommand, +} + +impl StyleTTS2Dependencies { + pub fn setup() -> AnyhowResult { + Ok(Self { + create_embedding_command: StyleTTS2CreateEmbeddingCommand::from_env()?, + downloaders: StyleTTS2Downloaders::build_all_from_env(), + inference_command: StyleTTS2InferenceCommand::from_env()?, + }) + } +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs new file mode 100644 index 0000000000..814cf97939 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/styletts2_inference_command.rs @@ -0,0 +1,497 @@ +use std::collections::HashSet; +use std::env; +use std::ffi::OsString; +use std::fs::File; +use std::path::{Path, PathBuf}; +use std::time::Duration; + +use anyhow::anyhow; +use log::info; +use once_cell::sync::Lazy; +use subprocess::{Popen, PopenConfig, Redirection}; + +use errors::AnyhowResult; +use filesys::path_to_string::path_to_string; +use subprocess_common::docker_options::{DockerFilesystemMount, DockerGpu, DockerOptions}; + +use crate::job::job_loop::command_exit_status::CommandExitStatus; + +#[derive(Clone)] +pub enum ExecutableOrCommand { + /// Eg. `inference.py` + Executable(PathBuf), + /// Eg. `python3 inference.py` + Command(String), +} + +// These environment vars are not copied over to the subprocess +// TODO/FIXME(bt, 2023-05-28): This is horrific security! +static IGNORED_ENVIRONMENT_VARS: Lazy> = Lazy::new(|| { + let env_var_names = ["MYSQL_URL", "ACCESS_KEY", "SECRET_KEY", "NEWRELIC_API_KEY"]; + + env_var_names + .iter() + .map(|value| value.to_string()) + .collect::>() +}); + +// The inference command assumes everything has been setup to run inference +// so we have all the files downloaded for the inputs. +pub struct InferenceArgs> { + /// --driven_audio: path to the input audio + pub input_embedding_file_path: P, // name of the embedding.npz in the work dir + pub input_text_file_path: P, // name of the text + pub output_file_name: P, // output file name in the output folder + pub stderr_output_file: P, +} + +pub struct CreateVoiceInferenceArgs> { + pub output_embedding_path: P, + pub audio_file: String, + pub stderr_output_file: P, +} + +#[derive(Clone)] +pub struct StyleTTS2InferenceCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2InferenceCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + pub fn from_env() -> AnyhowResult { + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_COMMAND" + ); + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_ALTERNATE_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: InferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: InferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" --input "); + command.push_str(&path_to_string(args.input_text_file_path)); + + command.push_str(" --input-style-npz "); + command.push_str(&path_to_string(args.input_embedding_file_path)); + + + command.push_str(" --output "); + command.push_str(&path_to_string(args.output_file_name)); + + //TODO(KS): Make this a parameter + command.push_str(" --vcsteps "); + command.push_str(format!("{}", 20).as_str()); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } +} + +#[derive(Clone)] +pub struct StyleTTS2CreateEmbeddingCommand { + /// Where the code lives + root_code_directory: PathBuf, + /// A single executable script or a much larger bash command. + executable_or_command: ExecutableOrCommand, + /// eg. `source python/bin/activate` + maybe_virtual_env_activation_command: Option, + /// Optional default config file to use + maybe_default_config_path: Option, + /// If this is run under Docker (eg. in development), these are the options. + maybe_docker_options: Option, + /// If the execution should be ended after a certain point. + maybe_execution_timeout: Option, + /// Inference arg. + /// --checkpoint_dir: optional location for checkpoints directory + pub alternate_checkpoint_dir: Option, +} + +impl StyleTTS2CreateEmbeddingCommand { + pub fn new>( + root_code_directory: P, + executable_or_command: ExecutableOrCommand, + maybe_virtual_env_activation_command: Option<&str>, + maybe_default_config_path: Option

, + maybe_docker_options: Option, + maybe_execution_timeout: Option, + alternate_checkpoint_dir: Option + ) -> Self { + Self { + root_code_directory: root_code_directory.as_ref().to_path_buf(), + executable_or_command, + maybe_virtual_env_activation_command: maybe_virtual_env_activation_command.map(|s| + s.to_string() + ), + maybe_default_config_path: maybe_default_config_path.map(|p| p.as_ref().to_path_buf()), + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + } + } + + // helper funciton to set up env + pub fn from_env() -> AnyhowResult { + // TODO: fix for now just for compilation + let root_code_directory = easyenv::get_env_pathbuf_required( + "STYLE_TTS2_INFERENCE_ROOT_DIRECTORY" + )?; + + let maybe_inference_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_COMMAND" + ); + + // EVENTUALLY should remove the check + // Optional, eg. `./infer.py`. Typically we'll use the command form instead. + let maybe_inference_executable = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_EXECUTABLE" + ); + + let executable_or_command = match maybe_inference_command { + Some(command) => ExecutableOrCommand::Command(command), + None => + match maybe_inference_executable { + Some(executable) => ExecutableOrCommand::Executable(executable), + None => { + return Err(anyhow!("neither command nor executable passed")); + } + } + }; + + let maybe_virtual_env_activation_command = easyenv::get_env_string_optional( + "STYLE_TTS2_INFERENCE_MAYBE_VENV_COMMAND" + ); + + let maybe_default_config_path = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_INFERENCE_MAYBE_DEFAULT_CONFIG_PATH" + ); + + let maybe_execution_timeout = easyenv::get_env_duration_seconds_optional("TIMEOUT_SECONDS"); + + // Probably for local + let maybe_docker_options = easyenv + ::get_env_string_optional("STYLE_TTS2_INFERENCE_MAYBE_DOCKER_IMAGE") + .map(|image_name| { + DockerOptions { + image_name, + maybe_bind_mount: Some(DockerFilesystemMount::tmp_to_tmp()), + maybe_environment_variables: None, + maybe_gpu: Some(DockerGpu::All), + } + }); + + // Override for --checkpoint_dir at inference time + let alternate_checkpoint_dir = easyenv::get_env_pathbuf_optional( + "STYLE_TTS2_CHECKPOINT_PATH" + ); + + Ok(Self { + root_code_directory, + executable_or_command, + maybe_virtual_env_activation_command, + maybe_default_config_path, + maybe_docker_options, + maybe_execution_timeout, + alternate_checkpoint_dir, + }) + } + + pub fn execute_inference>(&self, args: CreateVoiceInferenceArgs

) -> CommandExitStatus { + match self.do_execute_inference(args) { + Ok(exit_status) => exit_status, + Err(error) => + CommandExitStatus::FailureWithReason { reason: format!("error: {:?}", error) }, + } + } + + fn do_execute_inference>( + &self, + args: CreateVoiceInferenceArgs

+ ) -> AnyhowResult { + let mut command = String::new(); + command.push_str(&format!("cd {}", path_to_string(&self.root_code_directory))); + + if let Some(venv_command) = self.maybe_virtual_env_activation_command.as_deref() { + command.push_str(" && "); + command.push_str(venv_command); + command.push_str(" "); + } + + command.push_str(" && "); + + match self.executable_or_command { + ExecutableOrCommand::Executable(ref executable) => { + command.push_str(&path_to_string(executable)); + command.push_str(" "); + } + ExecutableOrCommand::Command(ref cmd) => { + command.push_str(cmd); + command.push_str(" "); + } + } + + // ===== Begin Python Args ===== + + command.push_str(" --output-style-npz "); + command.push_str(&path_to_string(&args.output_embedding_path)); + + command.push_str(" --voice "); + command.push_str(&args.audio_file); + + command.push_str(" --vcsteps "); + command.push_str(format!("{}", 20).as_str()); + + command.push_str(" --input "); + command.push_str("/dev/null"); + + // ===== End Python Args ===== + + if let Some(docker_options) = self.maybe_docker_options.as_ref() { + command = docker_options.to_command_string(&command); + } + + info!("Command: {:?}", command); + + let command_parts = ["bash", "-c", &command]; + + let mut env_vars = Vec::new(); + + // Copy all environment variables from the parent process. + // This is necessary to send all the kubernetes settings for Nvidia / CUDA. + for (env_key, env_value) in env::vars() { + if IGNORED_ENVIRONMENT_VARS.contains(&env_key) { + continue; + } + env_vars.push((OsString::from(env_key), OsString::from(env_value))); + } + + let mut config = PopenConfig::default(); + + info!("stderr will be written to file: {:?}", args.stderr_output_file.as_ref()); + + let stderr_file = File::create(&args.stderr_output_file)?; + config.stderr = Redirection::File(stderr_file); + + if !env_vars.is_empty() { + config.env = Some(env_vars); + } + + let mut p = Popen::create(&command_parts, config)?; + + info!("Subprocess PID: {:?}", p.pid()); + + match self.maybe_execution_timeout { + None => { + let exit_status = p.wait()?; + info!("Subprocess exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + Some(timeout) => { + info!("Executing with timeout: {:?}", &timeout); + let exit_status = p.wait_timeout(timeout.clone())?; + + match exit_status { + None => { + // NB: If the program didn't successfully terminate, kill it. + info!( + "Subprocess didn't end after timeout: {:?}; terminating...", + &timeout + ); + let _r = p.terminate()?; + Ok(CommandExitStatus::Timeout) + } + Some(exit_status) => { + info!("Subprocess timed wait exit status: {:?}", exit_status); + Ok(CommandExitStatus::from_exit_status(exit_status)) + } + } + } + } + } + +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs new file mode 100644 index 0000000000..41b76feac1 --- /dev/null +++ b/crates/service/job/inference_job/src/job/job_types/tts/styletts2/validate_job.rs @@ -0,0 +1,84 @@ +use anyhow::anyhow; + +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ + InferenceCategoryAbbreviated, + PolymorphicInferenceArgs, +}; +use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; + +use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; + +pub enum JobType { + Create, + Inference +} +// determines the job type voice inference or voice creation +pub struct JobArgs { + pub voice_token: Option, // infers + pub voice_dataset_token: Option, // creates + pub job_type: JobType +} + +pub fn validate_job(job: &AvailableInferenceJob) -> Result { + let inference_category = job.maybe_inference_args + .as_ref() + .map(|args| args.inference_category) + .flatten(); + + match inference_category { + Some(InferenceCategoryAbbreviated::TextToSpeech) => {} + Some(category) => { + return Err( + ProcessSingleJobError::from_anyhow_error( + anyhow!("wrong inference category for job: {:?}", category) + ) + ); + } + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference category for job!")) + ); + } + } + + let inference_args = job.maybe_inference_args + .as_ref() + .map(|args| args.args.as_ref()) + .flatten(); + + let inference_args = match inference_args { + Some(args) => args, + None => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("no inference args for job!")) + ); + } + }; + + let ttsArgs = match inference_args { + PolymorphicInferenceArgs::Tts(inference_args) => inference_args, + _ => { + return Err( + ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!")) + ); + } + }; + + if let Some(voice_token) = &ttsArgs.voice_token { + return Ok::(JobArgs { + voice_token: Some(voice_token.clone()), + voice_dataset_token: None, + job_type: JobType::Inference + }); + } + + if let Some(dataset_token) = &ttsArgs.dataset_token { + return Ok::(JobArgs { + voice_token: None, + voice_dataset_token: Some(dataset_token.clone()), + job_type: JobType::Create + }); + } + + Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Missing Voice Token!"))) +} diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs index 72272f4fab..7646fd7652 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/process_job_inference_voice.rs @@ -11,6 +11,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use cloud_storage::bucket_client::BucketClient; use cloud_storage::bucket_path_unifier::BucketPathUnifier; use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; use filesys::file_size::file_size; use hashing::sha256::sha256_hash_file::sha256_hash_file; use mysql_queries::queries::media_files::create::insert_media_file_from_zero_shot_tts::insert_media_file_from_zero_shot; @@ -224,6 +225,7 @@ pub async fn process_inference_voice( let (media_file_token, id) = insert_media_file_from_zero_shot(InsertArgs { pool: &args.job_dependencies.db.mysql_pool, job: &job, + origin_model_type: MediaFileOriginModelType::VallEX, maybe_mime_type: Some(&MIME_TYPE), file_size_bytes, sha256_checksum: &file_checksum, diff --git a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs index 09721a2a22..46516f93d7 100644 --- a/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs +++ b/crates/service/job/inference_job/src/job/job_types/tts/vall_e_x/vall_e_x_inference_command.rs @@ -72,6 +72,7 @@ pub struct VallEXInferenceCommand { pub alternate_checkpoint_dir: Option, } + impl VallEXInferenceCommand { pub fn new>( root_code_directory: P, diff --git a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs index e7bfdd0db3..81bc889486 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/rvc_v2/process_job.rs @@ -260,6 +260,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { override_f0_method, .. } => *override_f0_method, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, PolymorphicInferenceArgs::Wf(_) => None, }) @@ -271,6 +272,7 @@ pub async fn process_job(args: RvcV2ProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, PolymorphicInferenceArgs::Wf(_) => None, }) diff --git a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs index 485250fb99..110d3f9a1c 100644 --- a/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/vc/so_vits_svc/process_job.rs @@ -180,6 +180,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { auto_predict_f0, .. } => *auto_predict_f0, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, PolymorphicInferenceArgs::Wf(_) => None, }) @@ -192,6 +193,7 @@ pub async fn process_job(args: SoVitsSvcProcessJobArgs<'_>) -> Result None, PolymorphicInferenceArgs::Vc { transpose, .. } => *transpose, PolymorphicInferenceArgs::Rr(_) => None, + PolymorphicInferenceArgs::Ig(_) => None, PolymorphicInferenceArgs::Mc(_) => None, PolymorphicInferenceArgs::Wf(_) => None, }) diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs index bc5ded36a7..0cd5fa52f4 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_model_file.rs @@ -16,7 +16,7 @@ use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct ModelFile { pub filesystem_path: PathBuf, } - +// TODO replace this with remote cloud file manager. pub async fn download_model_file( model_record: &Option, public_bucket_client: &BucketClient, diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs index 6ec512d1ea..597757985a 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/download_video_file.rs @@ -22,7 +22,7 @@ use crate::util::scoped_temp_dir_creator::ScopedTempDirCreator; pub struct VideoFile { pub filesystem_path: PathBuf, } - +// TODO replace this with remote cloud file manager. pub async fn download_video_file( video_source: &VideofilterVideoSource, public_bucket_client: &BucketClient, @@ -63,23 +63,23 @@ pub async fn download_video_file( } } } - VideofilterVideoSource::U(media_upload_token) => { - let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); - let media_upload_result = get_media_upload_for_inference( - &media_upload_token, - mysql_pool - ).await; + VideofilterVideoSource::U(media_upload_token) => { + let media_upload_token = MediaUploadToken::new_from_str(media_upload_token); + let media_upload_result = get_media_upload_for_inference( + &media_upload_token, + mysql_pool + ).await; - let media_upload_result = match media_upload_result { - Ok(Some(result)) => result, - Ok(None) => { - warn!("could not find media upload: {:?}", media_upload_token); - return Err(ProcessSingleJobError::from_anyhow_error( - anyhow!("could not find media upload: {:?}", media_upload_token))) - } - Err(e) => { - error!("could not query media upload: {:?}", e); - return Err(ProcessSingleJobError::from_anyhow_error(e)) + let media_upload_result = match media_upload_result { + Ok(Some(result)) => result, + Ok(None) => { + warn!("could not find media upload: {:?}", media_upload_token); + return Err(ProcessSingleJobError::from_anyhow_error( + anyhow!("could not find media upload: {:?}", media_upload_token))) + } + Err(e) => { + error!("could not query media upload: {:?}", e); + return Err(ProcessSingleJobError::from_anyhow_error(e)) } }; diff --git a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs index 7de58ac925..f6e556dddf 100644 --- a/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/videofilter/rerender_a_video/process_job.rs @@ -35,7 +35,6 @@ const BUCKET_FILE_EXTENSION: &str = ".mp4"; pub struct RerenderProcessJobArgs<'a> { pub job_dependencies: &'a JobDependencies, pub job: &'a AvailableInferenceJob, - // pub media_file: MediaFile, } pub async fn process_job(args: RerenderProcessJobArgs<'_>) -> Result { diff --git a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs index 7408936e2f..ed43097993 100644 --- a/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs +++ b/crates/service/job/inference_job/src/job/job_types/workflow/comfy_ui/validate_job.rs @@ -2,14 +2,15 @@ use std::collections::HashMap; use anyhow::anyhow; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; -use mysql_queries::payloads::generic_inference_args::workflow_payload::{WorkflowArgs, FileSource, Dependency}; +use mysql_queries::payloads::generic_inference_args::workflow_payload::{NewValue}; use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob; +use tokens::tokens::model_weights::ModelWeightToken; use crate::job::job_loop::process_single_job_error::ProcessSingleJobError; pub struct JobArgs<'a> { - pub dependencies: &'a HashMap, - pub workflow_source: &'a FileSource, + pub workflow_source: &'a ModelWeightToken, + pub maybe_json_modifications: &'a Option>, } pub fn validate_job(job: &AvailableInferenceJob) -> Result { @@ -47,14 +48,7 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result args, - None => { - return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No dependency map provided!"))); - } - }; - - let workflow_source = match &inference_args.maybe_workflow_source { + let workflow_source = match &inference_args.maybe_workflow_config { Some(args) => args, None => { return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No workflow source provided!"))); @@ -63,7 +57,7 @@ pub fn validate_job(job: &AvailableInferenceJob) -> Result, pub maybe_vits_dependencies: Option, pub maybe_rerender_dependencies: Option, + pub maybe_stable_diffusion_dependencies: Option, pub maybe_mocapnet_dependencies: Option, + pub maybe_styletts2_dependencies: Option, pub maybe_comfy_ui_dependencies: Option, } @@ -34,7 +38,9 @@ impl JobSpecificDependencies { let mut maybe_vall_e_x_dependencies = None; let mut maybe_vits_dependencies = None; let mut maybe_rerender_dependencies = None; + let mut maybe_stable_diffusion_dependencies = None; let mut maybe_mocapnet_dependencies = None; + let mut maybe_styletts2_dependencies = None; let mut maybe_comfy_ui_dependencies = None; if scoped_execution.can_run_job(InferenceModelType::RvcV2) { @@ -62,6 +68,11 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies = Some(VallExDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::StyleTTS2) { + print_with_space("Setting StyleTTS2 dependencies..."); + maybe_styletts2_dependencies = Some(StyleTTS2Dependencies::setup()?); + } + if scoped_execution.can_run_job(InferenceModelType::Vits) { print_with_space("Setting Vits dependencies..."); maybe_vits_dependencies = Some(VitsDependencies::setup()?); @@ -72,6 +83,11 @@ impl JobSpecificDependencies { maybe_rerender_dependencies = Some(RerenderDependencies::setup()?); } + if scoped_execution.can_run_job(InferenceModelType::StableDiffusion) { + print_with_space("Setting Stable Diffusion dependencies..."); + maybe_stable_diffusion_dependencies = Some(StableDiffusionDependencies::setup()?); + } + if scoped_execution.can_run_job(InferenceModelType::MocapNet) { print_with_space("Setting MocapNet dependencies..."); maybe_mocapnet_dependencies = Some(MocapNetDependencies::setup()?); @@ -90,7 +106,9 @@ impl JobSpecificDependencies { maybe_vall_e_x_dependencies, maybe_vits_dependencies, maybe_rerender_dependencies, + maybe_stable_diffusion_dependencies, maybe_mocapnet_dependencies, + maybe_styletts2_dependencies, maybe_comfy_ui_dependencies, }) } diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index d8224d0269..0926381123 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -24,6 +24,7 @@ actix_helpers = { path = "../../../lib/actix_helpers" } bootstrap = { path = "../../../lib/bootstrap" } buckets = { path = "../../../_types/buckets" } cloud_storage = { path = "../../../lib/api_clients/cloud_storage" } +composite_identifiers = { path = "../../../_types/composite_identifiers" } config = { path = "../../../lib/deprecated/config" } container_common = { path = "../../../lib/deprecated/container_common" } cookies = { path = "../../../lib/cookies" } @@ -176,3 +177,5 @@ diesel = { version = "2.1.3", features = ["mysql"] } [dev-dependencies] speculoos = "0.9.0" # Fluent assertions +strum = { version = "0.25", features = ["derive"] } + diff --git a/crates/service/web/storyteller_web/src/api_doc.rs b/crates/service/web/storyteller_web/src/api_doc.rs index f8f0a63291..a6648d4983 100644 --- a/crates/service/web/storyteller_web/src/api_doc.rs +++ b/crates/service/web/storyteller_web/src/api_doc.rs @@ -16,6 +16,7 @@ use tokens::tokens::user_bookmarks::*; use tokens::tokens::users::*; use crate::http_server::common_responses::cover_image_details::*; +use crate::http_server::common_responses::media_file_origin_details::*; use crate::http_server::common_responses::media_file_social_meta_lite::MediaFileSocialMetaLight; use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::pagination_page::PaginationPage; @@ -90,6 +91,8 @@ use crate::http_server::web_utils::response_success_helpers::*; CoverImageDetails, DefaultAvatarInfo, DefaultCoverInfo, + MediaFileModelDetails, + MediaFileOriginDetails, MediaFileSocialMetaLight, PaginationCursors, PaginationPage, diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs index de1a9f7b3b..f0e7faf854 100644 --- a/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/cover_image_details.rs @@ -3,6 +3,7 @@ use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::users::UserToken; +use crate::util::placeholder_images::cover_images::default_cover_image_color_from_token::default_cover_image_color_from_token; use crate::util::placeholder_images::cover_images::default_cover_image_from_token::default_cover_image_from_token; use crate::util::placeholder_images::user_avatars::default_avatar_color_from_username::default_avatar_color_from_username; @@ -25,6 +26,8 @@ pub struct CoverImageDetails { pub struct DefaultCoverInfo { /// Which image to show. pub image_index: u8, + /// Which color to use. + pub color_index: u8, } impl CoverImageDetails { @@ -60,6 +63,7 @@ impl DefaultCoverInfo { pub fn from_token(model_weight_token: &ModelWeightToken) -> Self { Self { image_index: default_cover_image_from_token(model_weight_token), + color_index: default_cover_image_color_from_token(model_weight_token), } } } diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs new file mode 100644 index 0000000000..11aaaccbe1 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/media_file_origin_details.rs @@ -0,0 +1,104 @@ +use utoipa::ToSchema; + +use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory; +use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType; +use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory; +use tokens::tokens::model_weights::ModelWeightToken; + +/// Fields useful for enriching media file listings +#[derive(Clone, Serialize, ToSchema)] +pub struct MediaFileOriginDetails { + /// Where the file came from (broadly) + pub origin_category: MediaFileOriginCategory, + + /// Where the file came from (specifically, a product area, eg. "face_animator") + pub product_category: MediaFileOriginProductCategory, + + /// If the file was produced by a model, the details. + pub maybe_model: Option, +} + +/// Information about the model weights +/// https://serde.rs/enum-representations.html#untagged +#[derive(Clone, Serialize, ToSchema)] +#[serde(untagged)] +pub enum MediaFileModelDetails { + /// A system model is one that we maintain and update, + /// eg. SadTalker, Wav2Lip, MocapNet, etc. + SystemModel { + /// The type of model weight + model_type: MediaFileOriginModelType, + }, + /// A model weight is typically a user-submitted model. + /// We store lots of these. + ModelWeight { + /// The type of model weight + model_type: MediaFileOriginModelType, + /// The model token + token: ModelWeightToken, + /// The model title (typically only populated for `model_weights` models, not legacy tables such as `tts_models`.) + title: String, + }, +} + +impl MediaFileOriginDetails { + pub fn from_db_fields_owned( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option, + maybe_model_title: Option, + ) -> Self { + let maybe_model = match (maybe_model_type, maybe_model_token, maybe_model_title) { + (Some(model_type), Some(model_token), Some(model_title)) => { + Some(MediaFileModelDetails::ModelWeight { + model_type, + token: model_token, + title: model_title, + }) + }, + (Some(model_type), None, None) => { + Some(MediaFileModelDetails::SystemModel { + model_type, + }) + }, + _ => None, + }; + + Self { + origin_category, + product_category, + maybe_model, + } + } + + pub fn from_db_fields( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option<&ModelWeightToken>, + maybe_model_title: Option<&str>, + ) -> Self { + Self::from_db_fields_owned( + origin_category, + product_category, + maybe_model_type, + maybe_model_token.cloned(), + maybe_model_title.map(|s|s.to_string())) + } + + pub fn from_db_fields_str( + origin_category: MediaFileOriginCategory, + product_category: MediaFileOriginProductCategory, + maybe_model_type: Option, + maybe_model_token: Option<&str>, + maybe_model_title: Option<&str>, + ) -> Self { + Self::from_db_fields_owned( + origin_category, + product_category, + maybe_model_type, + maybe_model_token.map(|s| ModelWeightToken::new_from_str(s)), + maybe_model_title.map(|s| s.to_string())) + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs b/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs index 4d3f5bb33e..a46efd6cc9 100644 --- a/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/common_responses/mod.rs @@ -1,4 +1,5 @@ pub mod cover_image_details; +pub mod media_file_origin_details; pub mod media_file_social_meta_lite; pub mod pagination_cursors; pub mod pagination_page; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs new file mode 100644 index 0000000000..752c69fc86 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_generation_type.rs @@ -0,0 +1,76 @@ +use std::collections::BTreeSet; + +#[cfg(test)] +use strum::EnumCount; +#[cfg(test)] +use strum::EnumIter; +use utoipa::ToSchema; + + +#[cfg_attr(test, derive(EnumIter, EnumCount))] +#[derive(Clone, Copy, Eq, PartialEq, Hash, Ord, PartialOrd, Debug, Deserialize, Serialize, ToSchema)] +pub enum EnqueueImageGenType { + #[serde(rename = "upload_lora")] + UploadLoRA, + #[serde(rename = "upload_sd")] + Checkpoint, + #[serde(rename = "inference")] + Inference, +} + +impl EnqueueImageGenType { + pub fn to_str(&self) -> &'static str { + match self { + Self::UploadLoRA => "upload_lora", + Self::Checkpoint => "upload_sd", + Self::Inference => "inference", + } + } + + pub fn from_str(value: &str) -> Result { + match value { + "upload_lora" => Ok(Self::UploadLoRA), + "upload_sd" => Ok(Self::Checkpoint), + "inference" => Ok(Self::Inference), + _ => Err(format!("invalid value: {:?}", value)), + } + } + + pub fn all_variants() -> BTreeSet { + BTreeSet::from([ + Self::UploadLoRA, + Self::Checkpoint, + Self::Inference + ]) + } +} + + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_to_str() { + assert_eq!(EnqueueImageGenType::UploadLoRA.to_str(), "upload_lora"); + assert_eq!(EnqueueImageGenType::Checkpoint.to_str(), "upload_sd"); + assert_eq!(EnqueueImageGenType::Inference.to_str(), "inference"); + } + + #[test] + fn test_from_str() { + assert_eq!(EnqueueImageGenType::from_str("upload_lora").unwrap(), EnqueueImageGenType::UploadLoRA); + assert_eq!(EnqueueImageGenType::from_str("upload_sd").unwrap(), EnqueueImageGenType::Checkpoint); + assert_eq!(EnqueueImageGenType::from_str("inference").unwrap(), EnqueueImageGenType::Inference); + assert!(EnqueueImageGenType::from_str("invalid").is_err()); + } + + #[test] + fn test_all_variants() { + let variants = EnqueueImageGenType::all_variants(); + assert_eq!(variants.len(), 3); + assert!(variants.contains(&EnqueueImageGenType::UploadLoRA)); + assert!(variants.contains(&EnqueueImageGenType::Checkpoint)); + assert!(variants.contains(&EnqueueImageGenType::Inference)); + } +} \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs new file mode 100644 index 0000000000..2fe137a43f --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/enqueue_image_generation.rs @@ -0,0 +1,399 @@ +// #![forbid(unused_imports)] +#![forbid(unused_mut)] +#![forbid(unused_variables)] + +use std::fmt::{Display, Formatter}; +use std::fmt::Debug; +use std::sync::Arc; + +use actix_web::{HttpRequest, HttpResponse, web}; +use actix_web::error::ResponseError; +use actix_web::http::StatusCode; +use log::warn; +use serde::Deserialize; +use serde::Serialize; +use utoipa::ToSchema; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; +use http_server_common::request::get_request_header_optional::get_request_header_optional; +use http_server_common::request::get_request_ip::get_request_ip; +use mysql_queries::payloads::generic_inference_args::generic_inference_args::{ + GenericInferenceArgs, + InferenceCategoryAbbreviated, + PolymorphicInferenceArgs, +}; +use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs; +use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{ + insert_generic_inference_job, + InsertGenericInferenceArgs, +}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::users::UserToken; + +use crate::configs::plans::get_correct_plan_for_session::get_correct_plan_for_session; +use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; +use crate::server_state::ServerState; + +/// Debug requests can get routed to special "debug-only" workers, which can +/// be used to trial new code, run debugging, etc. +const DEBUG_HEADER_NAME: &str = "enable-debug-mode"; + +/// The routing tag header can send workloads to particular k8s hosts. +/// This is useful for catching the live logs or intercepting the job. +const ROUTING_TAG_HEADER_NAME: &str = "routing-tag"; + +#[derive(Debug,Deserialize, Clone, Copy, Eq, PartialEq)] +pub enum TypeOfInference { + #[serde(rename = "lora")] + UploadLoRA, + #[serde(rename = "checkpoint")] + CheckPoint, + #[serde(rename = "inference")] + Inference, +} +impl Display for TypeOfInference { + fn fmt(&self, f: &mut Formatter) -> std::fmt::Result { + match self { + TypeOfInference::UploadLoRA => write!(f, "lora"), + TypeOfInference::CheckPoint => write!(f, "checkpoint"), + TypeOfInference::Inference => write!(f, "inference"), + } + } +} + + +pub fn is_valid_string(input: &str) -> bool { + let valid_strings = [ + "DPM++ 2M Karras", + "DPM++ SDE Karras", + "DPM++ 2M SDE Exponential", + "DPM++ 2M SDE Karras", + "Euler a", + "Euler", + "LMS", + "Heun", + "DPM2", + "DPM2 a", + "DPM++ 2S a", + "DPM++ 2M", + "DPM++ SDE", + "DPM++ 2M SDE", + "DPM++ 2M SDE Heun", + "DPM++ 2M SDE Heun Karras", + "DPM++ 2M SDE Heun Exponential", + "DPM++ 3M SDE", + "DPM++ 3M SDE Karras", + "DPM++ 3M SDE Exponential", + "DPM fast", + "DPM adaptive", + "LMS Karras", + "DPM2 Karras", + "DPM2 a Karras", + "DPM++ 2S a Karras" + ]; + valid_strings.contains(&input) +} + + +#[derive(Deserialize,ToSchema)] +pub struct EnqueueImageGenRequest { + uuid_idempotency_token: String, + type_of_inference: TypeOfInference, // upload loRA / check point / standard inference + maybe_image_source: Option, + maybe_sd_model_token: Option, + maybe_lora_model_token: Option, + maybe_prompt: Option, + maybe_n_prompt: Option, + maybe_seed: Option, + maybe_width: Option, + maybe_height: Option, + maybe_sampler: Option, + maybe_upload_path: Option, + maybe_lora_upload_path: Option, + maybe_cfg_scale: Option, + maybe_number_of_samples: Option, + maybe_batch_count: Option, + name: String, + description: String, +} + +#[derive(Serialize,ToSchema)] +pub struct EnqueueImageGenRequestSuccessResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} + +#[derive(Debug,ToSchema)] +pub enum EnqueueImageGenRequestError { + BadInput(String), + NotAuthorized, + ServerError, + RateLimited, +} + +impl ResponseError for EnqueueImageGenRequestError { + fn status_code(&self) -> StatusCode { + match *self { + EnqueueImageGenRequestError::BadInput(_) => StatusCode::BAD_REQUEST, + EnqueueImageGenRequestError::NotAuthorized => StatusCode::UNAUTHORIZED, + EnqueueImageGenRequestError::ServerError => StatusCode::INTERNAL_SERVER_ERROR, + EnqueueImageGenRequestError::RateLimited => StatusCode::TOO_MANY_REQUESTS, + } + } + + fn error_response(&self) -> HttpResponse { + let error_reason = match self { + EnqueueImageGenRequestError::BadInput(reason) => reason.to_string(), + EnqueueImageGenRequestError::NotAuthorized => "unauthorized".to_string(), + EnqueueImageGenRequestError::ServerError => "server error".to_string(), + EnqueueImageGenRequestError::RateLimited => "rate limited".to_string(), + }; + + to_simple_json_error(&error_reason, self.status_code()) + } +} + +impl std::fmt::Display for EnqueueImageGenRequestError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + write!(f, "{:?}", self) + } +} + +// Implementation for enqueuing a TTS request +// Reference enqueue_infer_tts_handler.rs for checks: rate limiting / user sessions +// insert generic inference job.rs +// Need to convert it to generic inference job. +#[utoipa::path( +post, +path = "/inference/enqueue_image_gen/", +responses( +(status = 200, description = "Enqueue TTS generically", body = EnqueueImageGenRequestSuccessResponse), +(status = 400, description = "Bad input", body = EnqueueImageGenRequestError), +(status = 401, description = "Not authorized", body = EnqueueImageGenRequestError), +(status = 429, description = "Rate limited", body = EnqueueImageGenRequestError), +(status = 500, description = "Server error", body = EnqueueImageGenRequestError) +), +params( +("request" = EnqueueImageGenRequest, description = "Payload for TTS Request") +) +)] +pub async fn enqueue_image_generation_request( + http_request: HttpRequest, + request: web::Json, + server_state: web::Data> +) -> Result { + + // TODO: Brandon need to figure out premium vs not premium + + let mut maybe_user_token: Option = None; + let visbility = enums::common::visibility::Visibility::Public; + + let mut mysql_connection = server_state.mysql_pool.acquire().await.map_err(|err| { + warn!("MySql pool error: {:?}", err); + EnqueueImageGenRequestError::ServerError + })?; + + // ==================== USER SESSION ==================== // + + let maybe_user_session = server_state.session_checker + .maybe_get_user_session_extended_from_connection(&http_request, &mut mysql_connection).await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + EnqueueImageGenRequestError::ServerError + })?; + + + if let Some(user_session) = maybe_user_session.as_ref() { + maybe_user_token = Some(UserToken::new_from_str(&user_session.user_token)); + } + + // Plan should handle "first anonymous use" and "investor" cases. + let plan = get_correct_plan_for_session( + server_state.server_environment, + maybe_user_session.as_ref() + ); + + // Separate priority for animation. + let priority_level = plan.web_vc_base_priority_level(); + + // ==================== DEBUG MODE + ROUTING TAG ==================== // + + let is_debug_request = get_request_header_optional(&http_request, DEBUG_HEADER_NAME).is_some(); + + let maybe_routing_tag = get_request_header_optional(&http_request, ROUTING_TAG_HEADER_NAME).map( + |routing_tag| routing_tag.trim().to_string() + ); + + // ==================== BANNED USERS ==================== // + + if let Some(ref user) = maybe_user_session { + if user.role.is_banned { + return Err(EnqueueImageGenRequestError::NotAuthorized); + } + } + + // DETECT premium user and queue + + // ==================== RATE LIMIT ==================== // + + let rate_limiter = match maybe_user_session { + None => &server_state.redis_rate_limiters.logged_out, + Some(ref _user) => &server_state.redis_rate_limiters.logged_in, + }; + + if let Err(_err) = rate_limiter.rate_limit_request(&http_request) { + return Err(EnqueueImageGenRequestError::RateLimited); + } + + // Get up IP address + let ip_address = get_request_ip(&http_request); + + // Check the inference args to make sure everything is all there for upload loRA / model or standard inference + + // ==================== INFERENCE ARGS ==================== // + + if request.type_of_inference == TypeOfInference::Inference { + if request.maybe_sd_model_token.is_none() { + return Err(EnqueueImageGenRequestError::BadInput("No sd model token provided".to_string())); + } + + if request.maybe_prompt.is_none() { + return Err(EnqueueImageGenRequestError::BadInput("No prompt provided".to_string())); + } + + } else if request.type_of_inference == TypeOfInference::UploadLoRA { + if request.maybe_lora_upload_path.is_none() { + return Err(EnqueueImageGenRequestError::BadInput("No lora upload path provided".to_string())); + } + } else if request.type_of_inference == TypeOfInference::CheckPoint { + if request.maybe_upload_path.is_none() { + return Err(EnqueueImageGenRequestError::BadInput("No upload path provided".to_string())); + } + } + + let width = request.maybe_width.unwrap_or(512); + let height = request.maybe_height.unwrap_or(512); + + let width = match width { + 0..=512 => 512, + 513..=768 => 768, + 769..=1024 => 1024, + _ => 1024, + }; + + let height = match height { + 0..=512 => 512, + 513..=768 => 768, + 769..=1024 => 1024, + _ => 1024, + }; + + let cfg_scale = match request.maybe_cfg_scale { + Some(val) => if val > 32 { 32 } else { val }, + None => 7 + }; + + let number_of_samples = match request.maybe_number_of_samples { + Some(val) => if val < 0 { 20 } else if val > 128 { 128 } else { val }, + None => 20, + }; + + + let batch_count = match request.maybe_batch_count { + Some(val) => if val > 4 { 4 } else { val }, + None => 1, + }; + + let sampler = match request.maybe_sampler.clone() { + Some(val) => if is_valid_string(&val) { val.clone() } else { String::from("Euler a") }, + None =>String::from( "Euler a"), + }; + + let sd_weight_token = ModelWeightToken(request.maybe_sd_model_token.clone().unwrap_or_default()); + let lora_token = ModelWeightToken(request.maybe_lora_model_token.clone().unwrap_or_default()); + + + let n_prompt = request.maybe_n_prompt.clone().unwrap_or_default(); + + // we can only do 1 upload type at a time. + // if we are uploading a model. + let upload_path = request.maybe_upload_path.clone().unwrap_or_default(); + // if we are uploading a lora model. + let lora_upload_path = request.maybe_lora_upload_path.clone().unwrap_or_default(); + + let mut seed = -1; + + if let Some(s) = request.maybe_seed { + seed = s; + } + + let type_of_inference = request.type_of_inference.to_string().clone(); + let description = request.description.to_string().clone(); + let name = request.name.to_string().clone(); + let inference_args = StableDiffusionArgs { + maybe_sd_model_token: Some(sd_weight_token), + maybe_lora_model_token: Some(lora_token), + maybe_prompt: Some(request.maybe_prompt.clone().unwrap_or_default()), + maybe_n_prompt: Some(n_prompt), + maybe_seed: Some(seed), + maybe_upload_path: Some(upload_path), + maybe_lora_upload_path: Some(lora_upload_path), + type_of_inference, + maybe_cfg_scale: Some(cfg_scale), + maybe_number_of_samples: Some(number_of_samples), + maybe_batch_count: Some(batch_count), + maybe_width: Some(width), + maybe_height: Some(height), + maybe_sampler: Some(sampler), + description: Some(description), + name: Some(name), + }; + + // create the inference args here + let maybe_avt_token = server_state.avt_cookie_manager.get_avt_token_from_request(&http_request); + + // create the job record here! + let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { + uuid_idempotency_token: &request.uuid_idempotency_token, + inference_category: InferenceCategory::ImageGeneration, + maybe_model_type: Some(InferenceModelType::StableDiffusion), // NB: Model is static during inference + maybe_model_token: None, // NB: Model is static during inference + maybe_input_source_token: None, + maybe_input_source_token_type: None, + maybe_raw_inference_text: None, + maybe_max_duration_seconds: None, + maybe_inference_args: Some(GenericInferenceArgs { + inference_category: Some(InferenceCategoryAbbreviated::ImageGeneration), + args: Some(PolymorphicInferenceArgs::Ig(inference_args)), + }), + maybe_creator_user_token: maybe_user_token.as_ref(), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: visbility, + priority_level, + requires_keepalive: true, + is_debug_request, + maybe_routing_tag: maybe_routing_tag.as_deref(), + mysql_pool: &server_state.mysql_pool, + }).await; + + let job_token = match query_result { + Ok((job_token, _id)) => job_token, + Err(err) => { + warn!("New generic inference job creation DB error: {:?}", err); + return Err(EnqueueImageGenRequestError::ServerError); + } + }; + + let response: EnqueueImageGenRequestSuccessResponse = EnqueueImageGenRequestSuccessResponse { + success: true, + inference_job_token: job_token, + }; + + let body = serde_json::to_string(&response).map_err(|_e| EnqueueImageGenRequestError::ServerError)?; + + // Error handling 101 rust result type returned like so. + Ok(HttpResponse::Ok().content_type("application/json").body(body)) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt new file mode 100644 index 0000000000..67177d418f --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/json.txt @@ -0,0 +1,18 @@ +{ + "uuid_idempotency_token": "123", + "type_of_inference": "checkpoint", + "maybe_sd_model_token": "weight_dmmthavhawqc2hj7yqyemcbf8", + "maybe_lora_model_token": "weight_t7gz78fjg27m0wtw6r33gafxs", + "maybe_prompt": "a very well drawn, anime girl, with pink and purple hair sitting down on a chair relaxed.", + "maybe_a_prompt": "a anime girl, with pink and purple hair sitting down on a chair relaxed.", + "maybe_n_prompt": "nsfw, black and white, low quality, pixelated", + "maybe_seed": "-1", + "maybe_width": "512", + "maybe_height": "512", + "maybe_sampler": "DPM++ 2M SDE Heun", + "maybe_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", + "maybe_lora_upload_path": "https://drive.google.com/file/d/1WRgR2pn0Ky8ls5_9Zq6tQlHTBvyWeach/view?usp=sharing", + "maybe_cfg_scale": "7", + "maybe_number_of_samples": "32", + "maybe_batch_count": "4" +} \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs new file mode 100644 index 0000000000..460b32b57c --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/image_gen/mod.rs @@ -0,0 +1,2 @@ +pub mod enqueue_image_generation; +pub mod enqueue_generation_type; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs index 3de2c08b05..1332ac1790 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/get_inference_job_status.rs @@ -318,8 +318,11 @@ fn record_to_payload( } } } + InferenceCategory::ImageGeneration => { + "".to_string() + } }; - + ResultDetailsResponse { entity_type: result_details.entity_type, entity_token: result_details.entity_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs index b5be212aa0..883fc01212 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_featured_media_files.rs @@ -15,6 +15,7 @@ use enums::by_table::media_files::media_file_origin_product_category::MediaFileO use enums::by_table::media_files::media_file_type::MediaFileType; use mysql_queries::queries::media_files::list::list_media_files_by_tokens::list_media_files_by_tokens; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -37,11 +38,21 @@ pub struct MediaFile { /// URL to the media file pub public_bucket_path: String, + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_product_category: MediaFileOriginProductCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_token: Option, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, + pub maybe_creator: Option, /// Statistics about the media file @@ -142,6 +153,12 @@ pub async fn list_featured_media_files_handler( token: m.token, media_type: m.media_type, public_bucket_path, + origin: MediaFileOriginDetails::from_db_fields_str( + m.origin_category, + m.origin_product_category, + m.maybe_origin_model_type, + m.maybe_origin_model_token.as_deref(), + m.maybe_origin_model_title.as_deref()), origin_category: m.origin_category, origin_product_category: m.origin_product_category, maybe_origin_model_type: m.maybe_origin_model_type, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs index c999b90112..ddda36c3ff 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files.rs @@ -18,6 +18,7 @@ use enums::common::visibility::Visibility; use mysql_queries::queries::media_files::list::list_media_files::{list_media_files, ListMediaFilesArgs}; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::pagination_cursors::PaginationCursors; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -44,13 +45,22 @@ pub struct ListMediaFilesSuccessResponse { pub struct MediaFileListItem { pub token: MediaFileToken, + pub media_type: MediaFileType, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub origin_product_category: MediaFileOriginProductCategory, + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] pub maybe_origin_model_token: Option, - pub media_type: MediaFileType, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, /// URL to the media file. pub public_bucket_path: String, @@ -193,11 +203,17 @@ pub async fn list_media_files_handler( let results = results_page.records.into_iter() .map(|record| MediaFileListItem { token: record.token.clone(), + media_type: record.media_type, + origin: MediaFileOriginDetails::from_db_fields_str( + record.origin_category, + record.origin_product_category, + record.maybe_origin_model_type, + record.maybe_origin_model_token.as_deref(), + record.maybe_origin_model_title.as_deref()), origin_category: record.origin_category, origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, - media_type: record.media_type, public_bucket_path: MediaFileBucketPath::from_object_hash( &record.public_bucket_directory_hash, record.maybe_public_bucket_prefix.as_deref(), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs index 87fcecca7d..845087a2f8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/list_media_files_for_user.rs @@ -17,6 +17,7 @@ use enums::common::view_as::ViewAs; use enums::common::visibility::Visibility; use mysql_queries::queries::media_files::list::list_media_files_for_user::{list_media_files_for_user, ListMediaFileForUserArgs}; use tokens::tokens::media_files::MediaFileToken; +use crate::http_server::common_responses::media_file_origin_details::MediaFileOriginDetails; use crate::http_server::common_responses::pagination_page::PaginationPage; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; @@ -47,13 +48,23 @@ pub struct ListMediaFilesForUserSuccessResponse { pub struct MediaFileForUserListItem { pub token: MediaFileToken, - pub origin_category: MediaFileOriginCategory, - pub origin_product_category: MediaFileOriginProductCategory, - pub maybe_origin_model_type: Option, - pub maybe_origin_model_token: Option, - pub media_type: MediaFileType, + /// Details where the media file came from. + pub origin: MediaFileOriginDetails, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub origin_category: MediaFileOriginCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub origin_product_category: MediaFileOriginProductCategory, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub maybe_origin_model_type: Option, + + #[deprecated(note="Use MediaFileOriginDetails instead")] + pub maybe_origin_model_token: Option, + /// URL to the media file. pub public_bucket_path: String, @@ -168,11 +179,17 @@ pub async fn list_media_files_for_user_handler( let results = results_page.records.into_iter() .map(|record| MediaFileForUserListItem { token: record.token, + media_type: record.media_type, + origin: MediaFileOriginDetails::from_db_fields_str( + record.origin_category, + record.origin_product_category, + record.maybe_origin_model_type, + record.maybe_origin_model_token.as_deref(), + record.maybe_origin_model_title.as_deref()), origin_category: record.origin_category, origin_product_category: record.origin_product_category, maybe_origin_model_type: record.maybe_origin_model_type, maybe_origin_model_token: record.maybe_origin_model_token, - media_type: record.media_type, public_bucket_path: MediaFileBucketPath::from_object_hash( &record.public_bucket_directory_hash, record.maybe_public_bucket_prefix.as_deref(), diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs index f63e4eadaa..b0e5017e15 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/media_files/upload/process_upload_media_file.rs @@ -144,7 +144,7 @@ pub async fn process_upload_media_file( info!("Upload maybe filesize: {:?}", maybe_file_size_bytes); - let maybe_mimetype = upload_media_request.file_bytes + let mut maybe_mimetype = upload_media_request.file_bytes .as_ref() .map(|bytes| get_mimetype_for_bytes(bytes.as_ref())) .flatten(); @@ -256,6 +256,16 @@ pub async fn process_upload_media_file( } } + if media_file_type.is_none() && maybe_mimetype.is_none() { + // https://research.cs.wisc.edu/graphics/Courses/cs-838-1999/Jeff/BVH.html + const BVH_HEADER : &[u8] = "HIERARCHY".as_bytes(); + + if bytes.starts_with(BVH_HEADER) { + media_file_type = Some(MediaFileType::Mocap); + maybe_mimetype = Some("application/octet-stream"); + } + } + let media_file_type = match media_file_type { Some(m) => m, None => { @@ -279,9 +289,13 @@ pub async fn process_upload_media_file( }, }; - let extension = mimetype_to_extension(mime_type) + let mut extension = mimetype_to_extension(mime_type) .map(|extension| format!(".{extension}")); + if extension.is_none() && media_file_type == MediaFileType::Mocap { + extension = Some(".bvh".to_string()); + } + const PREFIX : Option<&str> = Some("upload_"); let public_upload_path = MediaFileBucketPath::generate_new(PREFIX, extension.as_deref()); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs index f53dd3a1c6..c9b64c5dac 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/mod.rs @@ -27,5 +27,6 @@ pub mod voice_designer; pub mod w2l; pub mod user_bookmarks; pub mod weights; +pub mod image_gen; pub(crate) mod mocap; -pub(crate) mod workflows; \ No newline at end of file +pub(crate) mod workflows; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs index 94a3854335..f3fbfdbe8b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/tts/list_tts_models.rs @@ -16,10 +16,11 @@ use log::{error, info, warn}; use sqlx::MySql; use sqlx::pool::PoolConnection; +use enums::by_table::tts_models::tts_model_type::TtsModelType; use enums::common::visibility::Visibility; use errors::AnyhowResult; +use migration::text_to_speech::list_tts_models_for_migration::list_tts_models_for_migration; use mysql_queries::queries::tts::tts_category_assignments::fetch_and_build_tts_model_category_map::fetch_and_build_tts_model_category_map_with_connection; -use mysql_queries::queries::tts::tts_models::list_tts_models::list_tts_models_with_connection; use mysql_queries::queries::users::user_sessions::get_user_session_by_token::SessionUserRecord; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; @@ -133,7 +134,7 @@ pub async fn list_tts_models_handler( })?; // TODO: Fail open in case the DB is down. Pull from expired cache if query fails. - let models = get_all_models(&mut mysql_connection) + let models = get_all_models(&mut mysql_connection, server_state.flags.switch_tts_to_model_weights) .await .map_err(|e| { error!("Error querying database: {:?}", e); @@ -212,47 +213,44 @@ pub fn render_response_payload(response: ListTtsModelsSuccessResponse) -> Result Ok(body) } -async fn get_all_models(mysql_connection: &mut PoolConnection) -> AnyhowResult> { - let mut models = list_tts_models_with_connection( - mysql_connection, - None, - false - ).await?; +async fn get_all_models(mysql_connection: &mut PoolConnection, use_weights_table: bool) -> AnyhowResult> { + let mut models = list_tts_models_for_migration(mysql_connection, use_weights_table).await?; + + // Make the list nice for human readers. + models.sort_by(|a, b| + natural_lexical_cmp(&a.title(), &b.title())); let model_categories_map = fetch_and_build_tts_model_category_map_with_connection(mysql_connection).await?; - // Make the list nice for human readers. - models.sort_by(|a, b| - natural_lexical_cmp(&a.title, &b.title)); - let models_for_response = models.into_iter() .map(|model| { - let model_token = model.model_token.clone(); + // NB: All the cloning is just for the migration of tts_models --> model_weights + let model_token = model.token().to_string(); TtsModelRecordForResponse { - model_token: model.model_token, - tts_model_type: model.tts_model_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, - ietf_language_tag: model.ietf_language_tag, - ietf_primary_language_subtag: model.ietf_primary_language_subtag, - is_front_page_featured: model.is_front_page_featured, - is_twitch_featured: model.is_twitch_featured, - maybe_suggested_unique_bot_command: model.maybe_suggested_unique_bot_command, - creator_set_visibility: model.creator_set_visibility, + model_token: model.token().to_string(), + tts_model_type: TtsModelType::Tacotron2.to_string(), // NB(bt): We only ever supported the TT2 value + creator_user_token: model.creator_user_token().to_string(), + creator_username: model.creator_username().to_string(), + creator_display_name: model.creator_display_name().to_string(), + creator_gravatar_hash: model.creator_gravatar_hash().to_string(), + title: model.title().to_string(), + ietf_language_tag: model.ietf_language_tag().to_string(), + ietf_primary_language_subtag: model.ietf_primary_language_subtag().to_string(), + is_front_page_featured: model.is_front_page_featured(), + is_twitch_featured: false, // NB: This is no longer used. + maybe_suggested_unique_bot_command: None, // NB: This is no longer used. + creator_set_visibility: model.creator_set_visibility(), user_ratings: UserRatingsStats { - positive_count: model.user_ratings_positive_count, - negative_count: model.user_ratings_negative_count, - total_count: model.user_ratings_total_count, + positive_count: model.user_ratings_positive_count(), + negative_count: model.user_ratings_negative_count(), + total_count: model.user_ratings_total_count(), }, category_tokens: model_categories_map.model_to_category_tokens.get(&model_token) .map(|hash| hash.clone()) .unwrap_or(HashSet::new()), - created_at: model.created_at, - updated_at: model.updated_at, + created_at: model.created_at().clone(), + updated_at: model.updated_at().clone(), } }) .collect::>(); diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs index 212d5b4e36..7f1ba945e6 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/batch_get_user_bookmarks_handler.rs @@ -9,7 +9,7 @@ use utoipa::{IntoParams, ToSchema}; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use http_server_common::response::serialize_as_json_error::serialize_as_json_error; -use mysql_queries::queries::user_bookmarks::batch_get_user_bookmarks::{batch_get_user_bookmarks, BatchUserBookmark}; +use mysql_queries::queries::users::user_bookmarks::batch_get_user_bookmarks::{batch_get_user_bookmarks, BatchUserBookmark}; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::tts_models::TtsModelToken; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs index 638ec1500b..49e9e391da 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/create_user_bookmark_handler.rs @@ -16,10 +16,10 @@ use utoipa::ToSchema; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; use mysql_queries::queries::entity_stats::stats_entity_token::StatsEntityToken; use mysql_queries::queries::entity_stats::upsert_entity_stats_on_bookmark_event::{BookmarkAction, upsert_entity_stats_on_bookmark_event, UpsertEntityStatsArgs}; -use mysql_queries::queries::user_bookmarks::upsert_user_bookmark::{upsert_user_bookmark, CreateUserBookmarkArgs}; -use mysql_queries::queries::user_bookmarks::get_total_bookmark_count_for_entity::get_total_bookmark_count_for_entity; -use mysql_queries::queries::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; -use mysql_queries::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use mysql_queries::queries::users::user_bookmarks::upsert_user_bookmark::{upsert_user_bookmark, CreateUserBookmarkArgs}; +use mysql_queries::queries::users::user_bookmarks::get_total_bookmark_count_for_entity::get_total_bookmark_count_for_entity; +use mysql_queries::queries::users::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; +use mysql_queries::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::model_weights::ModelWeightToken; use tokens::tokens::tts_models::TtsModelToken; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs index 5b3fd42eb1..dbadd9ddf8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/delete_user_bookmark_handler.rs @@ -16,8 +16,8 @@ use utoipa::ToSchema; use mysql_queries::queries::entity_stats::stats_entity_token::StatsEntityToken; use mysql_queries::queries::entity_stats::upsert_entity_stats_on_bookmark_event::{BookmarkAction, upsert_entity_stats_on_bookmark_event, UpsertEntityStatsArgs}; -use mysql_queries::queries::user_bookmarks::delete_user_bookmark::delete_user_bookmark; -use mysql_queries::queries::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; +use mysql_queries::queries::users::user_bookmarks::delete_user_bookmark::delete_user_bookmark; +use mysql_queries::queries::users::user_bookmarks::get_user_bookmark_transactional_locking::{BookmarkIdentifier, get_user_bookmark_transactional_locking}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs index abeb7fa0ca..b69dbf6a1e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_entity_handler.rs @@ -15,8 +15,8 @@ use log::warn; use utoipa::ToSchema; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks_for_entity::list_user_bookmarks_for_entity; +use mysql_queries::queries::users::user_bookmarks::user_bookmark_entity_token::UserBookmarkEntityToken; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks_for_entity::list_user_bookmarks_for_entity; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::common_responses::user_details_lite::{DefaultAvatarInfo, UserDetailsLight}; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs index 5c4a30fe01..3b10b10df9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_session_handler.rs @@ -10,7 +10,7 @@ use chrono::{DateTime, Utc}; use log::warn; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks::{ list_user_bookmarks_by_maybe_entity_type}; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks::{ list_user_bookmarks_by_maybe_entity_type}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::web_utils::response_error_helpers::to_simple_json_error; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs index a1771118fe..0bc22eb325 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/user_bookmarks/list_user_bookmarks_for_user_handler.rs @@ -19,7 +19,7 @@ use enums::by_table::media_files::media_file_type::MediaFileType; use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::by_table::user_bookmarks::user_bookmark_entity_type::UserBookmarkEntityType; -use mysql_queries::queries::user_bookmarks::list_user_bookmarks::{list_user_bookmarks_by_maybe_entity_type, ListUserBookmarksForUserArgs}; +use mysql_queries::queries::users::user_bookmarks::list_user_bookmarks::{list_user_bookmarks_by_maybe_entity_type, ListUserBookmarksForUserArgs}; use tokens::tokens::user_bookmarks::UserBookmarkToken; use crate::http_server::common_responses::pagination_page::PaginationPage; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs index 7e09f40782..5a8e7bb3f0 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/inference/enqueue_tts_request.rs @@ -189,7 +189,7 @@ pub async fn enqueue_tts_request( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs index 6654843007..f8dd1274b9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/create_voice.rs @@ -177,7 +177,7 @@ pub async fn create_voice_handler( let query_result = insert_generic_inference_job(InsertGenericInferenceArgs { uuid_idempotency_token: &request.uuid_idempotency_token, inference_category: InferenceCategory::TextToSpeech, - maybe_model_type: Some(InferenceModelType::VallEX), // NB: Model is static during inference + maybe_model_type: Some(InferenceModelType::StyleTTS2), // NB: Model is static during inference maybe_model_token: None, // NB: Model is static during inference maybe_input_source_token: None, maybe_input_source_token_type: None, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs index eb8efa27e5..28f1500449 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/get_weight.rs @@ -15,6 +15,7 @@ use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; use mysql_queries::queries::model_weights::get_weight::get_weight_by_token; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -40,8 +41,12 @@ pub struct GetWeightResponse { file_size_bytes: i32, file_checksum_sha2: String, + /// Information about the cover image. + cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -149,6 +154,13 @@ pub async fn get_weight_handler( } } + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &weight.token, + weight.maybe_cover_image_public_bucket_hash.as_deref(), + weight.maybe_cover_image_public_bucket_prefix.as_deref(), + weight.maybe_cover_image_public_bucket_extension.as_deref(), + ); + let maybe_cover_image = weight.maybe_cover_image_public_bucket_hash .as_deref() .map(|hash| { @@ -176,6 +188,7 @@ pub async fn get_weight_handler( // TODO(bt,2023-12-24): Migrated the column. We should return nullable fields, but I don't want to break the frontend description_markdown: weight.maybe_description_markdown.unwrap_or_else(|| "".to_string()), description_rendered_html: weight.maybe_description_rendered_html.unwrap_or_else(|| "".to_string()), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, creator, creator_set_visibility: weight.creator_set_visibility, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs index 91c86066e5..fa2a42e6c3 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_featured_weights.rs @@ -15,6 +15,7 @@ use enums::by_table::model_weights::{ }; use mysql_queries::queries::model_weights::list::list_weights_by_tokens::list_weights_by_tokens; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -37,8 +38,12 @@ pub struct ModelWeightForList { pub creator: UserDetailsLight, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] pub maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -128,6 +133,12 @@ pub async fn list_featured_weights_handler( success: true, results: weights.into_iter() .map(|w| { + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &w.token, + w.maybe_cover_image_public_bucket_hash.as_deref(), + w.maybe_cover_image_public_bucket_prefix.as_deref(), + w.maybe_cover_image_public_bucket_extension.as_deref(), + ); let maybe_cover_image = w.maybe_cover_image_public_bucket_hash .as_deref() @@ -145,6 +156,7 @@ pub async fn list_featured_weights_handler( title: w.title, weight_type: w.weights_type, weight_category: w.weights_category, + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, creator: UserDetailsLight::from_db_fields( &w.creator_user_token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs index 6347b81d27..b73e722da8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/list_weights_by_user.rs @@ -16,6 +16,7 @@ use enums::common::view_as::ViewAs; use enums::common::visibility::Visibility; use mysql_queries::queries::model_weights::list::list_weights_by_user::{list_weights_by_creator_username, ListWeightsForUserArgs}; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::pagination_page::PaginationPage; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; @@ -39,8 +40,12 @@ pub struct Weight { file_size_bytes: i32, file_checksum_sha2: String, + /// Information about the cover image. + cover_image: CoverImageDetails, + /// Cover images are small descriptive images that can be set for any model. /// If a cover image is set, this is the path to the asset. + #[deprecated(note="switch to CoverImageDetails")] maybe_cover_image_public_bucket_path: Option, /// Statistics about the weights @@ -179,7 +184,14 @@ pub async fn list_weights_by_user_handler( } }; - let weights:Vec = results_page.records.into_iter().map(|weight| { + let weights: Vec = results_page.records.into_iter().map(|weight| { + + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &weight.token, + weight.maybe_cover_image_public_bucket_hash.as_deref(), + weight.maybe_cover_image_public_bucket_prefix.as_deref(), + weight.maybe_cover_image_public_bucket_extension.as_deref(), + ); let maybe_cover_image = weight.maybe_cover_image_public_bucket_hash .as_deref() @@ -203,6 +215,7 @@ pub async fn list_weights_by_user_handler( &weight.creator_display_name, &weight.creator_email_gravatar_hash, ), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, file_size_bytes: weight.file_size_bytes, file_checksum_sha2: weight.file_checksum_sha2, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs index f8fb80f89f..e27ece3566 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs @@ -20,6 +20,7 @@ use enums::by_table::model_weights::weights_category::WeightsCategory; use enums::by_table::model_weights::weights_types::WeightsType; use enums::common::visibility::Visibility; use tokens::tokens::model_weights::ModelWeightToken; +use crate::http_server::common_responses::cover_image_details::CoverImageDetails; use crate::http_server::common_responses::simple_entity_stats::SimpleEntityStats; use crate::http_server::common_responses::user_details_lite::UserDetailsLight; @@ -35,8 +36,8 @@ pub struct SearchModelWeightsRequest { pub struct ModelWeightSearchResult { pub weight_token: ModelWeightToken, - pub weights_type: WeightsType, - pub weights_category: WeightsCategory, + pub weight_type: WeightsType, + pub weight_category: WeightsCategory, pub creator_set_visibility: Visibility, @@ -44,6 +45,10 @@ pub struct ModelWeightSearchResult { pub creator: UserDetailsLight, + /// Information about the cover image. + pub cover_image: CoverImageDetails, + + #[deprecated(note="switch to CoverImageDetails")] pub maybe_cover_image_public_bucket_path: Option, pub stats: SimpleEntityStats, @@ -117,6 +122,13 @@ pub async fn search_model_weights_handler( let results = results.into_iter() .map(|result| { + let cover_image_details = CoverImageDetails::from_optional_db_fields( + &result.token, + result.maybe_cover_image_public_bucket_hash.as_deref(), + result.maybe_cover_image_public_bucket_prefix.as_deref(), + result.maybe_cover_image_public_bucket_extension.as_deref(), + ); + let maybe_cover_image = result.maybe_cover_image_public_bucket_hash .as_deref() .map(|hash| { @@ -130,8 +142,8 @@ pub async fn search_model_weights_handler( ModelWeightSearchResult { weight_token: result.token, - weights_type: result.weights_type, - weights_category: result.weights_category, + weight_type: result.weights_type, + weight_category: result.weights_category, title: result.title, creator: UserDetailsLight::from_db_fields( &result.creator_user_token, @@ -139,6 +151,7 @@ pub async fn search_model_weights_handler( &result.creator_display_name, &result.creator_gravatar_hash, ), + cover_image: cover_image_details, maybe_cover_image_public_bucket_path: maybe_cover_image, stats: SimpleEntityStats { positive_rating_count: result.ratings_positive_count, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs index 89e0ed7d3b..8ff27f82bc 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/workflows/enqueue_comfy_ui.rs @@ -10,15 +10,13 @@ use actix_web::{HttpRequest, HttpResponse, web}; use actix_web::error::ResponseError; use actix_web::http::StatusCode; use log::{info, warn}; -use buckets::public::media_files::bucket_file_path::MediaFileBucketPath; - use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType; use enums::common::visibility::Visibility; use http_server_common::request::get_request_header_optional::get_request_header_optional; use http_server_common::request::get_request_ip::get_request_ip; use mysql_queries::payloads::generic_inference_args::generic_inference_args::{GenericInferenceArgs, InferenceCategoryAbbreviated, PolymorphicInferenceArgs}; -use mysql_queries::payloads::generic_inference_args::workflow_payload::{WorkflowArgs}; +use mysql_queries::payloads::generic_inference_args::workflow_payload::{NewValue, WorkflowArgs}; use mysql_queries::queries::generic_inference::web::insert_generic_inference_job::{insert_generic_inference_job, InsertGenericInferenceArgs}; use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::model_weights::ModelWeightToken; @@ -44,6 +42,7 @@ pub struct EnqueueComfyRequest { maybe_lora_model: Option, maybe_prompt: Option, maybe_negative_prompt: Option, + maybe_json_modifications: Option>, maybe_workflow_config: Option, creator_set_visibility: Option, @@ -179,6 +178,7 @@ pub async fn enqueue_comfy_ui_handler( maybe_lora_model: request.maybe_lora_model.clone(), maybe_prompt: request.maybe_prompt.clone(), maybe_negative_prompt: request.maybe_negative_prompt.clone(), + maybe_json_modifications: request.maybe_json_modifications.clone(), maybe_workflow_config: request.maybe_workflow_config.clone(), }; diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 457b4dc148..67bfaaac05 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -359,6 +359,7 @@ async fn main() -> AnyhowResult<()> { // Temporary flags enable_enqueue_generic_tts_job: easyenv::get_env_bool_or_default("FF_ENABLE_ENQUEUE_GENERIC_TTS_JOB", false), switch_voice_conversion_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_VOICE_CONVERSION_TO_MODEL_WEIGHTS", false), + switch_tts_to_model_weights: easyenv::get_env_bool_or_default("FF_SWITCH_TTS_TO_MODEL_WEIGHTS", false), }; let third_party_url_redirector = ThirdPartyUrlRedirector::new(server_environment); diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index 1a65c0d234..0dc5ece407 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -167,6 +167,7 @@ use crate::http_server::endpoints::weights::list_weights_by_user::list_weights_b use crate::http_server::endpoints::weights::search_model_weights_handler::search_model_weights_handler; use crate::http_server::endpoints::weights::set_model_weight_cover_image::set_model_weight_cover_image_handler; use crate::http_server::endpoints::weights::update_weight::update_weight_handler; +use crate::http_server::endpoints::image_gen::enqueue_image_generation::enqueue_image_generation_request; pub fn add_routes (app: App, server_environment: ServerEnvironment) -> App where @@ -199,12 +200,11 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> app = add_user_rating_routes(app); /* /v1/user_rating/... */ app = add_subscription_routes(app); /* /v1/subscriptions/... */ app = add_voice_designer_routes(app); /* /v1/voice_designer */ - app = add_weights_routes(app); /* /v1/weights/... */ - - //if server_environment == ServerEnvironment::Development { - // // ... - //} - + app = add_image_gen_routes(app); + app = add_weights_routes(app); + // if server_environment == ServerEnvironment::Development { + // + // } // ==================== Comments ==================== let mut app = RouteBuilder::from_app(app) @@ -224,6 +224,7 @@ pub fn add_routes (app: App, server_environment: ServerEnvironment) -> .add_get("/v1/user_bookmarks/list/entity/{entity_type}/{entity_token}", list_user_bookmarks_for_entity_handler) .into_app(); + // ==================== Animations ==================== let mut app = RouteBuilder::from_app(app) @@ -1279,6 +1280,32 @@ fn add_voice_designer_routes (app:App)-> App ) } + +fn add_image_gen_routes (app:App)-> App + where + B: MessageBody, + T: ServiceFactory< + ServiceRequest, + Config = (), + Response = ServiceResponse, + Error = Error, + InitError = (), + >, +{ + // + app.service( + web::scope("/v1/image_gen") + .service( + web::scope("/model") + .route("/upload", web::post().to(enqueue_image_generation_request)) + ) + .service( + web::scope("/inference") + .route("/enqueue_image_gen", web::post().to(enqueue_image_generation_request)) + ) + ) +} + // ==================== Weights ROUTES ==================== fn add_weights_routes(app: App) -> App where diff --git a/crates/service/web/storyteller_web/src/server_state.rs b/crates/service/web/storyteller_web/src/server_state.rs index 38264e89c2..2d9e96693b 100644 --- a/crates/service/web/storyteller_web/src/server_state.rs +++ b/crates/service/web/storyteller_web/src/server_state.rs @@ -263,6 +263,11 @@ pub struct StaticFeatureFlags { /// TEMPORARY: Move voice control model listing over to `model_weights` from `voice_conversion_models` /// This will control all downstream enqueuing, jobs, etc. pub switch_voice_conversion_to_model_weights: bool, + + // TODO(2024-01-13): Remove temporary flag when done. + /// TEMPORARY: Move voice control model listing over to `model_weights` from `tts_models` + /// This will control all downstream enqueuing, jobs, etc. + pub switch_tts_to_model_weights: bool, } /// Instead of top level service denial, these are bans against entities that instead return diff --git a/crates/service/web/storyteller_web/src/threads/poll_model_token_info_thread.rs b/crates/service/web/storyteller_web/src/threads/poll_model_token_info_thread.rs index c7dc38515f..724e6e5988 100644 --- a/crates/service/web/storyteller_web/src/threads/poll_model_token_info_thread.rs +++ b/crates/service/web/storyteller_web/src/threads/poll_model_token_info_thread.rs @@ -122,6 +122,7 @@ async fn query_model_weight_models(mysql_pool: &MySqlPool) WeightsType::StableDiffusionXL => continue, // TODO(bt,2023-12-18): Not yet set up with enums for inference WeightsType::LoRA => continue, // TODO(bt,2023-12-18): Not yet set up with enums for inference WeightsType::VallE => continue, // TODO(bt,2023-12-18): ??? We have vall-e weights?? Not sure what this means yet. + WeightsType::ComfyUi => continue, // Not supported for inference }; let info = ModelInfoForInferenceJob { diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs new file mode 100644 index 0000000000..bac7ff1806 --- /dev/null +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_color_from_token.rs @@ -0,0 +1,73 @@ +use std::collections::hash_map::DefaultHasher; +use std::hash::{Hash, Hasher}; + +use tokens::tokens::model_weights::ModelWeightToken; + +/// We'll return an index from [0, NUMBER_OF_COLORS) - [inclusive, exclusive) +const NUMBER_OF_COLORS: u64 = 9; + +/// Not that it matters, but this perturbs the hash. +const SALT_LIKE_OFFSET : u8 = 51; + +/// We return an index instead of a filename, that way the frontend can drive. +/// The hash should be stable with respect to username. +pub fn default_cover_image_color_from_token(token: &ModelWeightToken) -> u8 { + // Not as important as with usernames, but seems like a reasonable choice. + let token_string = token.as_str().to_lowercase(); + + let mut hasher = DefaultHasher::new(); + + token_string.hash(&mut hasher); + SALT_LIKE_OFFSET.hash(&mut hasher); + + let hash = hasher.finish(); + + let cover_image_color_index = hash % NUMBER_OF_COLORS; + cover_image_color_index as u8 +} + +#[cfg(test)] +mod tests { + use std::collections::HashSet; + + use rand::distributions::{Alphanumeric, DistString}; + + use tokens::tokens::model_weights::ModelWeightToken; + + use crate::util::placeholder_images::cover_images::default_cover_image_color_from_token::default_cover_image_color_from_token; + + #[test] + fn test_stability() { + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("TOKEN")), 3); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("foo")), 6); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("bar")), 1); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("baz")), 6); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("bin")), 1); + } + + #[test] + fn test_case_insensitivity() { + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_TOKen")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("T_tOkEn")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_ToKeN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_toKEN")), 7); + assert_eq!(default_cover_image_color_from_token(&ModelWeightToken::new_from_str("t_token")), 7); + } + + #[test] + fn test_range() { + let mut distribution = HashSet::new(); + for _ in 0..1000 { + let random_token = Alphanumeric.sample_string(&mut rand::thread_rng(), 16); + let random_token= ModelWeightToken::new_from_str(&random_token); + let avatar_id = default_cover_image_color_from_token(&random_token); + distribution.insert(avatar_id); + assert!(avatar_id >= 0); + assert!(avatar_id <= 8); + } + // NB: We could test frequency of the distribution too + assert_eq!(distribution.len(), 9); + } +} diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs index 2cf0b81211..711eeacbf1 100644 --- a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/default_cover_image_from_token.rs @@ -3,9 +3,7 @@ use std::hash::{Hash, Hasher}; use tokens::tokens::model_weights::ModelWeightToken; -// TODO(bt,2024-01-11): How many cover images are there? - -/// There are currently 25 avatars numbered 0 to 24 (0-indexed). +/// There are currently 25 cover images numbered 0 to 24 (0-indexed). /// The original dataset was numbered 1 - 25, but I renamed 25 to 0. const NUMBER_OF_COVER_IMAGES: u64 = 25; diff --git a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs index 574d4791c4..75c63a82a6 100644 --- a/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs +++ b/crates/service/web/storyteller_web/src/util/placeholder_images/cover_images/mod.rs @@ -1 +1,2 @@ +pub mod default_cover_image_color_from_token; pub mod default_cover_image_from_token; \ No newline at end of file