diff --git a/Cargo.lock b/Cargo.lock index e0d440b106..e670827390 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -11603,6 +11603,7 @@ dependencies = [ "utoipa-swagger-ui", "videos", "workspace-build-acceleration", + "world_labs_api", "wreq", "zip 0.6.6", ] diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/mod.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/mod.rs index b0a31af1c3..414181138d 100644 --- a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/mod.rs +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/mod.rs @@ -1,4 +1,5 @@ pub mod cost_estimate; pub mod image; pub mod object; +pub mod splat; pub mod video; \ No newline at end of file diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs new file mode 100644 index 0000000000..0048640caf --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs @@ -0,0 +1,25 @@ +use serde_derive::{Deserialize, Serialize}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use utoipa::ToSchema; + +pub const GENERATE_WORLDLABS_MARBLE_0P1_MINI_SPLAT_URL_PATH: &str = "/v1/generate/splat/worldlabs_marble_0p1_mini"; + +/// World Labs Marble 0.1-mini Splat Generation +#[derive(Serialize, Deserialize, ToSchema)] +pub struct GenerateWorldlabsMarble0p1MiniSplatRequest { + /// Idempotency token to prevent duplicate requests. + pub uuid_idempotency_token: String, + + /// Optional image to use as input for world generation. + pub image_media_file_token: Option, + + /// Optional text prompt for world generation. + pub prompt: Option, +} + +#[derive(Serialize, Deserialize, ToSchema)] +pub struct GenerateWorldlabsMarble0p1MiniSplatResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs new file mode 100644 index 0000000000..bed7ea606e --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs @@ -0,0 +1,25 @@ +use serde_derive::{Deserialize, Serialize}; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use utoipa::ToSchema; + +pub const GENERATE_WORLDLABS_MARBLE_0P1_PLUS_SPLAT_URL_PATH: &str = "/v1/generate/splat/worldlabs_marble_0p1_plus"; + +/// World Labs Marble 0.1-plus Splat Generation +#[derive(Serialize, Deserialize, ToSchema)] +pub struct GenerateWorldlabsMarble0p1PlusSplatRequest { + /// Idempotency token to prevent duplicate requests. + pub uuid_idempotency_token: String, + + /// Optional image to use as input for world generation. + pub image_media_file_token: Option, + + /// Optional text prompt for world generation. + pub prompt: Option, +} + +#[derive(Serialize, Deserialize, ToSchema)] +pub struct GenerateWorldlabsMarble0p1PlusSplatResponse { + pub success: bool, + pub inference_job_token: InferenceJobToken, +} diff --git a/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/mod.rs b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/mod.rs new file mode 100644 index 0000000000..68557f5370 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_api_defs/src/generate/splat/mod.rs @@ -0,0 +1,2 @@ +pub mod generate_worldlabs_marble_0p1_mini_splat; +pub mod generate_worldlabs_marble_0p1_plus_splat; diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/mod.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/mod.rs index b0a31af1c3..414181138d 100644 --- a/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/mod.rs +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/mod.rs @@ -1,4 +1,5 @@ pub mod cost_estimate; pub mod image; pub mod object; +pub mod splat; pub mod video; \ No newline at end of file diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs new file mode 100644 index 0000000000..95410177a4 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat.rs @@ -0,0 +1,18 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::basic_json_post_request::basic_json_post_request; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_mini_splat::{GenerateWorldlabsMarble0p1MiniSplatRequest, GenerateWorldlabsMarble0p1MiniSplatResponse, GENERATE_WORLDLABS_MARBLE_0P1_MINI_SPLAT_URL_PATH}; + +pub async fn generate_worldlabs_marble_0p1_mini_splat( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + request: GenerateWorldlabsMarble0p1MiniSplatRequest, +) -> Result { + Ok(basic_json_post_request( + api_host, + GENERATE_WORLDLABS_MARBLE_0P1_MINI_SPLAT_URL_PATH, + maybe_creds, + request, + ).await?) +} diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs new file mode 100644 index 0000000000..9dd97b21e3 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat.rs @@ -0,0 +1,18 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::basic_json_post_request::basic_json_post_request; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_plus_splat::{GenerateWorldlabsMarble0p1PlusSplatRequest, GenerateWorldlabsMarble0p1PlusSplatResponse, GENERATE_WORLDLABS_MARBLE_0P1_PLUS_SPLAT_URL_PATH}; + +pub async fn generate_worldlabs_marble_0p1_plus_splat( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + request: GenerateWorldlabsMarble0p1PlusSplatRequest, +) -> Result { + Ok(basic_json_post_request( + api_host, + GENERATE_WORLDLABS_MARBLE_0P1_PLUS_SPLAT_URL_PATH, + maybe_creds, + request, + ).await?) +} diff --git a/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/mod.rs b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/mod.rs new file mode 100644 index 0000000000..68557f5370 --- /dev/null +++ b/crates/api_clients/artcraft/artcraft_client/src/endpoints/generate/splat/mod.rs @@ -0,0 +1,2 @@ +pub mod generate_worldlabs_marble_0p1_mini_splat; +pub mod generate_worldlabs_marble_0p1_plus_splat; diff --git a/crates/api_clients/world_labs_api/src/api/requests/generate_world/mod.rs b/crates/api_clients/world_labs_api/src/api/requests/generate_world/mod.rs index 765d9bf619..c26446968e 100644 --- a/crates/api_clients/world_labs_api/src/api/requests/generate_world/mod.rs +++ b/crates/api_clients/world_labs_api/src/api/requests/generate_world/mod.rs @@ -1,2 +1,2 @@ pub mod generate_world; -pub(crate) mod http_request; +pub mod http_request; diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/mod.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/mod.rs index 75c08abda9..4946e447fa 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/mod.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/mod.rs @@ -2,3 +2,4 @@ pub mod fal; pub mod job; pub mod seedance2pro; pub mod web; +pub mod worldlabs; diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/batch_get_inference_job_status.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/batch_get_inference_job_status.rs index 9374d36ffb..5f3b7fa1aa 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/batch_get_inference_job_status.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/batch_get_inference_job_status.rs @@ -186,6 +186,7 @@ fn raw_records_to_public_result(records: Vec) -> V InferenceCategory::F5TTS => Some("F5 TTS"), InferenceCategory::SeedVc => Some("Seed VC"), InferenceCategory::ConvertBvhToWorkflow => Some("BVH to Workflow"), + InferenceCategory::SplatGeneration => Some("Splat Generation"), InferenceCategory::DeprecatedField => Some("Job"), // TODO(bt,2024-07-16): Fix }; } @@ -208,6 +209,7 @@ fn raw_records_to_public_result(records: Vec) -> V InferenceCategory::SeedVc => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::LivePortrait => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::ConvertBvhToWorkflow => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), + InferenceCategory::SplatGeneration => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::DeprecatedField => (true, None), // TODO(bt,2024-07-16): We'll need to read another field! }; diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs index a7d46edef1..39fcc4ea38 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/get_inference_job_status.rs @@ -162,6 +162,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::F5TTS => Some("F5 TTS"), InferenceCategory::SeedVc => Some("Seed VC"), InferenceCategory::ConvertBvhToWorkflow => Some("BVH to Workflow"), + InferenceCategory::SplatGeneration => Some("Splat Generation"), InferenceCategory::DeprecatedField => Some("Job"), // TODO(bt,2024-07-16): Fix }; } @@ -184,6 +185,7 @@ fn raw_record_to_public_result(record: RawGenericInferenceJobStatus) -> GenericI InferenceCategory::LivePortrait => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::SeedVc => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::ConvertBvhToWorkflow => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), + InferenceCategory::SplatGeneration => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::DeprecatedField => (true, None), // TODO(bt,2024-07-16): We'll need to read another field! }; diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/list_session_jobs.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/list_session_jobs.rs index 6b9f220002..9bae4be27b 100644 --- a/crates/schema/database/mysql_queries/src/queries/generic_inference/web/list_session_jobs.rs +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/web/list_session_jobs.rs @@ -252,6 +252,7 @@ fn raw_records_to_public_result(records: Vec) -> V InferenceCategory::ConvertBvhToWorkflow => Some("BVH to Workflow"), InferenceCategory::LivePortrait => Some("Live Portrait"), InferenceCategory::SeedVc => Some("Seed VC"), + InferenceCategory::SplatGeneration => Some("Splat Generation"), InferenceCategory::DeprecatedField => Some("Job"), // TODO(bt,2024-07-16): Fix }; } @@ -274,6 +275,7 @@ fn raw_records_to_public_result(records: Vec) -> V InferenceCategory::LivePortrait => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::SeedVc => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::ConvertBvhToWorkflow => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), + InferenceCategory::SplatGeneration => (true, record.maybe_media_file_public_bucket_directory_hash.as_deref()), InferenceCategory::DeprecatedField => (true, None), // TODO(bt,2024-07-16): We'll need to read another field! }; diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token.rs new file mode 100644 index 0000000000..25e53bab8e --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token.rs @@ -0,0 +1,143 @@ +use anyhow::anyhow; +use log::info; +use sqlx::{Executor, MySql}; +use std::marker::PhantomData; + +use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory; +use enums::by_table::generic_inference_jobs::inference_job_external_third_party::InferenceJobExternalThirdParty; +use enums::by_table::generic_inference_jobs::inference_job_product_category::InferenceJobProductCategory; +use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType; +use enums::common::job_status_plus::JobStatusPlus; +use enums::common::visibility::Visibility; +use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::prompts::PromptToken; +use tokens::tokens::users::UserToken; +use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken; + +use crate::errors::database_query_error::DatabaseQueryError; +use crate::payloads::generic_inference_args::generic_inference_args::GenericInferenceArgs; + +pub struct InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs<'e, 'c, E> + where E: 'e + Executor<'c, Database = MySql> +{ + pub uuid_idempotency_token: &'e str, + + pub apriori_job_token: &'e InferenceJobToken, + + /// The external operation_id returned by World Labs generate_world call. + pub maybe_external_third_party_id: &'e str, + + pub maybe_inference_args: Option, + + pub maybe_prompt_token: Option<&'e PromptToken>, + + pub maybe_wallet_ledger_entry_token: Option<&'e WalletLedgerEntryToken>, + + pub maybe_creator_user_token: Option<&'e UserToken>, + pub maybe_avt_token: Option<&'e AnonymousVisitorTrackingToken>, + pub creator_ip_address: &'e str, + pub creator_set_visibility: Visibility, + + pub mysql_executor: E, + + pub phantom: PhantomData<&'c E>, +} + +pub async fn insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token<'e, 'c : 'e, E>( + args: InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs<'e, 'c, E> +) -> Result + where E: 'e + Executor<'c, Database = MySql> +{ + let serialized_args_payload = serde_json::ser::to_string(&args.maybe_inference_args) + .map_err(|_e| anyhow!("could not encode inference args"))?; + + const JOB_TYPE: InferenceJobType = InferenceJobType::WorldlabsQueue; + const INFERENCE_CATEGORY: InferenceCategory = InferenceCategory::SplatGeneration; + const PRODUCT_CATEGORY: InferenceJobProductCategory = InferenceJobProductCategory::WorldlabsSplat; + const EXTERNAL_THIRD_PARTY: InferenceJobExternalThirdParty = InferenceJobExternalThirdParty::Worldlabs; + const STATUS: JobStatusPlus = JobStatusPlus::Pending; + + let query = sqlx::query!( + r#" +INSERT INTO generic_inference_jobs +SET + token = ?, + uuid_idempotency_token = ?, + + job_type = ?, + + maybe_external_third_party = ?, + maybe_external_third_party_id = ?, + + product_category = ?, + inference_category = ?, + + maybe_model_type = NULL, + maybe_model_token = NULL, + + maybe_input_source_token = NULL, + maybe_input_source_token_type = NULL, + + maybe_download_url = NULL, + maybe_cover_image_media_file_token = NULL, + + maybe_prompt_token = ?, + + maybe_wallet_ledger_entry_token = ?, + + maybe_raw_inference_text = NULL, + + maybe_inference_args = ?, + + maybe_creator_user_token = ?, + maybe_creator_anonymous_visitor_token = ?, + creator_ip_address = ?, + creator_set_visibility = ?, + + priority_level = 0, + is_keepalive_required = FALSE, + max_duration_seconds = 0, + + is_debug_request = FALSE, + maybe_routing_tag = NULL, + + status = ? + "#, + args.apriori_job_token.as_str(), + args.uuid_idempotency_token, + + JOB_TYPE.to_str(), + + EXTERNAL_THIRD_PARTY.to_str(), + args.maybe_external_third_party_id, + + PRODUCT_CATEGORY.to_str(), + INFERENCE_CATEGORY.to_str(), + + args.maybe_prompt_token.map(|t| t.to_string()), + + args.maybe_wallet_ledger_entry_token.map(|t| t.to_string()), + + serialized_args_payload, + + args.maybe_creator_user_token.map(|t| t.to_string()), + args.maybe_avt_token.map(|t| t.to_string()), + args.creator_ip_address, + args.creator_set_visibility.to_str(), + + STATUS.to_str(), + ); + + let query_result = query.execute(args.mysql_executor) + .await; + + let record_id = match query_result { + Err(err) => return Err(DatabaseQueryError::from(err)), + Ok(res) => res.last_insert_id(), + }; + + info!("Insert generic inference job for World Labs queue: {} with record ID {}", args.apriori_job_token, record_id); + + Ok(args.apriori_job_token.clone()) +} diff --git a/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs new file mode 100644 index 0000000000..259d0e2f25 --- /dev/null +++ b/crates/schema/database/mysql_queries/src/queries/generic_inference/worldlabs/mod.rs @@ -0,0 +1 @@ +pub mod insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token; diff --git a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_category.rs b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_category.rs index 517359484a..226a6af7db 100644 --- a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_category.rs +++ b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_category.rs @@ -47,7 +47,11 @@ pub enum InferenceCategory { /// FAL-powered 3D object generation #[serde(rename = "object_generation")] ObjectGeneration, - + + /// Gaussian Splat generation (eg. World Labs Marble) + #[serde(rename = "splat_generation")] + SplatGeneration, + /// FAL-powered image background removal #[serde(rename = "background_removal")] BackgroundRemoval, @@ -104,6 +108,7 @@ impl InferenceCategory { Self::ImageGeneration => "image_generation", Self::VideoGeneration => "video_generation", Self::ObjectGeneration => "object_generation", + Self::SplatGeneration => "splat_generation", Self::BackgroundRemoval => "background_removal", Self::Mocap => "mocap", Self::Workflow => "workflow", @@ -125,6 +130,7 @@ impl InferenceCategory { "image_generation" => Ok(Self::ImageGeneration), "video_generation" => Ok(Self::VideoGeneration), "object_generation" => Ok(Self::ObjectGeneration), + "splat_generation" => Ok(Self::SplatGeneration), "background_removal" => Ok(Self::BackgroundRemoval), "f5_tts" => Ok(Self::F5TTS), "mocap" => Ok(Self::Mocap), @@ -148,6 +154,7 @@ impl InferenceCategory { Self::VoiceConversion, Self::ImageGeneration, Self::ObjectGeneration, + Self::SplatGeneration, Self::VideoGeneration, Self::BackgroundRemoval, Self::Mocap, @@ -179,6 +186,7 @@ mod tests { assert_serialization(InferenceCategory::ImageGeneration, "image_generation"); assert_serialization(InferenceCategory::VideoGeneration, "video_generation"); assert_serialization(InferenceCategory::ObjectGeneration, "object_generation"); + assert_serialization(InferenceCategory::SplatGeneration, "splat_generation"); assert_serialization(InferenceCategory::BackgroundRemoval, "background_removal"); assert_serialization(InferenceCategory::Mocap, "mocap"); assert_serialization(InferenceCategory::F5TTS, "f5_tts"); @@ -199,6 +207,7 @@ mod tests { assert_eq!(InferenceCategory::ImageGeneration.to_str(), "image_generation"); assert_eq!(InferenceCategory::VideoGeneration.to_str(), "video_generation"); assert_eq!(InferenceCategory::ObjectGeneration.to_str(), "object_generation"); + assert_eq!(InferenceCategory::SplatGeneration.to_str(), "splat_generation"); assert_eq!(InferenceCategory::BackgroundRemoval.to_str(), "background_removal"); assert_eq!(InferenceCategory::F5TTS.to_str(), "f5_tts"); assert_eq!(InferenceCategory::SeedVc.to_str(), "seed_vc"); @@ -219,6 +228,7 @@ mod tests { assert_eq!(InferenceCategory::from_str("image_generation").unwrap(), InferenceCategory::ImageGeneration); assert_eq!(InferenceCategory::from_str("video_generation").unwrap(), InferenceCategory::VideoGeneration); assert_eq!(InferenceCategory::from_str("object_generation").unwrap(), InferenceCategory::ObjectGeneration); + assert_eq!(InferenceCategory::from_str("splat_generation").unwrap(), InferenceCategory::SplatGeneration); assert_eq!(InferenceCategory::from_str("background_removal").unwrap(), InferenceCategory::BackgroundRemoval); assert_eq!(InferenceCategory::from_str("f5_tts").unwrap(), InferenceCategory::F5TTS); assert_eq!(InferenceCategory::from_str("seed_vc").unwrap(), InferenceCategory::SeedVc); @@ -233,7 +243,7 @@ mod tests { #[test] fn all_variants() { // Static check - const EXPECTED_COUNT : usize = 16; + const EXPECTED_COUNT : usize = 17; assert_eq!(InferenceCategory::all_variants().len(), EXPECTED_COUNT); assert_eq!(InferenceCategory::iter().len(), EXPECTED_COUNT); diff --git a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_external_third_party.rs b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_external_third_party.rs index e0e9dcbc36..327a42ab39 100644 --- a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_external_third_party.rs +++ b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_external_third_party.rs @@ -20,6 +20,10 @@ pub enum InferenceJobExternalThirdParty { /// Seedance 2 Pro jobs #[serde(rename = "seedance2pro")] Seedance2Pro, + + /// World Labs jobs + #[serde(rename = "worldlabs")] + Worldlabs, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -32,6 +36,7 @@ impl InferenceJobExternalThirdParty { match self { Self::Fal => "fal", Self::Seedance2Pro => "seedance2pro", + Self::Worldlabs => "worldlabs", } } @@ -39,6 +44,7 @@ impl InferenceJobExternalThirdParty { match value { "fal" => Ok(Self::Fal), "seedance2pro" => Ok(Self::Seedance2Pro), + "worldlabs" => Ok(Self::Worldlabs), _ => Err(format!("invalid value: {:?}", value)), } } @@ -49,6 +55,7 @@ impl InferenceJobExternalThirdParty { BTreeSet::from([ Self::Fal, Self::Seedance2Pro, + Self::Worldlabs, ]) } } @@ -65,24 +72,27 @@ mod tests { fn test_serialization() { assert_serialization(InferenceJobExternalThirdParty::Fal, "fal"); assert_serialization(InferenceJobExternalThirdParty::Seedance2Pro, "seedance2pro"); + assert_serialization(InferenceJobExternalThirdParty::Worldlabs, "worldlabs"); } #[test] fn to_str() { assert_eq!(InferenceJobExternalThirdParty::Fal.to_str(), "fal"); assert_eq!(InferenceJobExternalThirdParty::Seedance2Pro.to_str(), "seedance2pro"); + assert_eq!(InferenceJobExternalThirdParty::Worldlabs.to_str(), "worldlabs"); } #[test] fn from_str() { assert_eq!(InferenceJobExternalThirdParty::from_str("fal").unwrap(), InferenceJobExternalThirdParty::Fal); assert_eq!(InferenceJobExternalThirdParty::from_str("seedance2pro").unwrap(), InferenceJobExternalThirdParty::Seedance2Pro); + assert_eq!(InferenceJobExternalThirdParty::from_str("worldlabs").unwrap(), InferenceJobExternalThirdParty::Worldlabs); } #[test] fn all_variants() { // Static check - const EXPECTED_COUNT : usize = 2; + const EXPECTED_COUNT : usize = 3; assert_eq!(InferenceJobExternalThirdParty::all_variants().len(), EXPECTED_COUNT); assert_eq!(InferenceJobExternalThirdParty::iter().len(), EXPECTED_COUNT); diff --git a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs index d4f540ccff..42fca57b78 100644 --- a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs +++ b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_product_category.rs @@ -35,6 +35,12 @@ pub enum InferenceJobProductCategory { #[serde(rename = "seedance2pro_video")] Seedance2ProVideo, + // =============== WORLD LABS =============== + + /// World Labs: Gaussian Splat Generation + #[serde(rename = "worldlabs_splat")] + WorldlabsSplat, + // =============== TEXT TO SPEECH =============== /// TTS: GptSoVits @@ -131,6 +137,7 @@ impl InferenceJobProductCategory { Self::FalObject => "fal_object", Self::FalBgRemoval => "fal_bg_removal", Self::Seedance2ProVideo => "seedance2pro_video", + Self::WorldlabsSplat => "worldlabs_splat", Self::TtsGptSoVits => "tts_gpt_so_vits", Self::TtsStyleTts2 => "tts_style_tts2", Self::TtsTacotron2 => "tts_tacotron2", @@ -164,6 +171,7 @@ impl InferenceJobProductCategory { "fal_object" => Ok(Self::FalObject), "fal_bg_removal" => Ok(Self::FalBgRemoval), "seedance2pro_video" => Ok(Self::Seedance2ProVideo), + "worldlabs_splat" => Ok(Self::WorldlabsSplat), "tts_gpt_so_vits" => Ok(Self::TtsGptSoVits), "tts_style_tts2" => Ok(Self::TtsStyleTts2), "tts_tacotron2" => Ok(Self::TtsTacotron2), @@ -200,6 +208,7 @@ impl InferenceJobProductCategory { Self::FalObject, Self::FalBgRemoval, Self::Seedance2ProVideo, + Self::WorldlabsSplat, Self::TtsGptSoVits, Self::TtsStyleTts2, Self::TtsTacotron2, @@ -242,6 +251,7 @@ mod tests { assert_serialization(InferenceJobProductCategory::FalObject, "fal_object"); assert_serialization(InferenceJobProductCategory::FalBgRemoval, "fal_bg_removal"); assert_serialization(InferenceJobProductCategory::Seedance2ProVideo, "seedance2pro_video"); + assert_serialization(InferenceJobProductCategory::WorldlabsSplat, "worldlabs_splat"); assert_serialization(InferenceJobProductCategory::TtsGptSoVits, "tts_gpt_so_vits"); assert_serialization(InferenceJobProductCategory::TtsStyleTts2, "tts_style_tts2"); assert_serialization(InferenceJobProductCategory::TtsTacotron2, "tts_tacotron2"); @@ -274,6 +284,7 @@ mod tests { assert_eq!(InferenceJobProductCategory::FalObject.to_str(), "fal_object"); assert_eq!(InferenceJobProductCategory::FalBgRemoval.to_str(), "fal_bg_removal"); assert_eq!(InferenceJobProductCategory::Seedance2ProVideo.to_str(), "seedance2pro_video"); + assert_eq!(InferenceJobProductCategory::WorldlabsSplat.to_str(), "worldlabs_splat"); assert_eq!(InferenceJobProductCategory::TtsGptSoVits.to_str(), "tts_gpt_so_vits"); assert_eq!(InferenceJobProductCategory::TtsStyleTts2.to_str(), "tts_style_tts2"); assert_eq!(InferenceJobProductCategory::TtsTacotron2.to_str(), "tts_tacotron2"); @@ -305,6 +316,7 @@ mod tests { assert_eq!(InferenceJobProductCategory::from_str("fal_object").unwrap(), InferenceJobProductCategory::FalObject); assert_eq!(InferenceJobProductCategory::from_str("fal_bg_removal").unwrap(), InferenceJobProductCategory::FalBgRemoval); assert_eq!(InferenceJobProductCategory::from_str("seedance2pro_video").unwrap(), InferenceJobProductCategory::Seedance2ProVideo); + assert_eq!(InferenceJobProductCategory::from_str("worldlabs_splat").unwrap(), InferenceJobProductCategory::WorldlabsSplat); assert_eq!(InferenceJobProductCategory::from_str("tts_gpt_so_vits").unwrap(), InferenceJobProductCategory::TtsGptSoVits); assert_eq!(InferenceJobProductCategory::from_str("tts_style_tts2").unwrap(), InferenceJobProductCategory::TtsStyleTts2); assert_eq!(InferenceJobProductCategory::from_str("tts_tacotron2").unwrap(), InferenceJobProductCategory::TtsTacotron2); @@ -331,7 +343,7 @@ mod tests { #[test] fn all_variants() { // Static check - const EXPECTED_COUNT : usize = 28; + const EXPECTED_COUNT : usize = 29; assert_eq!(InferenceJobProductCategory::all_variants().len(), EXPECTED_COUNT); assert_eq!(InferenceJobProductCategory::iter().len(), EXPECTED_COUNT); diff --git a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_type.rs b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_type.rs index 09c7b00d67..888f8a4889 100644 --- a/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_type.rs +++ b/crates/schema/public/enums/src/by_table/generic_inference_jobs/inference_job_type.rs @@ -23,6 +23,10 @@ pub enum InferenceJobType { #[serde(rename = "seedance2pro_queue")] Seedance2ProQueue, + /// World Labs jobs. We poll for results. + #[serde(rename = "worldlabs_queue")] + WorldlabsQueue, + /// Storyteller Studio and Video Style Transfer Jobs (which we may want to split). /// These run in Comfy. /// TODO(bt,2024-07-15): We may segregate these two job types in the future @@ -115,6 +119,7 @@ impl InferenceJobType { match self { Self::FalQueue => "fal_queue", Self::Seedance2ProQueue => "seedance2pro_queue", + Self::WorldlabsQueue => "worldlabs_queue", Self::VideoRender => "video_render", Self::LivePortrait => "live_portrait", Self::FaceFusion => "face_fusion", @@ -142,6 +147,7 @@ impl InferenceJobType { match value { "fal_queue" => Ok(Self::FalQueue), "seedance2pro_queue" => Ok(Self::Seedance2ProQueue), + "worldlabs_queue" => Ok(Self::WorldlabsQueue), "video_render" => Ok(Self::VideoRender), "live_portrait" => Ok(Self::LivePortrait), "face_fusion" => Ok(Self::FaceFusion), @@ -172,6 +178,7 @@ impl InferenceJobType { BTreeSet::from([ Self::FalQueue, Self::Seedance2ProQueue, + Self::WorldlabsQueue, Self::VideoRender, Self::LivePortrait, Self::FaceFusion, @@ -213,6 +220,7 @@ mod tests { fn test_serialization() { assert_serialization(InferenceJobType::FalQueue, "fal_queue"); assert_serialization(InferenceJobType::Seedance2ProQueue, "seedance2pro_queue"); + assert_serialization(InferenceJobType::WorldlabsQueue, "worldlabs_queue"); assert_serialization(InferenceJobType::VideoRender, "video_render"); assert_serialization(InferenceJobType::LivePortrait, "live_portrait"); assert_serialization(InferenceJobType::FaceFusion, "face_fusion"); @@ -239,6 +247,7 @@ mod tests { fn to_str() { assert_eq!(InferenceJobType::FalQueue.to_str(), "fal_queue"); assert_eq!(InferenceJobType::Seedance2ProQueue.to_str(), "seedance2pro_queue"); + assert_eq!(InferenceJobType::WorldlabsQueue.to_str(), "worldlabs_queue"); assert_eq!(InferenceJobType::VideoRender.to_str(), "video_render"); assert_eq!(InferenceJobType::LivePortrait.to_str(), "live_portrait"); assert_eq!(InferenceJobType::FaceFusion.to_str(), "face_fusion"); @@ -265,6 +274,7 @@ mod tests { fn from_str() { assert_eq!(InferenceJobType::from_str("fal_queue").unwrap(), InferenceJobType::FalQueue); assert_eq!(InferenceJobType::from_str("seedance2pro_queue").unwrap(), InferenceJobType::Seedance2ProQueue); + assert_eq!(InferenceJobType::from_str("worldlabs_queue").unwrap(), InferenceJobType::WorldlabsQueue); assert_eq!(InferenceJobType::from_str("video_render").unwrap(), InferenceJobType::VideoRender); assert_eq!(InferenceJobType::from_str("live_portrait").unwrap(), InferenceJobType::LivePortrait); assert_eq!(InferenceJobType::from_str("face_fusion").unwrap(), InferenceJobType::FaceFusion); @@ -290,7 +300,7 @@ mod tests { #[test] fn all_variants() { // Static check - const EXPECTED_COUNT : usize = 22; + const EXPECTED_COUNT : usize = 23; assert_eq!(InferenceJobType::all_variants().len(), EXPECTED_COUNT); assert_eq!(InferenceJobType::iter().len(), EXPECTED_COUNT); diff --git a/crates/schema/public/enums/src/common/model_type.rs b/crates/schema/public/enums/src/common/model_type.rs index d7b554b17c..caeb567333 100644 --- a/crates/schema/public/enums/src/common/model_type.rs +++ b/crates/schema/public/enums/src/common/model_type.rs @@ -127,6 +127,12 @@ pub enum ModelType { Hunyuan3d2_1, #[serde(rename = "hunyuan_3d_3")] Hunyuan3d3, + + // Splat generation models (World Labs) + #[serde(rename = "marble_0p1_mini")] + Marble0p1Mini, + #[serde(rename = "marble_0p1_plus")] + Marble0p1Plus, } impl_enum_display_and_debug_using_to_str!(ModelType); @@ -193,6 +199,10 @@ impl ModelType { Self::Hunyuan3d2_0 => "hunyuan_3d_2p0", Self::Hunyuan3d2_1 => "hunyuan_3d_2p1", Self::Hunyuan3d3 => "hunyuan_3d_3", + + // Splat generation models (World Labs) + Self::Marble0p1Mini => "marble_0p1_mini", + Self::Marble0p1Plus => "marble_0p1_plus", } } @@ -254,6 +264,10 @@ impl ModelType { "hunyuan_3d_2p1" => Ok(Self::Hunyuan3d2_1), "hunyuan_3d_3" => Ok(Self::Hunyuan3d3), + // Splat generation models (World Labs) + "marble_0p1_mini" => Ok(Self::Marble0p1Mini), + "marble_0p1_plus" => Ok(Self::Marble0p1Plus), + _ => Err(format!("invalid model_type: {:?}", job_status)), } } @@ -317,6 +331,10 @@ impl ModelType { Self::Hunyuan3d2_0, Self::Hunyuan3d2_1, Self::Hunyuan3d3, + + // Splat generation models (World Labs) + Self::Marble0p1Mini, + Self::Marble0p1Plus, ]) } } @@ -384,6 +402,9 @@ mod tests { assert_serialization(ModelType::Hunyuan3d2_0, "hunyuan_3d_2p0"); assert_serialization(ModelType::Hunyuan3d2_1, "hunyuan_3d_2p1"); assert_serialization(ModelType::Hunyuan3d3, "hunyuan_3d_3"); + // Splat generation models (World Labs) + assert_serialization(ModelType::Marble0p1Mini, "marble_0p1_mini"); + assert_serialization(ModelType::Marble0p1Plus, "marble_0p1_plus"); } #[test] @@ -443,6 +464,9 @@ mod tests { assert_eq!(ModelType::Hunyuan3d2_0.to_str(), "hunyuan_3d_2p0"); assert_eq!(ModelType::Hunyuan3d2_1.to_str(), "hunyuan_3d_2p1"); assert_eq!(ModelType::Hunyuan3d3.to_str(), "hunyuan_3d_3"); + // Splat generation models (World Labs) + assert_eq!(ModelType::Marble0p1Mini.to_str(), "marble_0p1_mini"); + assert_eq!(ModelType::Marble0p1Plus.to_str(), "marble_0p1_plus"); } #[test] @@ -501,12 +525,15 @@ mod tests { assert_eq!(ModelType::from_str("hunyuan_3d_2p0").unwrap(), ModelType::Hunyuan3d2_0); assert_eq!(ModelType::from_str("hunyuan_3d_2p1").unwrap(), ModelType::Hunyuan3d2_1); assert_eq!(ModelType::from_str("hunyuan_3d_3").unwrap(), ModelType::Hunyuan3d3); + // Splat generation models (World Labs) + assert_eq!(ModelType::from_str("marble_0p1_mini").unwrap(), ModelType::Marble0p1Mini); + assert_eq!(ModelType::from_str("marble_0p1_plus").unwrap(), ModelType::Marble0p1Plus); } #[test] fn all_variants() { let mut variants = ModelType::all_variants(); - assert_eq!(variants.len(), 50); + assert_eq!(variants.len(), 52); // Image models assert_eq!(variants.pop_first(), Some(ModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(ModelType::Flux1Schnell)); @@ -560,6 +587,9 @@ mod tests { assert_eq!(variants.pop_first(), Some(ModelType::Hunyuan3d2_0)); assert_eq!(variants.pop_first(), Some(ModelType::Hunyuan3d2_1)); assert_eq!(variants.pop_first(), Some(ModelType::Hunyuan3d3)); + // Splat generation models (World Labs) + assert_eq!(variants.pop_first(), Some(ModelType::Marble0p1Mini)); + assert_eq!(variants.pop_first(), Some(ModelType::Marble0p1Plus)); assert_eq!(variants.pop_first(), None); } diff --git a/crates/schema/public/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs b/crates/schema/public/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs index a8d7346004..c480b5ae79 100644 --- a/crates/schema/public/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs +++ b/crates/schema/public/enums_public/src/by_table/generic_inference_jobs/public_inference_job_type.rs @@ -104,6 +104,10 @@ pub enum PublicInferenceJobType { /// Seedance 2 Pro video generation jobs. We poll for results. #[serde(rename = "seedance2pro_queue")] Seedance2ProQueue, + + /// World Labs splat generation jobs. We poll for results. + #[serde(rename = "worldlabs_queue")] + WorldlabsQueue, } /// NB: Legacy API for older code. @@ -135,9 +139,10 @@ impl PublicInferenceJobType { InferenceJobType::RerenderAVideo => Self::RerenderAVideo, InferenceJobType::ImageGenApi => Self::ImageGenApi, InferenceJobType::Seedance2ProQueue => Self::Seedance2ProQueue, + InferenceJobType::WorldlabsQueue => Self::WorldlabsQueue, } } - + pub fn to_enum(&self) -> InferenceJobType { match self { // Renamed variants @@ -159,6 +164,7 @@ impl PublicInferenceJobType { Self::StableDiffusion => InferenceJobType::StableDiffusion, Self::ImageGenApi => InferenceJobType::ImageGenApi, Self::Seedance2ProQueue => InferenceJobType::Seedance2ProQueue, + Self::WorldlabsQueue => InferenceJobType::WorldlabsQueue, Self::StyleTTS2 => InferenceJobType::StyleTTS2, Self::Tacotron2 => InferenceJobType::Tacotron2, Self::Unknown => InferenceJobType::Unknown, diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index 20daf71ffd..e99a087152 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -82,6 +82,7 @@ url_config = { path = "../../../lib/url_config" } user_input_common = { path = "../../../lib/user_input_common" } users.workspace = true videos = { path = "../../../lib/files/videos" } +world_labs_api.workspace = true # Internal Components billing_component = { path = "../../plugins/billing" } diff --git a/crates/service/web/storyteller_web/src/docs/api_doc.rs b/crates/service/web/storyteller_web/src/docs/api_doc.rs index 664b287951..cd9c4f9d4b 100644 --- a/crates/service/web/storyteller_web/src/docs/api_doc.rs +++ b/crates/service/web/storyteller_web/src/docs/api_doc.rs @@ -159,6 +159,10 @@ use artcraft_api_defs::generate::object::generate_hunyuan_2_0_image_to_3d::Gener use artcraft_api_defs::generate::object::generate_hunyuan_2_0_image_to_3d::GenerateHunyuan20ImageTo3dResponse; use artcraft_api_defs::generate::object::generate_hunyuan_2_1_image_to_3d::GenerateHunyuan21ImageTo3dRequest; use artcraft_api_defs::generate::object::generate_hunyuan_2_1_image_to_3d::GenerateHunyuan21ImageTo3dResponse; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_mini_splat::GenerateWorldlabsMarble0p1MiniSplatRequest; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_mini_splat::GenerateWorldlabsMarble0p1MiniSplatResponse; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_plus_splat::GenerateWorldlabsMarble0p1PlusSplatRequest; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_plus_splat::GenerateWorldlabsMarble0p1PlusSplatResponse; use artcraft_api_defs::generate::video::generate_kling_1_6_pro_image_to_video::GenerateKling16ProAspectRatio; use artcraft_api_defs::generate::video::generate_kling_1_6_pro_image_to_video::GenerateKling16ProDuration; use artcraft_api_defs::generate::video::generate_kling_1_6_pro_image_to_video::GenerateKling16ProImageToVideoRequest; @@ -321,6 +325,8 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux crate::http_server::endpoints::generate::image::edit::gpt_image_1_edit_image_handler::gpt_image_1_edit_image_handler, crate::http_server::endpoints::generate::object::generate_hunyuan_2_1_image_to_3d_handler::generate_hunyuan_2_1_image_to_3d_handler, crate::http_server::endpoints::generate::object::generate_hunyuan_2_0_image_to_3d_handler::generate_hunyuan_2_0_image_to_3d_handler, + crate::http_server::endpoints::generate::splat::generate_worldlabs_marble_0p1_mini_splat_handler::generate_worldlabs_marble_0p1_mini_splat_handler, + crate::http_server::endpoints::generate::splat::generate_worldlabs_marble_0p1_plus_splat_handler::generate_worldlabs_marble_0p1_plus_splat_handler, crate::http_server::endpoints::generate::video::generate_kling_1_6_pro_video_handler::generate_kling_1_6_pro_video_handler, crate::http_server::endpoints::generate::video::generate_kling_2_1_master_video_handler::generate_kling_2_1_master_video_handler, crate::http_server::endpoints::generate::video::generate_kling_2_1_pro_video_handler::generate_kling_2_1_pro_video_handler, @@ -726,6 +732,10 @@ use crate::http_server::endpoints::media_files::list::list_batch_generated_redux GenerateHunyuan21ImageTo3dResponse, GenerateHunyuan20ImageTo3dRequest, GenerateHunyuan20ImageTo3dResponse, + GenerateWorldlabsMarble0p1MiniSplatRequest, + GenerateWorldlabsMarble0p1MiniSplatResponse, + GenerateWorldlabsMarble0p1PlusSplatRequest, + GenerateWorldlabsMarble0p1PlusSplatResponse, GenerateKling16ProAspectRatio, GenerateKling16ProDuration, GenerateKling16ProImageToVideoRequest, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/mod.rs new file mode 100644 index 0000000000..e21c42741e --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/mod.rs @@ -0,0 +1 @@ +pub mod refund_wallet_after_api_failure; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/refund_wallet_after_api_failure.rs b/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/refund_wallet_after_api_failure.rs new file mode 100644 index 0000000000..9e73e00fc7 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoint_helpers/refund_wallet_after_api_failure.rs @@ -0,0 +1,53 @@ +use crate::http_server::common_responses::common_web_error::CommonWebError; +use log::{error, info}; +use mysql_queries::queries::wallets::refund::try_to_refund_ledger_entry::{try_to_refund_ledger_entry, WalletRefundOutcome}; +use sqlx::Acquire; +use tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken; + +pub async fn refund_wallet_after_api_failure( + ledger_entry_token: &WalletLedgerEntryToken, + connection: &mut sqlx::pool::PoolConnection, +) -> Result<(), CommonWebError> { + let mut transaction = connection.begin().await.map_err(|err| { + error!( + "Failed to begin refund transaction after API failure (ledger {}): {:?}", + ledger_entry_token.as_str(), err + ); + CommonWebError::ServerError + })?; + + match try_to_refund_ledger_entry(ledger_entry_token, &mut transaction).await { + Ok(WalletRefundOutcome::Refunded(summary)) => { + info!( + "Refunded {} credits after API failure (ledger {} → refund ledger {}).", + summary.refund_amount, + ledger_entry_token.as_str(), + summary.refund_ledger_entry_token.as_str(), + ); + transaction.commit().await.map_err(|err| { + error!( + "Failed to commit refund after API failure (ledger {}): {:?}", + ledger_entry_token.as_str(), err + ); + CommonWebError::ServerError + })?; + } + Ok(WalletRefundOutcome::AlreadyRefunded) => { + info!( + "Ledger entry {} was already refunded; no action needed.", + ledger_entry_token.as_str() + ); + let _ = transaction.rollback().await; + } + Err(err) => { + error!( + "Failed to refund ledger entry {} after API failure: {:?}", + ledger_entry_token.as_str(), err + ); + let _ = transaction.rollback().await; + return Err(CommonWebError::ServerError); + } + } + + Ok(()) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/mod.rs index cf693fa8c1..2f5f3a93e8 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/mod.rs @@ -2,4 +2,5 @@ pub mod common; pub mod cost_estimate; pub mod image; pub mod object; +pub mod splat; pub mod video; \ No newline at end of file diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat_handler.rs new file mode 100644 index 0000000000..df17394fa2 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_mini_splat_handler.rs @@ -0,0 +1,249 @@ +use std::sync::Arc; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure; +use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; +use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_mini_splat::{GenerateWorldlabsMarble0p1MiniSplatRequest, GenerateWorldlabsMarble0p1MiniSplatResponse}; +use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::common::visibility::Visibility; +use http_server_common::request::get_request_ip::get_request_ip; +use log::{error, info, warn}; +use mysql_queries::queries::generic_inference::worldlabs::insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token::{insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token, InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs}; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; +use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; +use sqlx::Acquire; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use world_labs_api::api::api_types::world_labs_model::WorldLabsModel; +use world_labs_api::api::requests::generate_world::generate_world::{generate_world, GenerateWorldArgs}; +use world_labs_api::api::requests::generate_world::http_request::{ContentReference, WorldPrompt}; +use world_labs_api::credentials::world_labs_api_creds::WorldLabsApiCreds; +use world_labs_api::pricing::check_pricing::{calculate_cost, InputType}; + +/// World Labs Marble 0.1-mini Splat Generation +#[utoipa::path( + post, + tag = "Generate Splats", + path = "/v1/generate/splat/worldlabs_marble_0p1_mini", + responses( + (status = 200, description = "Success", body = GenerateWorldlabsMarble0p1MiniSplatResponse), + ), + params( + ("request" = GenerateWorldlabsMarble0p1MiniSplatRequest, description = "Payload for Request"), + ) +)] +pub async fn generate_worldlabs_marble_0p1_mini_splat_handler( + http_request: HttpRequest, + request: Json, + server_state: web::Data> +) -> Result, CommonWebError> { + + payments_error_test(&request.prompt.as_deref().unwrap_or(""))?; + + if request.image_media_file_token.is_none() && request.prompt.is_none() { + return Err(CommonWebError::BadInputWithSimpleMessage( + "At least one of image_media_file_token or prompt must be provided".to_string() + )); + } + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(CommonWebError::BadInputWithSimpleMessage(reason)); + } + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await?; + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CommonWebError::ServerError + })?; + + let maybe_avt_token = server_state + .avt_cookie_manager + .get_avt_token_from_request(&http_request); + + let user_token = match maybe_user_session.as_ref() { + Some(session) => &session.user_token, + None => { + return Err(CommonWebError::NotAuthorized); + } + }; + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string()) + })?; + + let apriori_job_token = InferenceJobToken::generate(); + + // Determine input type and calculate cost + let input_type = if request.image_media_file_token.is_some() { + InputType::ImageNonPanorama + } else { + InputType::Text + }; + + let cost = calculate_cost(WorldLabsModel::Marble0p1Mini, input_type); + let cost_in_cents = cost.us_dollar_cents as u64; + + info!("Charging wallet: {} cents for Marble 0.1-mini splat", cost_in_cents); + + let wallet_deduction = attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost_in_cents, + &mut mysql_connection, + ).await?; + + // Build WorldPrompt based on input + let world_prompt = if let Some(image_token) = &request.image_media_file_token { + let image_urls = lookup_image_urls_as_map( + &http_request, + &mut mysql_connection, + server_state.server_environment, + &[image_token.clone()], + ).await?; + + let cdn_url = image_urls.get(image_token) + .ok_or_else(|| { + warn!("Image token not found in lookup results: {:?}", image_token); + CommonWebError::ServerError + })?; + + WorldPrompt::Image { + image_prompt: ContentReference::Uri { uri: cdn_url.clone() }, + text_prompt: request.prompt.clone(), + is_pano: None, + disable_recaption: None, + } + } else { + WorldPrompt::Text { + text_prompt: request.prompt.clone(), + disable_recaption: None, + } + }; + + // Call World Labs API to start generation + let creds = WorldLabsApiCreds::new(server_state.worldlabs.api_key.clone()); + + let generate_result = match generate_world(GenerateWorldArgs { + creds: &creds, + world_prompt, + display_name: None, + model: WorldLabsModel::Marble0p1Mini, + seed: None, + tags: None, + permission: None, + request_timeout: None, + }).await { + Ok(result) => result, + Err(err) => { + warn!("World Labs generate_world error: {:?}", err); + refund_wallet_after_api_failure(&wallet_deduction.ledger_entry_token, &mut mysql_connection).await?; + return Err(CommonWebError::ServerError); + } + }; + + let ip_address = get_request_ip(&http_request); + + let mut transaction = mysql_connection + .begin() + .await + .map_err(|err| { + error!("Error starting MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + let prompt_result = insert_prompt(InsertPromptArgs { + maybe_apriori_prompt_token: None, + prompt_type: PromptType::ArtcraftApp, + maybe_creator_user_token: Some(user_token), + maybe_model_type: Some(ModelType::Marble0p1Mini), + maybe_generation_provider: Some(GenerationProvider::Artcraft), + maybe_positive_prompt: request.prompt.as_deref(), + maybe_negative_prompt: None, + maybe_other_args: None, + creator_ip_address: &ip_address, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let prompt_token = match prompt_result { + Ok(token) => Some(token), + Err(err) => { + warn!("Error inserting prompt: {:?}", err); + None + } + }; + + if let Some(image_token) = &request.image_media_file_token { + if let Some(token) = prompt_token.as_ref() { + let result = insert_batch_prompt_context_items(InsertBatchArgs { + prompt_token: token.clone(), + items: vec![ + PromptContextItem { + media_token: image_token.clone(), + context_semantic_type: PromptContextSemanticType::Imgref, + } + ], + transaction: &mut transaction, + }).await; + + if let Err(err) = result { + warn!("Error inserting batch prompt context items: {:?}", err); + } + } + } + + let db_result = insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token(InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs { + apriori_job_token: &apriori_job_token, + uuid_idempotency_token: &request.uuid_idempotency_token, + maybe_external_third_party_id: generate_result.operation_id.as_str(), + maybe_inference_args: None, + maybe_prompt_token: prompt_token.as_ref(), + maybe_wallet_ledger_entry_token: Some(&wallet_deduction.ledger_entry_token), + maybe_creator_user_token: Some(user_token), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: Visibility::Public, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let job_token = match db_result { + Ok(token) => token, + Err(err) => { + warn!("Error inserting generic inference job for WorldLabs queue: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + let _r = transaction + .commit() + .await + .map_err(|err| { + error!("Error committing MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + Ok(Json(GenerateWorldlabsMarble0p1MiniSplatResponse { + success: true, + inference_job_token: job_token, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat_handler.rs new file mode 100644 index 0000000000..89b1210959 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/generate_worldlabs_marble_0p1_plus_splat_handler.rs @@ -0,0 +1,249 @@ +use std::sync::Arc; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure; +use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; +use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; +use crate::state::server_state::ServerState; +use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; +use actix_web::web::Json; +use actix_web::{web, HttpRequest}; +use artcraft_api_defs::generate::splat::generate_worldlabs_marble_0p1_plus_splat::{GenerateWorldlabsMarble0p1PlusSplatRequest, GenerateWorldlabsMarble0p1PlusSplatResponse}; +use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::common::visibility::Visibility; +use http_server_common::request::get_request_ip::get_request_ip; +use log::{error, info, warn}; +use mysql_queries::queries::generic_inference::worldlabs::insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token::{insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token, InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs}; +use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token; +use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem}; +use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; +use sqlx::Acquire; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use world_labs_api::api::api_types::world_labs_model::WorldLabsModel; +use world_labs_api::api::requests::generate_world::generate_world::{generate_world, GenerateWorldArgs}; +use world_labs_api::api::requests::generate_world::http_request::{ContentReference, WorldPrompt}; +use world_labs_api::credentials::world_labs_api_creds::WorldLabsApiCreds; +use world_labs_api::pricing::check_pricing::{calculate_cost, InputType}; + +/// World Labs Marble 0.1-plus Splat Generation +#[utoipa::path( + post, + tag = "Generate Splats", + path = "/v1/generate/splat/worldlabs_marble_0p1_plus", + responses( + (status = 200, description = "Success", body = GenerateWorldlabsMarble0p1PlusSplatResponse), + ), + params( + ("request" = GenerateWorldlabsMarble0p1PlusSplatRequest, description = "Payload for Request"), + ) +)] +pub async fn generate_worldlabs_marble_0p1_plus_splat_handler( + http_request: HttpRequest, + request: Json, + server_state: web::Data> +) -> Result, CommonWebError> { + + payments_error_test(&request.prompt.as_deref().unwrap_or(""))?; + + if request.image_media_file_token.is_none() && request.prompt.is_none() { + return Err(CommonWebError::BadInputWithSimpleMessage( + "At least one of image_media_file_token or prompt must be provided".to_string() + )); + } + + if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) { + return Err(CommonWebError::BadInputWithSimpleMessage(reason)); + } + + let mut mysql_connection = server_state.mysql_pool + .acquire() + .await?; + + let maybe_user_session = server_state + .session_checker + .maybe_get_user_session_from_connection(&http_request, &mut mysql_connection) + .await + .map_err(|e| { + warn!("Session checker error: {:?}", e); + CommonWebError::ServerError + })?; + + let maybe_avt_token = server_state + .avt_cookie_manager + .get_avt_token_from_request(&http_request); + + let user_token = match maybe_user_session.as_ref() { + Some(session) => &session.user_token, + None => { + return Err(CommonWebError::NotAuthorized); + } + }; + + insert_idempotency_token(&request.uuid_idempotency_token, &mut *mysql_connection) + .await + .map_err(|err| { + error!("Error inserting idempotency token: {:?}", err); + CommonWebError::BadInputWithSimpleMessage("repeated idempotency token".to_string()) + })?; + + let apriori_job_token = InferenceJobToken::generate(); + + // Determine input type and calculate cost + let input_type = if request.image_media_file_token.is_some() { + InputType::ImageNonPanorama + } else { + InputType::Text + }; + + let cost = calculate_cost(WorldLabsModel::Marble0p1Plus, input_type); + let cost_in_cents = cost.us_dollar_cents as u64; + + info!("Charging wallet: {} cents for Marble 0.1-plus splat", cost_in_cents); + + let wallet_deduction = attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost_in_cents, + &mut mysql_connection, + ).await?; + + // Build WorldPrompt based on input + let world_prompt = if let Some(image_token) = &request.image_media_file_token { + let image_urls = lookup_image_urls_as_map( + &http_request, + &mut mysql_connection, + server_state.server_environment, + &[image_token.clone()], + ).await?; + + let cdn_url = image_urls.get(image_token) + .ok_or_else(|| { + warn!("Image token not found in lookup results: {:?}", image_token); + CommonWebError::ServerError + })?; + + WorldPrompt::Image { + image_prompt: ContentReference::Uri { uri: cdn_url.clone() }, + text_prompt: request.prompt.clone(), + is_pano: None, + disable_recaption: None, + } + } else { + WorldPrompt::Text { + text_prompt: request.prompt.clone(), + disable_recaption: None, + } + }; + + // Call World Labs API to start generation + let creds = WorldLabsApiCreds::new(server_state.worldlabs.api_key.clone()); + + let generate_result = match generate_world(GenerateWorldArgs { + creds: &creds, + world_prompt, + display_name: None, + model: WorldLabsModel::Marble0p1Plus, + seed: None, + tags: None, + permission: None, + request_timeout: None, + }).await { + Ok(result) => result, + Err(err) => { + warn!("World Labs generate_world error: {:?}", err); + refund_wallet_after_api_failure(&wallet_deduction.ledger_entry_token, &mut mysql_connection).await?; + return Err(CommonWebError::ServerError); + } + }; + + let ip_address = get_request_ip(&http_request); + + let mut transaction = mysql_connection + .begin() + .await + .map_err(|err| { + error!("Error starting MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + let prompt_result = insert_prompt(InsertPromptArgs { + maybe_apriori_prompt_token: None, + prompt_type: PromptType::ArtcraftApp, + maybe_creator_user_token: Some(user_token), + maybe_model_type: Some(ModelType::Marble0p1Plus), + maybe_generation_provider: Some(GenerationProvider::Artcraft), + maybe_positive_prompt: request.prompt.as_deref(), + maybe_negative_prompt: None, + maybe_other_args: None, + creator_ip_address: &ip_address, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let prompt_token = match prompt_result { + Ok(token) => Some(token), + Err(err) => { + warn!("Error inserting prompt: {:?}", err); + None + } + }; + + if let Some(image_token) = &request.image_media_file_token { + if let Some(token) = prompt_token.as_ref() { + let result = insert_batch_prompt_context_items(InsertBatchArgs { + prompt_token: token.clone(), + items: vec![ + PromptContextItem { + media_token: image_token.clone(), + context_semantic_type: PromptContextSemanticType::Imgref, + } + ], + transaction: &mut transaction, + }).await; + + if let Err(err) = result { + warn!("Error inserting batch prompt context items: {:?}", err); + } + } + } + + let db_result = insert_generic_inference_job_for_worldlabs_queue_with_apriori_job_token(InsertGenericInferenceForWorldlabsWithAprioriJobTokenArgs { + apriori_job_token: &apriori_job_token, + uuid_idempotency_token: &request.uuid_idempotency_token, + maybe_external_third_party_id: generate_result.operation_id.as_str(), + maybe_inference_args: None, + maybe_prompt_token: prompt_token.as_ref(), + maybe_wallet_ledger_entry_token: Some(&wallet_deduction.ledger_entry_token), + maybe_creator_user_token: Some(user_token), + maybe_avt_token: maybe_avt_token.as_ref(), + creator_ip_address: &ip_address, + creator_set_visibility: Visibility::Public, + mysql_executor: &mut *transaction, + phantom: Default::default(), + }).await; + + let job_token = match db_result { + Ok(token) => token, + Err(err) => { + warn!("Error inserting generic inference job for WorldLabs queue: {:?}", err); + return Err(CommonWebError::ServerError); + } + }; + + let _r = transaction + .commit() + .await + .map_err(|err| { + error!("Error committing MySQL transaction: {:?}", err); + CommonWebError::ServerError + })?; + + Ok(Json(GenerateWorldlabsMarble0p1PlusSplatResponse { + success: true, + inference_job_token: job_token, + })) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/mod.rs new file mode 100644 index 0000000000..75f8229fb4 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/splat/mod.rs @@ -0,0 +1,2 @@ +pub mod generate_worldlabs_marble_0p1_mini_splat_handler; +pub mod generate_worldlabs_marble_0p1_plus_splat_handler; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/multi_function/seedance_2p0_multi_function_video_gen_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/multi_function/seedance_2p0_multi_function_video_gen_handler.rs index d8752c358a..cffaae017f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/multi_function/seedance_2p0_multi_function_video_gen_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/video/multi_function/seedance_2p0_multi_function_video_gen_handler.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; use crate::http_server::common_responses::common_web_error::CommonWebError; +use crate::http_server::endpoint_helpers::refund_wallet_after_api_failure::refund_wallet_after_api_failure; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::state::server_state::ServerState; @@ -29,7 +30,6 @@ use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_it insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem, }; use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; -use mysql_queries::queries::wallets::refund::try_to_refund_ledger_entry::{try_to_refund_ledger_entry, WalletRefundOutcome}; use seedance2pro::creds::seedance2pro_session::Seedance2ProSession; use seedance2pro::requests::generate_video::generate_video::{ generate_video, BatchCount, GenerateVideoArgs, Resolution, @@ -221,7 +221,7 @@ pub async fn seedance_2p0_multi_function_video_gen_handler( Ok(resp) => resp, Err(err) => { warn!("Error calling seedance2pro generate_video: {:?}", err); - refund_after_api_failure(&deduction_result.ledger_entry_token, &mut mysql_connection).await; + refund_wallet_after_api_failure(&deduction_result.ledger_entry_token, &mut mysql_connection).await?; // TODO: We need to start returning robust failure messages to the caller. return Err(CommonWebError::ServerError); @@ -375,53 +375,6 @@ pub async fn seedance_2p0_multi_function_video_gen_handler( })) } -async fn refund_after_api_failure( - ledger_entry_token: &tokens::tokens::wallet_ledger_entries::WalletLedgerEntryToken, - connection: &mut sqlx::pool::PoolConnection, -) { - let mut transaction = match connection.begin().await { - Ok(tx) => tx, - Err(err) => { - error!( - "Failed to begin refund transaction after API failure (ledger {}): {:?}", - ledger_entry_token.as_str(), err - ); - return; - } - }; - - match try_to_refund_ledger_entry(ledger_entry_token, &mut transaction).await { - Ok(WalletRefundOutcome::Refunded(summary)) => { - info!( - "Refunded {} credits after API failure (ledger {} → refund ledger {}).", - summary.refund_amount, - ledger_entry_token.as_str(), - summary.refund_ledger_entry_token.as_str(), - ); - if let Err(err) = transaction.commit().await { - error!( - "Failed to commit refund after API failure (ledger {}): {:?}", - ledger_entry_token.as_str(), err - ); - } - } - Ok(WalletRefundOutcome::AlreadyRefunded) => { - info!( - "Ledger entry {} was already refunded; no action needed.", - ledger_entry_token.as_str() - ); - let _ = transaction.rollback().await; - } - Err(err) => { - error!( - "Failed to refund ledger entry {} after API failure: {:?}", - ledger_entry_token.as_str(), err - ); - let _ = transaction.rollback().await; - } - } -} - async fn upload_to_seedance2pro( session: &Seedance2ProSession, our_cdn_url: &str, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs index 35179e2219..2365e439d5 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/inference_job/utils/estimates/estimate_job_progress.rs @@ -85,6 +85,7 @@ pub fn estimate_job_progress(job: &GenericInferenceJobStatus, maybe_args: Option InferenceCategory::Mocap => 0, InferenceCategory::VideoFilter => 0, InferenceCategory::ConvertBvhToWorkflow => 0, + InferenceCategory::SplatGeneration => percent(duration_seconds, 60*5), InferenceCategory::DeprecatedField => 0, // TODO(bt,2024-07-16): Read job type instead. }; diff --git a/crates/service/web/storyteller_web/src/http_server/mod.rs b/crates/service/web/storyteller_web/src/http_server/mod.rs index 672a4343c9..b3bef446c3 100644 --- a/crates/service/web/storyteller_web/src/http_server/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/mod.rs @@ -1,5 +1,6 @@ pub mod common_requests; pub mod common_responses; +pub mod endpoint_helpers; pub mod cookies; pub mod deprecated_endpoints; pub mod endpoints; diff --git a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs index a0522ddf04..012ef22d26 100644 --- a/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs +++ b/crates/service/web/storyteller_web/src/http_server/routes/application_routes/generate_routes.rs @@ -21,6 +21,8 @@ use crate::http_server::endpoints::generate::image::multi_function::nano_banana_ use crate::http_server::endpoints::generate::image::multi_function::nano_banana_pro_multi_function_image_gen_handler::nano_banana_pro_multi_function_image_gen_handler; use crate::http_server::endpoints::generate::image::remove_image_background_handler::remove_image_background_handler; use crate::http_server::endpoints::generate::object::generate_hunyuan_2_0_image_to_3d_handler::generate_hunyuan_2_0_image_to_3d_handler; +use crate::http_server::endpoints::generate::splat::generate_worldlabs_marble_0p1_mini_splat_handler::generate_worldlabs_marble_0p1_mini_splat_handler; +use crate::http_server::endpoints::generate::splat::generate_worldlabs_marble_0p1_plus_splat_handler::generate_worldlabs_marble_0p1_plus_splat_handler; use crate::http_server::endpoints::generate::object::generate_hunyuan_2_1_image_to_3d_handler::generate_hunyuan_2_1_image_to_3d_handler; use crate::http_server::endpoints::generate::object::multi_function::generate_hunyuan3d_v3_multi_function_object_handler::generate_hunyuan3d_v3_multi_function_object_handler; use crate::http_server::endpoints::generate::video::generate_kling_1_6_pro_video_handler::generate_kling_1_6_pro_video_handler; @@ -248,5 +250,15 @@ where .route(web::head().to(|| HttpResponse::Ok())) ) ) + .service(web::scope("/splat") + .service(web::resource("/worldlabs_marble_0p1_mini") + .route(web::post().to(generate_worldlabs_marble_0p1_mini_splat_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + .service(web::resource("/worldlabs_marble_0p1_plus") + .route(web::post().to(generate_worldlabs_marble_0p1_plus_splat_handler)) + .route(web::head().to(|| HttpResponse::Ok())) + ) + ) ) } diff --git a/crates/service/web/storyteller_web/src/main.rs b/crates/service/web/storyteller_web/src/main.rs index 5b6df8e2cd..ddbcd39526 100644 --- a/crates/service/web/storyteller_web/src/main.rs +++ b/crates/service/web/storyteller_web/src/main.rs @@ -80,7 +80,7 @@ use crate::http_server::web_utils::handle_multipart_error::handle_multipart_erro use crate::http_server::web_utils::scoped_temp_dir_creator::ScopedTempDirCreator; use crate::state::certs::google_sign_in_cert::GoogleSignInCert; use crate::state::memory_cache::model_token_to_info_cache::ModelTokenToInfoCache; -use crate::state::server_state::{DurableInMemoryCaches, EnvConfig, EphemeralInMemoryCaches, FalData, InMemoryCaches, OpenAiData, ResendData, Seedance2ProData, ServerInfo, ServerState, StaticFeatureFlags, StripeSettings, TrollBans}; +use crate::state::server_state::{DurableInMemoryCaches, EnvConfig, EphemeralInMemoryCaches, FalData, InMemoryCaches, OpenAiData, ResendData, Seedance2ProData, ServerInfo, ServerState, StaticFeatureFlags, StripeSettings, TrollBans, WorldLabsData}; use crate::threads::db_health_checker_thread::db_health_check_status::HealthCheckStatus; use crate::threads::db_health_checker_thread::db_health_checker_thread::db_health_checker_thread; use crate::threads::poll_ip_banlist_thread::poll_ip_bans; @@ -413,6 +413,8 @@ async fn main() -> AnyhowResult<()> { let resend_api_key = easyenv::get_env_string_required("RESEND_API_KEY")?; + let worldlabs_api_key = easyenv::get_env_string_required("WORLDLABS_API_KEY")?; + let startup_time = Utc::now(); let server_state = ServerState { @@ -477,6 +479,9 @@ async fn main() -> AnyhowResult<()> { resend: ResendData { api_key: resend_api_key, }, + worldlabs: WorldLabsData { + api_key: worldlabs_api_key, + }, caches: InMemoryCaches { durable: DurableInMemoryCaches { model_token_info: model_token_info_cache, diff --git a/crates/service/web/storyteller_web/src/state/server_state.rs b/crates/service/web/storyteller_web/src/state/server_state.rs index 291daf35c1..7fc4a0a35f 100644 --- a/crates/service/web/storyteller_web/src/state/server_state.rs +++ b/crates/service/web/storyteller_web/src/state/server_state.rs @@ -98,7 +98,9 @@ pub struct ServerState { pub openai: OpenAiData, pub resend: ResendData, - + + pub worldlabs: WorldLabsData, + /// Where to store audio uploads for w2l pub audio_uploads_bucket_root: String, @@ -331,3 +333,9 @@ pub struct OpenAiData { pub struct ResendData { pub api_key: String, } + +/// World Labs integration +#[derive(Clone)] +pub struct WorldLabsData { + pub api_key: String, +}