diff --git a/Cargo.lock b/Cargo.lock index c877641655..74dee62c16 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -698,6 +698,7 @@ dependencies = [ "url_utils", "web_base64", "workspace-build-acceleration", + "world_labs_client", ] [[package]] @@ -14192,6 +14193,7 @@ dependencies = [ "futures 0.3.31", "http_headers", "indexmap 2.9.0", + "jwt_light", "log", "log4rs", "once_cell", diff --git a/Cargo.toml b/Cargo.toml index 65bb724fc2..82799245e9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -220,6 +220,7 @@ testing = { path = "crates/lib/testing" } tokens = { path = "crates/schema/public/tokens" } url_utils = { path = "crates/lib/url_utils" } web_base64 = { path = "crates/lib/data/web_base64" } +world_labs_client = { path = "crates/api_clients/world_labs_client" } # External libs # Centrally pinned versions so that we can uniformly upgrade/downgrade versions across the workspace diff --git a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs new file mode 100644 index 0000000000..f66c4d3aa3 --- /dev/null +++ b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs @@ -0,0 +1,143 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::api_error::ApiError; +use crate::error::client_error::ClientError; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::constants::{APPLICATION_JSON, USER_AGENT}; +use crate::utils::filter_bad_response::filter_bad_response; +use crate::utils::http_get_anonymous::http_get_anonymous; +use chrono::{DateTime, Utc}; +use enums::by_table::media_files::media_file_animation_type::MediaFileAnimationType; +use enums::by_table::media_files::media_file_class::MediaFileClass; +use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::common::generation_provider::GenerationProvider; +use enums::common::visibility::Visibility; +use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; +use enums_public::by_table::model_weights::public_weights_types::PublicWeightsType; +use idempotency::uuid::generate_random_uuid; +use log::debug; +use reqwest::multipart::Form; +use reqwest::Client; +use serde_derive::{Deserialize, Serialize}; +use std::fs::File; +use std::io::Read; +use std::path::Path; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::prompts::PromptToken; +use uuid::uuid; + +pub struct LegacyUploadMediaFileFromFileArgs<'a, P: AsRef> { + pub api_host: &'a ApiHost, + pub maybe_creds: Option<&'a StorytellerCredentialSet>, + + // NB: Path needs to be owned for the request. + pub path: P, + + // /// If true, we should hide the image from the user's gallery. + // pub is_intermediate_system_file: bool, + + // /// If provided, this is the prompt that this image is associated with. + // /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set. + // pub maybe_prompt_token: Option<&'a PromptToken>, + + // /// If provided, this is the service provider that created the image. + // /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set. + // pub maybe_generation_provider: Option, + + // /// If provided, this groups the file into a batch + // /// TODO: This shouldn't be set clientside without the backend generating the token + // /// and cryptographically securing it. But we need to go fast here. + // pub maybe_batch_token: Option<&'a BatchGenerationToken>, +} + + +/// Upload using the legacy endpoint (spz, etc.) +pub async fn legacy_upload_media_file_from_file>( + args: LegacyUploadMediaFileFromFileArgs<'_, P>, +) -> Result { + + validate_args(&args)?; + + let url = get_route(args.api_host); + + debug!("Requesting {:?}", &url); + + //let mut file = File::open(path)?; + + let client = Client::builder() + .gzip(true) + .build() + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + let mut form = Form::new() + .text("uuid_idempotency_token", generate_random_uuid()) + .file("file", args.path) + .await + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + //if args.is_intermediate_system_file { + // form = form.text("is_intermediate_system_file", "true"); + //} + + //if let Some(prompt_token) = &args.maybe_prompt_token { + // form = form.text("maybe_prompt_token", prompt_token.to_string()); + //} + + //if let Some(batch_token) = &args.maybe_batch_token { + // form = form.text("maybe_batch_token", batch_token.to_string()); + //} + + let mut request_builder = client.post(url) + .header("User-Agent", USER_AGENT) + .header("Accept", APPLICATION_JSON); + + if let Some(creds) = args.maybe_creds { + if let Some(header) = &creds.maybe_as_cookie_header() { + request_builder = request_builder.header("Cookie", header); + } + } + + let response = request_builder + .multipart(form) + .send() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response = filter_bad_response(response).await?; + let response_body = &response.text().await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let media_file = serde_json::from_str(&response_body) + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + Ok(media_file) +} + +fn get_route(api_host: &ApiHost) -> String { + let api_hostname_and_scheme = api_host.to_api_hostname_and_scheme(); + format!("{}/v1/media_files/upload", api_hostname_and_scheme) +} + +// TODO(bt,2025-04-22): Share API definitions between client and server in common crate. + +#[derive(Deserialize, Debug)] +pub struct UploadImageMediaFileSuccessResponse { + pub success: bool, + pub media_file_token: MediaFileToken, +} + +fn validate_args>( + args: &LegacyUploadMediaFileFromFileArgs<'_, P>, +) -> Result<(), StorytellerError> { + // if args.is_intermediate_system_file && args.maybe_prompt_token.is_some() { + // return Err(StorytellerError::Client(ClientError::InvalidPreflightRequest( + // "Cannot set `is_intermediate_system_file` to true if `maybe_prompt_token` is provided." + // .to_string()))); + // } + + Ok(()) +} diff --git a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs index d7cf4c1378..99cc13aa5f 100644 --- a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs +++ b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs @@ -5,3 +5,4 @@ pub mod upload_image_media_file_from_bytes; pub mod upload_image_media_file_from_file; pub mod upload_new_engine_asset_from_file; pub mod upload_video_media_file_from_file; +pub mod legacy_upload_media_file_from_file; diff --git a/crates/api_clients/world_labs_client/Cargo.toml b/crates/api_clients/world_labs_client/Cargo.toml index 0470d89d01..4c3fa182a8 100644 --- a/crates/api_clients/world_labs_client/Cargo.toml +++ b/crates/api_clients/world_labs_client/Cargo.toml @@ -20,7 +20,9 @@ cloudflare_errors.workspace = true cloudflare_mitigation.workspace = true cookie_store.workspace = true errors.workspace = true +filesys.workspace = true http_headers.workspace = true +jwt_light.workspace = true # External anyhow.workspace = true diff --git a/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs b/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs index 9ab01960d9..d5c0341995 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/objects/begin_object_image_upload.rs @@ -62,7 +62,7 @@ pub async fn begin_object_image_upload(args: BeginObjectImageUploadArgs<'_>) -> .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") .header(REFERER, REFERER_VALUE) .header(CONTENT_TYPE, CONTENT_TYPE_APPLICATION_JSON) - .header(AUTHORIZATION, args.bearer_token.to_bearer_token_string()) + .header(AUTHORIZATION, args.bearer_token.to_bearer_token_header_string()) .header(ORIGIN, ORIGIN_VALUE) .header(SEC_GPC, "1") .header(CONNECTION, CONNECTION_KEEP_ALIVE) diff --git a/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs b/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs index 05a4da77e9..11d9b9f3dc 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/objects/create_run_object.rs @@ -54,7 +54,7 @@ pub async fn create_run_object(args: CreateRunObjectArgs<'_>) -> Result) -> Result .header(ACCEPT_ENCODING, "gzip, deflate, br, zstd") .header(REFERER, REFERER_VALUE) .header(CONTENT_TYPE, CONTENT_TYPE_APPLICATION_JSON) - .header(AUTHORIZATION, args.bearer_token.to_bearer_token_string()) + .header(AUTHORIZATION, args.bearer_token.to_bearer_token_header_string()) .header(ORIGIN, ORIGIN_VALUE) .header(SEC_GPC, "1") .header(CONNECTION, CONNECTION_KEEP_ALIVE) diff --git a/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs b/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs index 64aed8760f..dc9d2db97c 100644 --- a/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs +++ b/crates/api_clients/world_labs_client/src/api/requests/recaption/recaption_image.rs @@ -58,7 +58,7 @@ pub async fn recaption_image(args: RecaptionImageArgs<'_>) -> Result) -> Result) -> Result, + + /// From the "exp" field + pub expiration: DateTime, + + /// From the "user_id" field + pub user_id: Option, + + /// From the "email" field + pub email: Option, +} + +impl ParseJwtClaims for JwtClaims { + fn extract_claims(common_claims: CommonClaims, extra_claims: Map) -> Result { + let user_id = extra_claims.get("user_id") + .map(|val| val.as_str()) + .flatten() + .map(|val| val.to_string()); + + let mut email = extra_claims.get("email") + .map(|val| val.as_str()) + .flatten() + .map(|val| val.to_string()); + + if email.is_none() { + // Parse out email from this field: + // "firebase": { + // "identities": { + // "email": [ + // "sabalid234@crsay.com" + // ] + // }, + // }, + email = extra_claims.get("firebase") + .and_then(|val| val.as_object()) + .and_then(|fb| fb.get("identities")) + .and_then(|val| val.as_object()) + .and_then(|id| id.get("email")) + .and_then(|emails| emails.as_array()) + .and_then(|emails| emails.get(0)) + .and_then(|email| email.as_str()) + .map(|email| email.to_string()); + } + + Ok(Self { + created: common_claims.created, + expiration: common_claims.expiration, + user_id, + email, + }) + } +} + +#[cfg(test)] +mod tests { + use crate::credentials::jwt_claims::JwtClaims; + use jwt_light::parse_jwt_claims_trait::ParseJwtClaims; + + #[test] + fn test_token() { + /* + { + "iss": "https://securetoken.google.com/wlt-training-gsc", + "aud": "wlt-training-gsc", + "auth_time": 1766002923, + "user_id": "CArJJ2PQFFd8l2Rf8uptZhl86d33", + "sub": "CArJJ2PQFFd8l2Rf8uptZhl86d33", + "iat": 1766099113, + "exp": 1766102713, + "email": "sabalid234@crsay.com", + "email_verified": true, + "firebase": { + "identities": { + "email": [ + "sabalid234@crsay.com" + ] + }, + "sign_in_provider": "password" + } + } + */ + let jwt = "eyJhbGciOiJSUzI1NiIsImtpZCI6IjM4MTFiMDdmMjhiODQxZjRiNDllNDgyNTg1ZmQ2NmQ1NWUzOGRiNWQiLCJ0eXAiOiJKV1QifQ.eyJpc3MiOiJodHRwczovL3NlY3VyZXRva2VuLmdvb2dsZS5jb20vd2x0LXRyYWluaW5nLWdzYyIsImF1ZCI6IndsdC10cmFpbmluZy1nc2MiLCJhdXRoX3RpbWUiOjE3NjYwMDI5MjMsInVzZXJfaWQiOiJDQXJKSjJQUUZGZDhsMlJmOHVwdFpobDg2ZDMzIiwic3ViIjoiQ0FySkoyUFFGRmQ4bDJSZjh1cHRaaGw4NmQzMyIsImlhdCI6MTc2NjA5OTExMywiZXhwIjoxNzY2MTAyNzEzLCJlbWFpbCI6InNhYmFsaWQyMzRAY3JzYXkuY29tIiwiZW1haWxfdmVyaWZpZWQiOnRydWUsImZpcmViYXNlIjp7ImlkZW50aXRpZXMiOnsiZW1haWwiOlsic2FiYWxpZDIzNEBjcnNheS5jb20iXX0sInNpZ25faW5fcHJvdmlkZXIiOiJwYXNzd29yZCJ9fQ.byDW2QiRfajxDgQMD89YNdlRxO5rsHves3nwK_Ry5ZLFThARbdU9-wTgbCz1mzz9P3amUkM_XM9E4rEAnj2u-kCv5iW_FIfRDeJ0Z8FlwqxCt6ekvJgwYE80BUIedEvTiNpP_AilThz1ILqAu3gVQ2GXUcolquvS9i12iGTl-thJH6gnZgv3UdP2Vrr5ZP9VhKgMi5gQER5umFQ4Ud64kD9P9SLttRI2RJdolavJengQBAfec8hOEk6uikw20rEBeT_xzl8TpLNG6i4SIGYqJQHfRALGNniYoU9RlIoFR9VGeyk_6BFFGRd8I_EgL-fxnfmpV8w-riXHlOUg6uNRgw"; + let claims = JwtClaims::parse_claims(jwt).unwrap(); + + assert_eq!(claims.user_id, Some("CArJJ2PQFFd8l2Rf8uptZhl86d33".to_string())); + assert_eq!(claims.email, Some("sabalid234@crsay.com".to_string())); + assert_eq!(claims.created.timestamp(), 1766099113); + assert_eq!(claims.expiration.timestamp(), 1766102713); + } +} diff --git a/crates/api_clients/world_labs_client/src/credentials/mod.rs b/crates/api_clients/world_labs_client/src/credentials/mod.rs index fbab09d193..f752f37ea7 100644 --- a/crates/api_clients/world_labs_client/src/credentials/mod.rs +++ b/crates/api_clients/world_labs_client/src/credentials/mod.rs @@ -1,2 +1,4 @@ pub mod world_labs_bearer_token; pub mod world_labs_cookies; +pub mod worldlabs_refresh_token; +pub mod jwt_claims; diff --git a/crates/api_clients/world_labs_client/src/credentials/world_labs_bearer_token.rs b/crates/api_clients/world_labs_client/src/credentials/world_labs_bearer_token.rs index 516991c43a..779b3968f0 100644 --- a/crates/api_clients/world_labs_client/src/credentials/world_labs_bearer_token.rs +++ b/crates/api_clients/world_labs_client/src/credentials/world_labs_bearer_token.rs @@ -1,3 +1,29 @@ +use crate::credentials::jwt_claims::JwtClaims; +use crate::error::world_labs_client_error::WorldLabsClientError; +use crate::error::world_labs_error::WorldLabsError; +use jwt_light::parse_jwt_claims_trait::ParseJwtClaims; +use log::error; + +/* + +This is issued by Google, even for email+password login: + + https://identitytoolkit.googleapis.com/v1/accounts:signInWithPassword?key={ENTROPY} + +Value: + +{ + "kind": "identitytoolkit#VerifyPasswordResponse", + "localId": "{{WorldLabsUserId}}", + "email": "{{email@whatever.com}}", + "displayName": "", + "idToken": "{{BEARER_TOKEN}}", + "registered": true, + "refreshToken": "{{REFRESH_TOKEN}}", + "expiresIn": "3600" +} + +*/ #[derive(Clone)] pub struct WorldLabsBearerToken { @@ -29,7 +55,15 @@ impl WorldLabsBearerToken { self.bearer_token.clone() } - pub fn to_bearer_token_string(&self) -> String { + pub fn to_bearer_token_header_string(&self) -> String { format!("Bearer {}", self.bearer_token) } + + pub fn parse_jwt_claims(&self) -> Result { + JwtClaims::parse_claims(&self.bearer_token) + .map_err(|err| { + error!("Failed to parse bearer token into JWT claims: {}", err); + WorldLabsError::Client(WorldLabsClientError::FailedToParseJwtClaims(err)) + }) + } } diff --git a/crates/api_clients/world_labs_client/src/credentials/worldlabs_refresh_token.rs b/crates/api_clients/world_labs_client/src/credentials/worldlabs_refresh_token.rs new file mode 100644 index 0000000000..5e305ec425 --- /dev/null +++ b/crates/api_clients/world_labs_client/src/credentials/worldlabs_refresh_token.rs @@ -0,0 +1,25 @@ + +#[derive(Clone)] +pub struct WorldLabsRefreshToken { + /// The refresh token + /// Used to get a new bearer token + refresh_token: String, +} + +impl WorldLabsRefreshToken { + pub fn new(refresh_token: String) -> Self { + Self { refresh_token } + } + + pub fn as_raw_str(&self) -> &str { + &self.refresh_token + } + + pub fn as_raw_bytes(&self) -> &[u8] { + self.refresh_token.as_bytes() + } + + pub fn to_raw_string(&self) -> String { + self.refresh_token.clone() + } +} diff --git a/crates/api_clients/world_labs_client/src/error/world_labs_client_error.rs b/crates/api_clients/world_labs_client/src/error/world_labs_client_error.rs index 8261b28ec0..abcfb7f7af 100644 --- a/crates/api_clients/world_labs_client/src/error/world_labs_client_error.rs +++ b/crates/api_clients/world_labs_client/src/error/world_labs_client_error.rs @@ -7,6 +7,9 @@ pub enum WorldLabsClientError { /// An error was encountered in building the Wreq client WreqClientError(wreq::Error), + + /// Failed to parse JWT claims. + FailedToParseJwtClaims(jwt_light::error::JwtError), /// Error serializing a message to send the websocket WebsocketRequestSerializationError(serde_json::Error), @@ -114,6 +117,7 @@ impl Display for WorldLabsClientError { fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result { match self { Self::WreqClientError(err) => write!(f, "Wreq client error (during client creation): {}", err), + Self::FailedToParseJwtClaims(err) => write!(f, "Failed to parse JWT claims: {}", err), Self::WebsocketRequestSerializationError(err) => write!(f, "Websocket request serialization error: {}", err), Self::WebsocketLockError => write!(f, "Websocket lock error"), Self::WebsocketReadError(err) => write!(f, "Websocket read error: {}", err), diff --git a/crates/api_clients/world_labs_client/src/recipes/mod.rs b/crates/api_clients/world_labs_client/src/recipes/mod.rs index 84affa5b06..d17ac4bb12 100644 --- a/crates/api_clients/world_labs_client/src/recipes/mod.rs +++ b/crates/api_clients/world_labs_client/src/recipes/mod.rs @@ -1 +1 @@ -mod upload_image_and_create_world_with_retry; \ No newline at end of file +pub mod upload_image_and_create_world_with_retry; \ No newline at end of file diff --git a/crates/api_clients/world_labs_client/src/recipes/upload_image_and_create_world_with_retry.rs b/crates/api_clients/world_labs_client/src/recipes/upload_image_and_create_world_with_retry.rs index 3dd0080f4f..01983301a5 100644 --- a/crates/api_clients/world_labs_client/src/recipes/upload_image_and_create_world_with_retry.rs +++ b/crates/api_clients/world_labs_client/src/recipes/upload_image_and_create_world_with_retry.rs @@ -1,3 +1,4 @@ +use std::path::PathBuf; use crate::api::api_types::run_object_id::RunObjectId; use crate::api::api_types::upload_mime_type::UploadMimeType; use crate::api::requests::google_upload::google_upload_image::{google_upload_image, GoogleUploadImageArgs}; @@ -14,6 +15,7 @@ use log::{error, info}; use serde::Deserialize; use std::time::Duration; use uuid::Uuid; +use filesys::file_read_bytes::file_read_bytes; use crate::api::api_types::image_input_object_id::ImageInputObjectId; use crate::api::api_types::pano_object_id::PanoObjectId; use crate::api::api_types::meta_world_object_id::MetaWorldObjectId; @@ -22,12 +24,18 @@ use crate::api::requests::objects::update_run_object_with_upload::{update_run_ob use crate::api::requests::objects::update_run_object_with_world::{update_run_object_with_world, UpdateRunObjectWithWorldArgs, UpdateRunObjectWithWorldPayloadArgs}; use crate::api::requests::recaption::recaption_image::{recaption_image, RecaptionImageArgs}; use crate::api::requests::worlds::create_world::{create_world, CreateWorldArgs}; +use crate::error::world_labs_client_error::WorldLabsClientError; pub struct UploadImageAndCreateWorldWithRetryArgs<'a> { pub cookies: &'a WorldLabsCookies, pub bearer_token: &'a WorldLabsBearerToken, pub individual_request_timeout: Option, - pub file_bytes: Vec, + pub file: FileBytesOrPath, +} + +pub enum FileBytesOrPath { + Bytes(Vec), + Path(PathBuf), } pub struct UploadImageAndCreateWorldWithRetryResponse { @@ -41,6 +49,25 @@ pub struct UploadImageAndCreateWorldWithRetryResponse { /// Request #5 (of ~10) pub async fn upload_image_and_create_world_with_retry(args: UploadImageAndCreateWorldWithRetryArgs<'_>) -> Result { + info!("Checking file input..."); + + let file_bytes = match args.file { + FileBytesOrPath::Bytes(bytes) => { + info!("File bytes provided"); + bytes + } + FileBytesOrPath::Path(path) => { + info!("File path provided"); + match file_read_bytes(&path) { + Ok(bytes) => bytes, + Err(err) => { + error!("Error reading file bytes from path: {:?} - error: {:?}", path, err); + return Err(WorldLabsClientError::CannotReadLocalFileForUpload(err).into()); + } + } + } + }; + info!("Request #1 of 10: create run object ..."); let response = create_run_object(CreateRunObjectArgs { @@ -89,7 +116,7 @@ pub async fn upload_image_and_create_world_with_retry(args: UploadImageAndCreate let _response = google_upload_image(GoogleUploadImageArgs { upload_url: &google_upload_url, upload_mime_type, - file_bytes: args.file_bytes, + file_bytes, request_timeout: args.individual_request_timeout, }).await?; @@ -188,7 +215,7 @@ pub async fn upload_image_and_create_world_with_retry(args: UploadImageAndCreate #[cfg(test)] mod tests { - use crate::recipes::upload_image_and_create_world_with_retry::{upload_image_and_create_world_with_retry, UploadImageAndCreateWorldWithRetryArgs}; + use crate::recipes::upload_image_and_create_world_with_retry::{upload_image_and_create_world_with_retry, FileBytesOrPath, UploadImageAndCreateWorldWithRetryArgs}; use crate::test_utils::get_test_bearer_token::get_test_bearer_token; use crate::test_utils::get_test_cookies::get_typed_test_cookies; use crate::test_utils::setup_test_logging::setup_test_logging; @@ -213,7 +240,7 @@ mod tests { cookies: &cookies, bearer_token: &bearer_token, individual_request_timeout: None, - file_bytes: file_bytes, + file: FileBytesOrPath::Bytes(file_bytes), }).await.unwrap(); println!("Upload URL: {}", results.image_upload_url); diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index 053bd3a9be..59f8741d92 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -45,6 +45,7 @@ storyteller_client.workspace = true tokens.workspace = true url_utils.workspace = true web_base64.workspace = true +world_labs_client.workspace = true # External anyhow.workspace = true diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs index 28b2e501e5..3d9a1dd73b 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/generate_error.rs @@ -6,6 +6,7 @@ use grok_client::error::grok_error::GrokError; use midjourney_client::error::midjourney_error::MidjourneyError; use openai_sora_client::error::sora_error::SoraError; use storyteller_client::error::storyteller_error::StorytellerError; +use world_labs_client::error::world_labs_error::WorldLabsError; //use fal_client::error::fal_error_plus::FalErrorPlus; #[derive(Debug)] @@ -62,6 +63,7 @@ pub enum MissingCredentialsReason { NeedsMidjourneyUserInfo, NeedsSoraCredentials, NeedsStorytellerCredentials, + NeedsWorldLabsCredentials, } #[derive(Debug)] @@ -88,6 +90,7 @@ pub enum ProviderFailureReason { MidjourneyJobEnqueueFailed, SoraError(SoraError), StorytellerError(StorytellerError), + WorldLabsError(WorldLabsError), } impl GenerateError { @@ -187,3 +190,9 @@ impl From for GenerateError { Self::ProviderFailure(ProviderFailureReason::StorytellerError(value)) } } + +impl From for GenerateError { + fn from(value: WorldLabsError) -> Self { + Self::ProviderFailure(ProviderFailureReason::WorldLabsError(value)) + } +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs new file mode 100644 index 0000000000..5722069891 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/enqueue_image_to_gaussian_command.rs @@ -0,0 +1,233 @@ +use crate::core::commands::enqueue::common::notify_frontend_of_errors::notify_frontend_of_errors; +use crate::core::commands::enqueue::generate_error::{BadInputReason, GenerateError, MissingCredentialsReason}; +use crate::core::commands::enqueue::image_to_gaussian::worldlabs::handle_worldlabs_marble::handle_worldlabs_marble; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus}; +use crate::core::commands::response::shorthand::Response; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::credits_balance_changed_event::CreditsBalanceChangedEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent; +use crate::core::events::sendable_event_trait::SendableEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::provider_priority::{Provider, ProviderPriorityStore}; +use crate::core::state::task_database::TaskDatabase; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; +use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::state::sora_task_queue::SoraTaskQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_status::TaskStatus; +use enums::tauri::tasks::task_type::TaskType; +use enums::tauri::ux::tauri_command_caller::TauriCommandCaller; +use log::{error, info, warn}; +use serde::{Deserialize, Serialize}; +use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs}; +use tauri::{AppHandle, State}; +use tokens::tokens::media_files::MediaFileToken; + +/// This is used in the Tauri command bridge. +/// Don't change the serializations without coordinating with the frontend. +#[derive(Deserialize, Debug, Copy, Clone)] +#[serde(rename_all = "snake_case")] +pub enum GaussianModel { + #[serde(rename = "world_labs_marble")] + WorldLabsMarble, +} + +#[derive(Deserialize)] +pub struct EnqueueImageToGaussianRequest { + /// REQUIRED. + /// The model to use. + pub model: Option, + + /// OPTIONAL. + /// Text prompt for the gaussian generation. Required. + pub prompt: Option, + + /// OPTIONAL. + /// Images for the gaussian generation + pub image_media_tokens: Option>, + + /// OPTIONAL. + /// Name of the frontend caller. + /// We'll use this to selectively trigger events. + pub frontend_caller: Option, + + /// OPTIONAL. + /// A frontend-defined identifier that we'll send back to the frontend + /// as a Tauri event on task completion. + pub frontend_subscriber_id: Option, + + /// OPTIONAL. + /// A frontend-defined payload that we'll send back to the frontend + /// as a Tauri event on task completion. + pub frontend_subscriber_payload: Option, +} + + +#[derive(Serialize)] +pub struct EnqueueImageToGaussianSuccessResponse { +} + +impl SerializeMarker for EnqueueImageToGaussianSuccessResponse {} + +#[derive(Serialize, Debug)] +#[serde(rename_all = "snake_case")] +pub enum EnqueueImageToGaussianErrorType { + /// Caller didn't specify a model + ModelNotSpecified, + /// No model available for image generation + NoProviderAvailable, + /// Generic server error + ServerError, + /// Needs to be logged into Artcraft + NeedsStorytellerCredentials, +} + + +#[tauri::command] +pub async fn enqueue_image_to_gaussian_command( + request: EnqueueImageToGaussianRequest, + app: AppHandle, + app_data_root: State<'_, AppDataRoot>, + app_env_configs: State<'_, AppEnvConfigs>, + task_database: State<'_, TaskDatabase>, + storyteller_creds_manager: State<'_, StorytellerCredentialManager>, + worldlabs_creds_manager: State<'_, WorldlabsCredentialManager>, +) -> Response { + + info!("enqueue_image_to_gaussian called"); + + let result = handle_request( + request, + &app, + &app_data_root, + &task_database, + &storyteller_creds_manager, + &worldlabs_creds_manager, + &app_env_configs, + ).await; + + match result { + Err(err) => { + error!("error: {:?}", err); + + notify_frontend_of_errors(&app, &err).await; + + let mut status = CommandErrorStatus::ServerError; + let mut error_type = EnqueueImageToGaussianErrorType::ServerError; + let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem."; + + match err { + GenerateError::BadInput(BadInputReason::NoModelSpecified) => { + status = CommandErrorStatus::BadRequest; + error_type = EnqueueImageToGaussianErrorType::ModelNotSpecified; + error_message = "No model specified for image generation"; + } + GenerateError::NoProviderAvailable => { + status = CommandErrorStatus::ServerError; + error_type = EnqueueImageToGaussianErrorType::NoProviderAvailable; + error_message = "No configured provider available for image generation"; + } + _ => {}, // Fall-through + } + + Err(CommandErrorResponseWrapper { + status, + error_message: Some(error_message.to_string()), + error_type: Some(error_type), + error_details: None, + }) + } + Ok(event) => { + let event = GenerationEnqueueSuccessEvent { + action: event.to_frontend_event_action(), + service: event.to_frontend_event_service(), + model: event.model, + }; + + if let Err(err) = event.send(&app) { + error!("Failed to emit event: {:?}", err); // Fail open. + } + + CreditsBalanceChangedEvent{}.send_infallible(&app); + + Ok(EnqueueImageToGaussianSuccessResponse {}.into()) + } + } +} + + +pub async fn handle_request( + request: EnqueueImageToGaussianRequest, + app: &AppHandle, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + worldlabs_creds_manager: &WorldlabsCredentialManager, + app_env_configs: &AppEnvConfigs, +) -> Result { + + let result = dispatch_request( + &request, + &app, + &app_data_root, + &storyteller_creds_manager, + &worldlabs_creds_manager, + &app_env_configs, + ).await; + + let success_event = match result { + Err(err) => return Err(err), + Ok(event) => event, + }; + + let result = success_event + .insert_into_task_database_with_frontend_payload( + task_database, + request.frontend_caller, + request.frontend_subscriber_id.as_deref(), + request.frontend_subscriber_payload.as_deref(), + ) + .await; + + if let Err(err) = result { + error!("Failed to create task in database: {:?}", err); + // NB: Fail open, but find a way to flag this. + } + + Ok(success_event) +} + +pub async fn dispatch_request( + request: &EnqueueImageToGaussianRequest, + app: &AppHandle, + app_data_root: &AppDataRoot, + storyteller_creds_manager: &StorytellerCredentialManager, + worldlabs_creds_manager: &WorldlabsCredentialManager, + app_env_configs: &AppEnvConfigs, +) -> Result { + + match request.model { + None => { + return Err(GenerateError::BadInput(BadInputReason::NoModelSpecified)); + } + Some(GaussianModel::WorldLabsMarble) => { + return handle_worldlabs_marble( + app, + app_data_root, + app_env_configs, + request, + worldlabs_creds_manager, + ).await; + } + }; + + Err(GenerateError::NoProviderAvailable) +} diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/mod.rs new file mode 100644 index 0000000000..b4487f3526 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/mod.rs @@ -0,0 +1,2 @@ +pub mod enqueue_image_to_gaussian_command; +mod worldlabs; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs new file mode 100644 index 0000000000..977036ac15 --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/handle_worldlabs_marble.rs @@ -0,0 +1,215 @@ +use crate::core::commands::enqueue::generate_error::{BadInputReason, GenerateError, MissingCredentialsReason}; +use crate::core::commands::enqueue::image_to_gaussian::enqueue_image_to_gaussian_command::EnqueueImageToGaussianRequest; +use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess; +use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize}; +use crate::core::events::generation_events::common::GenerationModel; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::utils::get_url_file_extension::get_url_file_extension; +use crate::core::utils::simple_http_download::simple_http_download; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use enums::common::generation_provider::GenerationProvider; +use enums::tauri::tasks::task_type::TaskType; +use idempotency::uuid::generate_random_uuid; +use log::{error, info}; +use std::path::PathBuf; +use storyteller_client::endpoints::media_files::get_media_file::get_media_file; +use tauri::AppHandle; +use tokens::tokens::media_files::MediaFileToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; +use world_labs_client::recipes::upload_image_and_create_world_with_retry::{upload_image_and_create_world_with_retry, FileBytesOrPath, UploadImageAndCreateWorldWithRetryArgs}; + +pub(super) const MAX_IMAGES: usize = 10; + +pub async fn handle_worldlabs_marble( + app: &AppHandle, + app_data_root: &AppDataRoot, + app_env_configs: &AppEnvConfigs, + request: &EnqueueImageToGaussianRequest, + worldlabs_creds_manager: &WorldlabsCredentialManager, +) -> Result { + + let world_labs_cookies = match worldlabs_creds_manager.maybe_copy_typed_cookies()? { + Some(cookies) => cookies, + None => { + error!("No WorldLabs cookies!"); + return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsWorldLabsCredentials)); + } + }; + + let world_labs_bearer = match worldlabs_creds_manager.maybe_copy_bearer_token()? { + Some(bearer) => bearer, + None => { + error!("No WorldLabs bearer token!"); + return Err(GenerateError::MissingCredentials(MissingCredentialsReason::NeedsWorldLabsCredentials)); + } + }; + + let job_id = generate_random_uuid(); + + let maybe_prompt = request.prompt + .as_deref() + .map(|prompt| prompt.trim().to_string()); + + info!("Downloading image file..."); + + let file_path = download_file( + app_data_root, + app_env_configs, + request, + ).await?; + + info!("Enqueueing WorldLabs request..."); + + let response = upload_image_and_create_world_with_retry(UploadImageAndCreateWorldWithRetryArgs { + cookies: &world_labs_cookies, + bearer_token: &world_labs_bearer, + individual_request_timeout: None, + file: FileBytesOrPath::Path(file_path), + }).await?; + + let run_id = response.run_id; + let world_id = response.world_id; + + info!("Run ID: {:?}", run_id); + info!("World ID: {:?}", world_id); + + Ok(TaskEnqueueSuccess { + provider: GenerationProvider::WorldLabs, + model: Some(GenerationModel::WorldlabsMarble), + provider_job_id: Some(world_id.0), + task_type: TaskType::ImageGeneration, + }) +} + +async fn download_file( + app_data_root: &AppDataRoot, + app_env_configs: &AppEnvConfigs, + request: &EnqueueImageToGaussianRequest, +) -> Result { + + let maybe_media_token = request.image_media_tokens + .as_ref() + .and_then(|tokens| tokens.get(0)); + + let media_token = match maybe_media_token { + Some(token) => token, + None => { + return Err(GenerateError::BadInput(BadInputReason::InvalidNumberOfInputImages { + min: 1, + max: MAX_IMAGES as u32, + provided: 0, + })); + } + }; + + info!("Calling get media file API: {:?}", app_env_configs.storyteller_host); + + // TODO(bt,2025-12-18): Add a cache. + info!("Using media token: {:?}", media_token); + + let response = get_media_file( + &app_env_configs.storyteller_host, + media_token + ).await?; + + let media_file_url = &response.media_file.media_links.cdn_url; + let extension_with_dot = get_url_file_extension(media_file_url) + .map(|ext| format!(".{}", ext)) + .unwrap_or_else(|| ".png".to_string()); + + let filename = format!("{}{}", response.media_file.token.as_str(), extension_with_dot); + let filename = app_data_root.downloads_dir().path().join(&filename); + + simple_http_download(&media_file_url, &filename).await?; + + Ok(filename) +} + +//async fn download_files( +// app_data_root: &AppDataRoot, +// app_env_configs: &AppEnvConfigs, +// request: &EnqueueImageToGaussianRequest, +//) -> Result, GenerateError> { +// +// let maybe_media_token = request.image_media_tokens +// .as_ref() +// .and_then(|tokens| tokens.get(0)); +// +// +// let media_tokens = match request.image_media_tokens.as_ref() { +// Some(media_tokens) => media_tokens.to_vec(), +// None => { +// return Err(GenerateError::BadInput(BadInputReason::InvalidNumberOfInputImages { +// min: 1, +// max: MAX_IMAGES as u32, +// provided: 0, +// })); +// } +// }; +// +// if media_tokens.len() > MAX_IMAGES { +// return Err(GenerateError::BadInput(BadInputReason::InvalidNumberOfInputImages { +// min: 1, +// max: MAX_IMAGES as u32, +// provided: media_tokens.len() as u32, +// })); +// } +// +// info!("Calling get media file API: {:?}", app_env_configs.storyteller_host); +// +// let mut local_files_to_upload = Vec::with_capacity(10); +// +// // TODO(bt,2025-07-07): This is inefficient. Cache and parallelize this. +// // TODO(bt,2025-12-18): Add a cache. +// for media_token in media_tokens.iter() { +// info!("Using media token: {:?}", media_token); +// +// let response = get_media_file( +// &app_env_configs.storyteller_host, +// media_token +// ).await?; +// +// let media_file_url = &response.media_file.media_links.cdn_url; +// let extension_with_dot = get_url_file_extension(media_file_url) +// .map(|ext| format!(".{}", ext)) +// .unwrap_or_else(|| ".png".to_string()); +// +// let filename = format!("{}{}", response.media_file.token.as_str(), extension_with_dot); +// let filename = app_data_root.downloads_dir().path().join(&filename); +// +// simple_http_download(&media_file_url, &filename).await?; +// +// local_files_to_upload.push(filename); +// } +// +// Ok(local_files_to_upload) +//} + +// TODO: +// fn handle_midjourney_errors( +// app: &AppHandle, +// maybe_errors: Option> +// ) -> Result { +// if let Some(errors) = maybe_errors { +// if !errors.is_empty() { +// let messages: Vec = errors.iter() +// .map(|e| format!("{:?}", e)) +// .collect(); +// +// let combined_message = messages.join("; "); +// +// let event = FlashUserInputErrorEvent { +// message: format!("Midjourney Error: {}", combined_message), +// }; +// +// if let Err(err) = event.send(&app) { +// error!("Failed to send FlashUserInputErrorEvent: {:?}", err); // Fail open +// } +// } +// } +// +// Err(GenerateError::ProviderFailure(ProviderFailureReason::MidjourneyJobEnqueueFailed)) +// } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/mod.rs new file mode 100644 index 0000000000..73ec59edec --- /dev/null +++ b/crates/desktop/artcraft/src/core/commands/enqueue/image_to_gaussian/worldlabs/mod.rs @@ -0,0 +1 @@ +pub (super) mod handle_worldlabs_marble; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/mod.rs b/crates/desktop/artcraft/src/core/commands/enqueue/mod.rs index 9ec6f747be..b2f9aa2fc3 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/mod.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/mod.rs @@ -6,4 +6,5 @@ pub mod image_inpaint; pub mod image_to_object; pub mod image_to_video; pub mod task_enqueue_success; -pub mod text_to_image; \ No newline at end of file +pub mod text_to_image; +pub mod image_to_gaussian; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs index ac76434588..176e788178 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/task_enqueue_success.rs @@ -21,6 +21,7 @@ impl TaskEnqueueSuccess{ pub fn to_frontend_event_action(&self) -> GenerationAction { match self.task_type { TaskType::ImageGeneration => GenerationAction::GenerateImage, + TaskType::GaussianGeneration => GenerationAction::GenerateGaussian, TaskType::VideoGeneration => GenerationAction::GenerateVideo, TaskType::BackgroundRemoval => GenerationAction::RemoveBackground, TaskType::ObjectGeneration => GenerationAction::ImageTo3d, @@ -35,6 +36,7 @@ impl TaskEnqueueSuccess{ GenerationProvider::Grok => GenerationServiceProvider::Grok, GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney, GenerationProvider::Sora => GenerationServiceProvider::Sora, + GenerationProvider::WorldLabs => GenerationServiceProvider::WorldLabs, } } @@ -81,6 +83,7 @@ impl TaskEnqueueSuccess{ Some(GenerationModel::Veo3Fast) => Some(TaskModelType::Veo3Fast), Some(GenerationModel::Hunyuan3d2_0) => Some(TaskModelType::Hunyuan3d2_0), Some(GenerationModel::Hunyuan3d2_1) => Some(TaskModelType::Hunyuan3d2_1), + Some(GenerationModel::WorldlabsMarble) => Some(TaskModelType::WorldlabsMarble), Some(GenerationModel::Midjourney) => Some(TaskModelType::Midjourney), // NB: This is a generic Midjourney model, version unknown. // TODO: These seem wrong - diff --git a/crates/desktop/artcraft/src/core/events/functional_events/gaussian_generation_complete_event.rs b/crates/desktop/artcraft/src/core/events/functional_events/gaussian_generation_complete_event.rs new file mode 100644 index 0000000000..51ffda4794 --- /dev/null +++ b/crates/desktop/artcraft/src/core/events/functional_events/gaussian_generation_complete_event.rs @@ -0,0 +1,26 @@ +use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent}; +use enums::tauri::ux::tauri_event_name::TauriEventName; +use serde_derive::Serialize; +use tokens::tokens::media_files::MediaFileToken; +use url::Url; + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct GaussianGenerationCompleteEvent { + pub generated_gaussian: Option, + pub maybe_frontend_subscriber_id: Option, + pub maybe_frontend_subscriber_payload: Option, +} + +#[derive(Clone, Debug, Serialize)] +#[serde(rename_all = "snake_case")] +pub struct GeneratedGaussian { + pub media_token: MediaFileToken, + pub cdn_url: Url, + pub maybe_thumbnail_template: Option, +} + +impl BasicSendableEvent for GaussianGenerationCompleteEvent { + const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GaussianGenerationCompleteEvent; + const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success; +} diff --git a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs index 605d0a82c2..f6f1773d22 100644 --- a/crates/desktop/artcraft/src/core/events/functional_events/mod.rs +++ b/crates/desktop/artcraft/src/core/events/functional_events/mod.rs @@ -1,5 +1,6 @@ pub mod canvas_background_removal_complete_event; pub mod credits_balance_changed_event; +pub mod gaussian_generation_complete_event; pub mod image_edit_complete_event; pub mod object_generation_complete_event; pub mod refresh_account_state_event; diff --git a/crates/desktop/artcraft/src/core/events/generation_events/common.rs b/crates/desktop/artcraft/src/core/events/generation_events/common.rs index eb6ec67c53..93c358fca6 100644 --- a/crates/desktop/artcraft/src/core/events/generation_events/common.rs +++ b/crates/desktop/artcraft/src/core/events/generation_events/common.rs @@ -8,6 +8,7 @@ pub enum GenerationServiceProvider { Grok, Midjourney, Sora, + WorldLabs, } #[derive(Debug, Copy, Clone, Serialize)] @@ -44,6 +45,9 @@ pub enum GenerationModel { #[serde(rename = "hunyuan_3d_2_1")] Hunyuan3d2_1, + #[serde(rename = "worldlabs_marble")] + WorldlabsMarble, + // Generic Midjourney model, version unknown. #[serde(rename = "midjourney")] Midjourney, @@ -84,6 +88,7 @@ pub enum GenerationAction { RemoveBackground, #[serde(rename = "image_to_3d")] ImageTo3d, + GenerateGaussian, #[serde(rename = "image_inpaint_edit")] ImageInpaintEdit, } diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index 3c33ef011b..bbdb7b3f06 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs @@ -18,6 +18,9 @@ use crate::services::midjourney::threads::midjourney_long_polling_thread::midjou use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; use crate::services::sora::state::sora_task_queue::SoraTaskQueue; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use crate::services::worldlabs::threads::worldlabs_marble_task_polling::worldlabs_marble_task_polling; use errors::AnyhowResult; use tauri::{AppHandle, Manager}; @@ -32,6 +35,8 @@ pub async fn handle_tauri_startup( mj_creds_manager: MidjourneyCredentialManager, grok_creds_manager: GrokCredentialManager, grok_image_prompt_queue: GrokImagePromptQueue, + worldlabs_bearer_bridge: WorldlabsBearerBridge, + worldlabs_creds_manager: WorldlabsCredentialManager, ) -> AnyhowResult<()> { set_app_log_level( @@ -99,6 +104,15 @@ pub async fn handle_tauri_startup( storyteller_creds_manager.clone(), )); + tauri::async_runtime::spawn(worldlabs_marble_task_polling( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + worldlabs_creds_manager.clone(), + storyteller_creds_manager.clone(), + )); + spawn_discord_presence_thread()?; initially_size_and_position_windows(&app, &root); diff --git a/crates/desktop/artcraft/src/core/state/data_dir/app_credentials_dir.rs b/crates/desktop/artcraft/src/core/state/data_dir/app_credentials_dir.rs index 3e52aea794..b7fe13b1af 100644 --- a/crates/desktop/artcraft/src/core/state/data_dir/app_credentials_dir.rs +++ b/crates/desktop/artcraft/src/core/state/data_dir/app_credentials_dir.rs @@ -61,4 +61,12 @@ impl AppCredentialsDir { pub fn get_midjourney_state_path(&self) -> PathBuf { self.path.join("midjourney_state.json") } + + pub fn get_worldlabs_state_path(&self) -> PathBuf { + self.path.join("worldlabs_state.json") + } + + pub fn get_worldlabs_cookies_path(&self) -> PathBuf { + self.path.join("worldlabs_cookies.txt") + } } diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs index 1ec6cc92c3..991747a68e 100644 --- a/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/generation_provider.rs @@ -9,5 +9,6 @@ pub fn to_generation_service_provider(provider: GenerationProvider) -> Generatio GenerationProvider::Grok => GenerationServiceProvider::Grok, GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney, GenerationProvider::Sora => GenerationServiceProvider::Sora, + GenerationProvider::WorldLabs => GenerationServiceProvider::WorldLabs, } } diff --git a/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs b/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs index 55a4c63584..79054dbc2b 100644 --- a/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs +++ b/crates/desktop/artcraft/src/core/utils/enum_conversion/task_type.rs @@ -7,6 +7,7 @@ pub fn to_generation_action(task_type: TaskType) -> GenerationAction { TaskType::VideoGeneration => GenerationAction::GenerateVideo, TaskType::BackgroundRemoval => GenerationAction::RemoveBackground, TaskType::ObjectGeneration => GenerationAction::ImageTo3d, + TaskType::GaussianGeneration => GenerationAction::GenerateGaussian, TaskType::ImageInpaintEdit => GenerationAction::ImageInpaintEdit, } } diff --git a/crates/desktop/artcraft/src/core/utils/window/get_webview_window_url_path.rs b/crates/desktop/artcraft/src/core/utils/window/get_webview_window_url_path.rs new file mode 100644 index 0000000000..0cdec49494 --- /dev/null +++ b/crates/desktop/artcraft/src/core/utils/window/get_webview_window_url_path.rs @@ -0,0 +1,7 @@ +use errors::AnyhowResult; +use tauri::WebviewWindow; + +pub fn get_webview_window_url_path(webview: &WebviewWindow) -> AnyhowResult { + let url = webview.url()?; + Ok(url.path().to_string()) +} diff --git a/crates/desktop/artcraft/src/core/utils/window/mod.rs b/crates/desktop/artcraft/src/core/utils/window/mod.rs index 11abd59ee4..bc3d5e1a84 100644 --- a/crates/desktop/artcraft/src/core/utils/window/mod.rs +++ b/crates/desktop/artcraft/src/core/utils/window/mod.rs @@ -2,3 +2,4 @@ pub mod get_webview_window_hostname; pub mod get_window_size_heuristic; pub mod position_main_window; pub mod resize_main_window; +pub mod get_webview_window_url_path; diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index ccf8ef10d0..4a587dd280 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -8,6 +8,7 @@ use crate::core::commands::app_preferences::update_app_preference_command::updat use crate::core::commands::enqueue::image_bg_removal::enqueue_image_bg_removal_command::enqueue_image_bg_removal_command; use crate::core::commands::enqueue::image_edit::enqueue_edit_image_command::enqueue_edit_image_command; use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::enqueue_image_inpaint_command; +use crate::core::commands::enqueue::image_to_gaussian::enqueue_image_to_gaussian_command::enqueue_image_to_gaussian_command; use crate::core::commands::enqueue::image_to_object::enqueue_image_to_3d_object_command::enqueue_image_to_3d_object_command; use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::enqueue_image_to_video_command; use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::enqueue_text_to_image_command; @@ -32,8 +33,8 @@ use crate::core::threads::main_window_thread::main_window_thread::main_window_th use crate::services::grok::commands::grok_clear_credentials_command::grok_clear_credentials_command; use crate::services::grok::commands::grok_get_credential_info_command::grok_get_credential_info_command; use crate::services::grok::commands::grok_open_login_command::grok_open_login_command; -use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command; use crate::services::midjourney::commands::midjourney_get_credential_info_command::midjourney_get_credential_info_command; use crate::services::midjourney::commands::midjourney_open_login_command::midjourney_open_login_command; @@ -55,6 +56,12 @@ use crate::services::storyteller::commands::stripe_customer_portal::storyteller_ use crate::services::storyteller::commands::stripe_customer_portal::storyteller_open_customer_portal_switch_plan_command::storyteller_open_customer_portal_switch_plan_command; use crate::services::storyteller::commands::stripe_customer_portal::storyteller_open_customer_portal_update_payment_method_command::storyteller_open_customer_portal_update_payment_method_command; use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::commands::worldlabs_clear_credentials_command::worldlabs_clear_credentials_command; +use crate::services::worldlabs::commands::worldlabs_get_credential_info_command::worldlabs_get_credential_info_command; +use crate::services::worldlabs::commands::worldlabs_open_login_command::worldlabs_open_login_command; +use crate::services::worldlabs::commands::worldlabs_receive_bearer_command::worldlabs_receive_bearer_command; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; use log::error; use tauri_plugin_dialog; @@ -108,6 +115,12 @@ pub fn run() { let grok_prompt_queue = GrokImagePromptQueue::new(); let grok_prompt_queue_2 = grok_prompt_queue.clone(); + let worldlabs_creds_manager = WorldlabsCredentialManager::initialize_from_disk_infallible(&app_data_root); + let worldlabs_creds_manager_2 = worldlabs_creds_manager.clone(); + + let worldlabs_bearer_bridge = WorldlabsBearerBridge::empty(); + let worldlabs_bearer_bridge_2 = worldlabs_bearer_bridge.clone(); + println!("Initializing backend runtime..."); let builder = tauri::Builder::default() @@ -147,6 +160,8 @@ pub fn run() { midjourney_creds_manager_2, grok_creds_manager_2, grok_prompt_queue_2, + worldlabs_bearer_bridge_2, + worldlabs_creds_manager_2, ).await; if let Err(err) = result { @@ -166,7 +181,9 @@ pub fn run() { .manage(midjourney_creds_manager) .manage(sora_creds_manager) .manage(sora_task_queue) - .manage(storyteller_creds_manager_3); + .manage(storyteller_creds_manager_3) + .manage(worldlabs_bearer_bridge) + .manage(worldlabs_creds_manager); // TODO: Break this out into another module, because RustRover/IntelliJ lags with these macros. // My first attempt at naively doing this didn't work because the macros can't find their codegen'd targets. @@ -176,6 +193,7 @@ pub fn run() { enqueue_image_bg_removal_command, enqueue_image_inpaint_command, enqueue_image_to_3d_object_command, + enqueue_image_to_gaussian_command, enqueue_image_to_video_command, enqueue_text_to_image_command, flip_image, @@ -207,6 +225,10 @@ pub fn run() { storyteller_open_subscription_purchase_command, storyteller_purge_credentials_command, update_app_preferences_command, + worldlabs_clear_credentials_command, + worldlabs_get_credential_info_command, + worldlabs_open_login_command, + worldlabs_receive_bearer_command, ]); builder.run(tauri::generate_context!("tauri.conf.json")) diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs b/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs index cb1aed6817..4379a4c1f9 100644 --- a/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs +++ b/crates/desktop/artcraft/src/services/grok/state/grok_serializable_state.rs @@ -3,6 +3,8 @@ use cookie_store::serialized_cookie_store::SerializableCookieStore; use serde_derive::{Deserialize, Serialize}; use std::fs::read_to_string; +/// History +/// v1 - initial pub (super) const SERIALIZABLE_GROK_STATE_VERSION: u32 = 1; /// We only need to serialize Grok cookies. diff --git a/crates/desktop/artcraft/src/services/mod.rs b/crates/desktop/artcraft/src/services/mod.rs index f7480cf88b..cea58b5941 100644 --- a/crates/desktop/artcraft/src/services/mod.rs +++ b/crates/desktop/artcraft/src/services/mod.rs @@ -1,4 +1,5 @@ pub mod grok; pub mod midjourney; pub mod sora; -pub mod storyteller; \ No newline at end of file +pub mod storyteller; +pub mod worldlabs; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs index 2484f32e1b..353396cfb9 100644 --- a/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs +++ b/crates/desktop/artcraft/src/services/storyteller/threads/events/maybe_handle_frontend_caller_notification.rs @@ -1,4 +1,6 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; +use crate::core::events::functional_events::gaussian_generation_complete_event::{GaussianGenerationCompleteEvent, GeneratedGaussian}; use crate::core::events::functional_events::image_edit_complete_event::{EditedImage, ImageEditCompleteEvent}; use crate::core::events::functional_events::object_generation_complete_event::{GeneratedObject, ObjectGenerationCompleteEvent}; use crate::core::events::functional_events::text_to_image_generation_complete_event::{GeneratedImage, TextToImageGenerationCompleteEvent}; @@ -18,7 +20,6 @@ use storyteller_client::endpoints::media_files::list_batch_generated_redux_media use tauri::AppHandle; use tokens::tokens::batch_generations::BatchGenerationToken; use tokens::tokens::media_files::MediaFileToken; -use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; pub async fn maybe_handle_frontend_caller_notification( app: &AppHandle, @@ -69,6 +70,13 @@ pub async fn maybe_handle_frontend_caller_notification( job_result, ).await?; } + TaskType::GaussianGeneration => { + let _r = handle_gaussian_generation( + app, + task, + job_result, + ).await?; + } TaskType::BackgroundRemoval => { let _r = handle_background_removal_generation( app, @@ -225,6 +233,29 @@ async fn handle_object_generation( Ok(()) } +async fn handle_gaussian_generation( + app: &AppHandle, + task: &Task, + job_result: &ListSessionResultDetailsResponse, +) -> AnyhowResult<()> { + + // NB: For now, we only generate one object (gaussian) at a time. + let event = GaussianGenerationCompleteEvent { + generated_gaussian: Some(GeneratedGaussian { + media_token: MediaFileToken::new_from_str(&job_result.entity_token), + cdn_url: job_result.media_links.cdn_url.clone(), + maybe_thumbnail_template: media_links_to_thumbnail_template(&job_result.media_links) + .map(|s| s.to_owned()), + }), + maybe_frontend_subscriber_id: task.frontend_subscriber_id.clone(), + maybe_frontend_subscriber_payload: task.frontend_subscriber_payload.clone(), + }; + + event.send_infallible(&app); + + Ok(()) +} + async fn handle_background_removal_generation( app: &AppHandle, task: &Task, diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/mod.rs new file mode 100644 index 0000000000..77c3621ea1 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/mod.rs @@ -0,0 +1,4 @@ +pub mod worldlabs_open_login_command; +pub mod worldlabs_get_credential_info_command; +pub mod worldlabs_clear_credentials_command; +pub mod worldlabs_receive_bearer_command; diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_clear_credentials_command.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_clear_credentials_command.rs new file mode 100644 index 0000000000..f829c45249 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_clear_credentials_command.rs @@ -0,0 +1,44 @@ +use crate::core::commands::response::shorthand::{Response, SimpleResponse}; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use errors::AnyhowResult; +use log::{error, info}; +use serde_derive::Deserialize; +use tauri::State; + +#[tauri::command] +pub async fn worldlabs_clear_credentials_command( + root: State<'_, AppDataRoot>, + creds_manager: State<'_, WorldlabsCredentialManager>, +) -> SimpleResponse { + info!("worldlabs_clear_credentials_command called"); + + clear_creds(&root, &creds_manager) + .map_err(|err| { + error!("Error clearing creds: {:?}", err); + "error clearing creds" + })?; + + Ok(().into()) +} + +fn clear_creds( + root: &AppDataRoot, + creds: &WorldlabsCredentialManager, +) -> AnyhowResult<()> { + + info!("Clearing credentials..."); + creds.clear_credentials()?; + + info!("Persisting to disk..."); + creds.persist_to_disk()?; + + let creds_path = root.credentials_dir().get_worldlabs_state_path(); + + if creds_path.exists() { + info!("Removing from disk..."); + std::fs::remove_file(creds_path)?; + } + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs new file mode 100644 index 0000000000..4ebd87334d --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_get_credential_info_command.rs @@ -0,0 +1,63 @@ +use crate::core::commands::response::shorthand::{Response, ResponseOrErrorMessage}; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use errors::AnyhowResult; +use log::{error, info}; +use serde_derive::Serialize; +use tauri::State; +use world_labs_client::credentials::jwt_claims::JwtClaims; +use world_labs_client::error::world_labs_error::WorldLabsError; + +#[derive(Serialize)] +pub struct WorldlabsGetCredentialInfoResponse { + pub maybe_email: Option, + pub can_clear_state: bool, +} + +impl SerializeMarker for WorldlabsGetCredentialInfoResponse {} + +#[tauri::command] +pub async fn worldlabs_get_credential_info_command( + creds_manager: State<'_, WorldlabsCredentialManager>, +) -> ResponseOrErrorMessage { + info!("worldlabs_get_credential_info_command called"); + + let info = get_info(&creds_manager) + .map_err(|err| { + error!("Error getting info: {:?}", err); + "error getting info" + })?; + + Ok(info.into()) +} + +fn get_info( + creds: &WorldlabsCredentialManager, +) -> AnyhowResult { + let mut can_clear_state = true; + + let maybe_cookies = creds.maybe_copy_cookie_store()?; + let maybe_bearer = creds.maybe_copy_bearer_token()?; + + if maybe_cookies.is_none() && maybe_bearer.is_none() { + can_clear_state = false; + } + + let mut maybe_email = None; + + if let Some(bearer) = maybe_bearer { + match bearer.parse_jwt_claims() { + Ok(claims) => { + maybe_email = claims.email; + } + Err(err) => { + error!("Failed to parse JWT bearer claims: {}", err); + } + } + } + + Ok(WorldlabsGetCredentialInfoResponse { + maybe_email, + can_clear_state, + }) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_open_login_command.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_open_login_command.rs new file mode 100644 index 0000000000..e0a7b5b7bd --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_open_login_command.rs @@ -0,0 +1,45 @@ +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use crate::services::worldlabs::windows::worldlabs_login_window::worldlabs_login_window_open::worldlabs_login_window_open; +use errors::AnyhowResult; +use log::{error, info}; +use tauri::{AppHandle, State}; + +#[tauri::command] +pub async fn worldlabs_open_login_command( + app: AppHandle, + app_data_root: State<'_, AppDataRoot>, + worldlabs_bearer_bridge: State<'_, WorldlabsBearerBridge>, + worldlabs_creds_manager: State<'_, WorldlabsCredentialManager>, +) -> Result { + info!("worldlabs_open_login_command called"); + + do_open_login(&app, &app_data_root, &worldlabs_bearer_bridge, &worldlabs_creds_manager) + .await + .map_err(|err| { + error!("Error opening login: {:?}", err); + format!("Error opening login: {:?}", err) + })?; + + Ok("result".to_string()) +} + +async fn do_open_login( + app: &AppHandle, + app_data_root: &AppDataRoot, + worldlabs_bearer_bridge: &WorldlabsBearerBridge, + worldlabs_creds_manager: &WorldlabsCredentialManager, +) -> AnyhowResult<()> { + info!("Building login window..."); + + worldlabs_login_window_open( + app, + app_data_root, + worldlabs_bearer_bridge, + worldlabs_creds_manager + ).await?; + + info!("Done."); + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_receive_bearer_command.rs b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_receive_bearer_command.rs new file mode 100644 index 0000000000..0a9004df86 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/commands/worldlabs_receive_bearer_command.rs @@ -0,0 +1,59 @@ +use crate::core::commands::response::shorthand::{Response, SimpleResponse}; +use crate::core::commands::response::success_response_wrapper::SerializeMarker; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::{WorldlabsBearerBridge, WorldlabsBearerBridgeInner}; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use errors::AnyhowResult; +use log::{error, info}; +use serde_derive::{Deserialize, Serialize}; +use tauri::State; +use tokens::tokens::media_files::MediaFileToken; + + +#[derive(Deserialize, Debug)] +pub struct WorldlabsReceiveBearerRequest { + /// REQUIRED. + pub bearer_token: String, + + /// REQUIRED. + pub refresh_token: String, +} + +#[derive(Serialize)] +pub struct WorldlabsReceiveBearerResponse { +} + +impl SerializeMarker for WorldlabsReceiveBearerResponse {} + + +#[tauri::command] +pub async fn worldlabs_receive_bearer_command( + root: State<'_, AppDataRoot>, + request: WorldlabsReceiveBearerRequest, + bearer_bridge: State<'_, WorldlabsBearerBridge>, +) -> SimpleResponse { + info!("worldlabs_receive_bearer_command called"); + + info!("Request: {:?}", request); + + set_bearer(request, &bearer_bridge) + .map_err(|err| { + error!("Error setting bearer: {:?}", err); + "error setting bearer" + })?; + + Ok(().into()) +} + +fn set_bearer( + request: WorldlabsReceiveBearerRequest, + bearer_bridge: &WorldlabsBearerBridge, +) -> AnyhowResult<()> { + + bearer_bridge.set(WorldlabsBearerBridgeInner { + bearer_token: request.bearer_token, + refresh_token: request.refresh_token, + })?; + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/mod.rs new file mode 100644 index 0000000000..2393fb390d --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/mod.rs @@ -0,0 +1,4 @@ +pub mod commands; +pub mod state; +pub mod threads; +pub mod windows; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs new file mode 100644 index 0000000000..ef22d19e2b --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/state/mod.rs @@ -0,0 +1,6 @@ +pub mod worldlabs_credential_manager; +pub mod worldlabs_serializable_state; +pub mod worldlabs_bearer_bridge; + +// Private +mod worldlabs_credential_holder; diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_bearer_bridge.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_bearer_bridge.rs new file mode 100644 index 0000000000..f177a867ff --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_bearer_bridge.rs @@ -0,0 +1,60 @@ +use std::sync::{Arc, RwLock}; +use log::error; +use crate::core::artcraft_error::ArtcraftError; + +#[derive(Clone)] +pub struct WorldlabsBearerBridge { + state: Arc>> +} + +#[derive(Clone, Debug)] +pub struct WorldlabsBearerBridgeInner { + pub bearer_token: String, + pub refresh_token: String, +} + +impl WorldlabsBearerBridge { + pub fn empty() -> Self { + Self { + state: Arc::new(RwLock::new(None)), + } + } + + pub fn clear(&self) -> Result<(), ArtcraftError> { + match self.state.write() { + Ok(mut state) => { + *state = None; + Ok(()) + } + Err(err) => { + error!("Lock poisoned: {:?}", err); + Err(ArtcraftError::RwLockWriteError) + } + } + } + + pub fn get(&self) -> Result, ArtcraftError> { + match self.state.read() { + Ok(state) => { + Ok(state.clone()) + } + Err(err) => { + error!("Lock poisoned: {:?}", err); + Err(ArtcraftError::RwLockReadError) + } + } + } + + pub fn set(&self, new_state: WorldlabsBearerBridgeInner) -> Result<(), ArtcraftError> { + match self.state.write() { + Ok(mut state) => { + *state = Some(new_state); + Ok(()) + } + Err(err) => { + error!("Lock poisoned: {:?}", err); + Err(ArtcraftError::RwLockReadError) + } + } + } +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs new file mode 100644 index 0000000000..167fecf2e5 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_holder.rs @@ -0,0 +1,35 @@ +use cookie_store::cookie_store::CookieStore; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; +use world_labs_client::credentials::worldlabs_refresh_token::WorldLabsRefreshToken; + +#[derive(Clone)] +pub struct WorldlabsCredentialHolder { + /// Directly off the Tauri browser session. + /// Read once, write to disk + /// The Worldlabs client consumes a string-only form (rather than this cookie jar) + pub browser_cookies: Option, + + /// Copy of cookies, specifically for the client library. + /// NOT PERSISTED TO DISK. + pub world_labs_cookies: Option, + + /// Bearer token + /// This needs to be persisted to disk. + pub world_labs_bearer_token: Option, + + /// Refresh token + /// This needs to be persisted to disk. + pub world_labs_refresh_token: Option, +} + +impl WorldlabsCredentialHolder { + pub fn empty() -> Self { + Self { + browser_cookies: None, + world_labs_cookies: None, + world_labs_bearer_token: None, + world_labs_refresh_token: None, + } + } +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs new file mode 100644 index 0000000000..8f23128774 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs @@ -0,0 +1,203 @@ +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::services::worldlabs::state::worldlabs_credential_holder::WorldlabsCredentialHolder; +use crate::services::worldlabs::state::worldlabs_serializable_state::{WorldlabsSerializableState, SERIALIZABLE_WORLDLABS_STATE_VERSION}; +use cookie_store::cookie_store::CookieStore; +use errors::AnyhowResult; +use log::{info, warn}; +use std::fs::read_to_string; +use std::sync::{Arc, RwLock}; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; +use world_labs_client::credentials::worldlabs_refresh_token::WorldLabsRefreshToken; + +#[derive(Clone)] +pub struct WorldlabsCredentialManager { + // TODO: Put last write/last read timestamps on these. + credential_data: Arc>, + + //user_info: Arc>>, + app_data_root: AppDataRoot, +} + +impl WorldlabsCredentialManager { + pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self { + Self { + credential_data: Arc::new(RwLock::new(WorldlabsCredentialHolder::empty())), + app_data_root: app_data_root.clone(), + } + } + + pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self { + let mut credential_data; + + let result = WorldlabsSerializableState::read_from_disk(app_data_root); + match result { + Err(err) => { + warn!("Failed to read worldlabs state from disk: {:?}", err); + credential_data = Arc::new(RwLock::new(WorldlabsCredentialHolder::empty())); + }, + Ok(None) => { + credential_data = Arc::new(RwLock::new(WorldlabsCredentialHolder::empty())); + } + Ok(Some(state)) => { + let maybe_browser_cookies = state.user_cookies + .as_ref() + .map(|cookies| cookies.to_cookie_store()); + + let maybe_cookies = maybe_browser_cookies + .as_ref() + .map(|cookies| cookies.to_cookie_string()) + .map(|cookies| WorldLabsCookies::new(cookies)); + + let maybe_bearer = state.bearer_token + .map(|bearer| WorldLabsBearerToken::new(bearer)); + + let maybe_refresh = state.refresh_token + .map(|bearer| WorldLabsRefreshToken::new(bearer)); + + credential_data = Arc::new(RwLock::new(WorldlabsCredentialHolder { + browser_cookies: maybe_browser_cookies, + world_labs_cookies: maybe_cookies, + world_labs_bearer_token: maybe_bearer, + world_labs_refresh_token: maybe_refresh, + })); + } + }; + + Self { + credential_data, + app_data_root: app_data_root.clone(), + } + } + + pub fn maybe_copy_cookie_store(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.browser_cookies.clone()), + } + } + + pub fn maybe_copy_cookie_header_string(&self) -> anyhow::Result> { + let maybe_cookies = self.maybe_copy_cookie_store()?; + let maybe_cookies = maybe_cookies.map(|cookies| { + cookies.to_cookie_string() + }); + Ok(maybe_cookies) + } + + pub fn maybe_copy_typed_cookies(&self) -> anyhow::Result> { + let maybe_cookies = self.maybe_copy_cookie_header_string()?; + let maybe_cookies = maybe_cookies.map(|cookies| { + WorldLabsCookies::new(cookies) + }); + Ok(maybe_cookies) + } + + pub fn maybe_copy_bearer_token(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.world_labs_bearer_token.clone()), + } + } + + pub fn maybe_copy_refresh_token(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.world_labs_refresh_token.clone()), + } + } + + pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + holder.browser_cookies = Some(store); + Ok(()) + } + } + } + + pub fn replace_bearer_token(&self, token: WorldLabsBearerToken) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + holder.world_labs_bearer_token = Some(token); + Ok(()) + } + } + } + + pub fn replace_refresh_token(&self, token: WorldLabsRefreshToken) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut holder) => { + holder.world_labs_refresh_token = Some(token); + Ok(()) + } + } + } + + + // NB: This is just a heuristic. We'll add better checks later. + pub fn do_task_polling(&self) -> anyhow::Result { + self.session_appears_active() + } + + // NB: This is just a heuristic. We'll add better checks later. + pub fn session_appears_active(&self) -> anyhow::Result { + let holder = match self.credential_data.read() { + Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(store) => store.clone(), + }; + + if holder.world_labs_bearer_token.is_some() { + return Ok(true); + } + + if holder.browser_cookies.is_none() { + return Ok(false); + } + + Ok(false) + } + + pub fn clear_credentials(&self) -> anyhow::Result<()> { + match self.credential_data.write() { + Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), + Ok(mut store) => *store = WorldlabsCredentialHolder::empty(), + } + Ok(()) + } + + pub fn persist_to_disk(&self) -> anyhow::Result<()> { + let creds = match self.credential_data.read() { + Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(store) => store.clone(), + }; + + let state = WorldlabsSerializableState { + version: SERIALIZABLE_WORLDLABS_STATE_VERSION, + user_cookies: creds.browser_cookies + .as_ref() + .map(|cookies| cookies.to_serializable()), + user_id: None, // TODO + user_email: None, // TODO + bearer_token: creds.world_labs_bearer_token.map(|token| token.to_raw_string()), + refresh_token: creds.world_labs_refresh_token.map(|token| token.to_raw_string()), + }; + + let path = self.app_data_root.credentials_dir().get_worldlabs_state_path(); + let serialized = serde_json::to_string(&state)?; + + std::fs::write(path, serialized)?; + + // TODO: This is just for building the client and testing. + if let Some(cookies) = creds.browser_cookies.as_ref() { + let cookies_header = cookies.to_cookie_string(); + let path = self.app_data_root.credentials_dir().get_worldlabs_cookies_path(); + std::fs::write(path, cookies_header)?; + } + + Ok(()) + } +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs new file mode 100644 index 0000000000..cbf57429ba --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_serializable_state.rs @@ -0,0 +1,49 @@ +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use cookie_store::serialized_cookie_store::SerializableCookieStore; +use serde_derive::{Deserialize, Serialize}; +use std::fs::read_to_string; + +/// History +/// v1 - initial +pub (super) const SERIALIZABLE_WORLDLABS_STATE_VERSION: u32 = 1; + +/// We only need to serialize Grok cookies. +/// Most of the other state (verification_token, etc.) can be pulled at runtime. +/// We will cache the `user_id` and `user_email`. +#[derive(Serialize, Deserialize)] +pub (super) struct WorldlabsSerializableState { + pub (super) version: u32, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) user_cookies: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) bearer_token: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) refresh_token: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) user_id: Option, + + #[serde(skip_serializing_if = "Option::is_none")] + pub (super) user_email: Option, +} + +impl WorldlabsSerializableState { + pub fn read_from_disk(root: &AppDataRoot) -> anyhow::Result> { + let state_path = root.credentials_dir().get_worldlabs_state_path(); + + if !state_path.exists() { + return Ok(None); + } + + let contents = read_to_string(&state_path)? + .trim() + .to_string(); + + let state: WorldlabsSerializableState = serde_json::from_str(&contents)?; + + Ok(Some(state)) + } +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs new file mode 100644 index 0000000000..e03acdf874 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs @@ -0,0 +1 @@ +pub mod worldlabs_marble_task_polling; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs new file mode 100644 index 0000000000..ef24dde5f0 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs @@ -0,0 +1,386 @@ +use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason}; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::util::get_or_upgrade_grok_full_credentials::get_or_update_grok_full_credentials; +use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; +use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; +use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +use cookie_store::cookie_store::CookieStore; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_status::TaskStatus; +use errors::AnyhowResult; +use futures::poll; +use grok_client::credentials::grok_full_credentials::GrokFullCredentials; +use grok_client::error::grok_error::GrokError; +use grok_client::requests::download_video_file::download_video_file::{download_video_file, DownloadVideoFileArgs}; +use grok_client::requests::download_video_file::grok_download_video::GrokDownloadVideo; +use grok_client::requests::media_posts::list_media_posts::grok_list_media_posts::{GrokMediaPostList, GrokMediaPostListRequest, VideoData}; +use idempotency::uuid::generate_random_uuid; +use log::{error, info}; +use midjourney_client::client::midjourney_hostname::MidjourneyHostname; +use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId; +use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest, MidjourneyJobType}; +use midjourney_client::utils::get_image_url::get_image_url; +use midjourney_client::utils::image_downloader_client::ImageDownloaderClient; +use once_cell::sync::Lazy; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs}; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::endpoints::media_files::get_media_file::get_media_file; +use storyteller_client::endpoints::media_files::legacy_upload_media_file_from_file::{legacy_upload_media_file_from_file, LegacyUploadMediaFileFromFileArgs}; +use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs}; +use storyteller_client::endpoints::media_files::upload_video_media_file_from_file::{upload_video_media_file_from_file, UploadVideoFromFileArgs}; +use storyteller_client::endpoints::prompts::create_prompt::create_prompt; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use storyteller_client::recipes::upload_media_file_from_file::upload_media_file_from_file; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; +use url::Url; +use url_utils::extension::extract_download_extension_from_url::extract_download_extension_from_url_str; +use world_labs_client::api::api_types::world_id::WorldObjectId; +use world_labs_client::api::requests::worlds::poll_world_status::{poll_world_status, PollWorldStatusArgs}; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; +use world_labs_client::credentials::worldlabs_refresh_token::WorldLabsRefreshToken; + +pub async fn worldlabs_marble_task_polling( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + creds: WorldlabsCredentialManager, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &creds, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + worldlabs_creds: &WorldlabsCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + loop { + if !worldlabs_creds.do_task_polling()? { + tokio::time::sleep(Duration::from_millis(10_000)).await; + continue; + } + + // TODO: Graceful wait, fix this long function body + let storyteller_creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + error!("No Storyteller credentials found. Cannot proceed with WorldLabs polling."); + tokio::time::sleep(Duration::from_millis(5_000)).await; + continue; + } + }; + + let world_labs_cookies = match worldlabs_creds.maybe_copy_typed_cookies()? { + Some(cookies) => cookies, + None => { + info!("No full WorldLabs cookies"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let world_labs_bearer = match worldlabs_creds.maybe_copy_bearer_token()? { + Some(bearer) => bearer, + None => { + info!("No full WorldLabs bearer"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let world_labs_refresh = match worldlabs_creds.maybe_copy_refresh_token()? { + Some(bearer) => bearer, + None => { + info!("No full WorldLabs refresh"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { + db: task_database.get_connection(), + provider: GenerationProvider::WorldLabs, + task_statuses: &TASK_DATABASE_PENDING_STATUSES, + }).await?; + + poll_grok_tasks( + app_handle, + app_env_configs, + app_data_root, + task_database, + &world_labs_cookies, + &world_labs_bearer, + &world_labs_refresh, + &storyteller_creds, + local_tasks, + ).await?; + + tokio::time::sleep(Duration::from_millis(2_000)).await; + } +} + +async fn poll_grok_tasks( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + world_labs_cookies: &WorldLabsCookies, + world_labs_bearer: &WorldLabsBearerToken, + world_labs_refresh: &WorldLabsRefreshToken, + storyteller_creds: &StorytellerCredentialSet, + local_tasks: TaskList, +) -> AnyhowResult<()> { + let local_tasks = local_tasks.tasks; + + if local_tasks.is_empty() { + return Ok(()) + } + + info!("WorldLabs tasks waiting: {:?}", local_tasks.len()); + + // Map of WorldLabs World ID to Local Task. + let local_tasks_by_world_labs_world_id = local_tasks.iter() + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + .collect::>(); + + + for (world_id, local_task) in local_tasks_by_world_labs_world_id.iter() { + let world_id = WorldObjectId(world_id.to_string()); + + let poll_world_response = poll_world_status(PollWorldStatusArgs { + cookies: &world_labs_cookies, + bearer_token: &world_labs_bearer, + world_id: &world_id, + request_timeout: None, + }).await?; + + if !poll_world_response.is_complete { + tokio::time::sleep(Duration::from_millis(2_000)).await; + continue; + } + + let spz_url = match poll_world_response.spz_splat_url { + Some(url) => url, + None => { + error!("No spz splat URL despite being marked complete"); + continue; + } + }; + + + upload_spz_splat( + &app_handle, + &app_env_configs, + app_data_root, + task_database, + &storyteller_creds, + &local_task, + &spz_url, + ).await?; + } + + + tokio::time::sleep(Duration::from_millis(5_000)).await; + + Ok(()) +} + +async fn upload_spz_splat( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds: &StorytellerCredentialSet, + local_task: &Task, + spz_url: &str, +) -> AnyhowResult<()> { + + let mut maybe_primary_media_file_token = None; + + info!("Downloading generated spz splat ..."); + + let spz_download_filename = download_spz(spz_url, app_data_root).await?; + + let mut wait_delay = 0; + + loop { + info!("Uploading to backend..."); + + // TODO: media_files.origin_category + // TODO: media_files.maybe_prompt_token + // TODO: media_files.maybe_generation_provider + // TODO: media_files.maybe_origin_model_type + // TODO: media_files.maybe_origin_model_token (sref?) + // TODO: media_files.is_batch_generated + // TODO: media_files.maybe_batch_token + // TODO: media_files.is_user_upload + + // TODO: batch_generations.token + // TODO: batch_generations.entity_type + // TODO: batch_generations.entity_token + + let result = legacy_upload_media_file_from_file(LegacyUploadMediaFileFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: &spz_download_filename, + }).await; + + match result { + Ok(result) => { + info!("Successfully uploaded to backend: {:?}", result.media_file_token); + if maybe_primary_media_file_token.is_none() { + maybe_primary_media_file_token = Some(result.media_file_token); + } + break; + }, + Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => { + error!("Too many requests, retrying upload after delay..."); + // If we hit a rate limit, we can retry after a short delay. + wait_delay += 10; + if wait_delay > 60 { + wait_delay = 60; + } + tokio::time::sleep(Duration::from_secs(wait_delay)).await; + continue; // Retry the upload. + } + Err(err) => { + error!("Failed to upload to backend: {:?}", err); + return Err(err.into()) + }, + } + } // End loop + + let mut maybe_cdn_url = None; + let mut maybe_thumbnail_url_template = None; + + if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() { + info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token); + + let lookup_result = get_media_file( + &app_env_configs.storyteller_host, + media_file_token, + ).await; + match lookup_result { + Ok(response) => { + maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string()); + maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links) + .map(|s| s.to_string()); + } + Err(err) => { + error!("Failed to look up media file after upload: {:?} (failing open)", err); + } + } + } + + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { + db: task_database.get_connection(), + task_id: &local_task.id, + maybe_batch_token: None, + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_primary_media_file_class: Some(TaskMediaFileClass::Dimensional), + maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), + maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(), + }).await?; + + if !updated { + return Ok(()); // If anything breaks with queries, don't spam events. + } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateGaussian), + service: GenerationServiceProvider::WorldLabs, + model: None, + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + + //let result = maybe_handle_text_to_image_complete_event( + // app_handle, + // app_env_configs, + // Some(storyteller_creds), + // local_task, + // &batch_token, + //).await; + + //if let Err(err) = result { + // error!("Failed to send text-to-image complete event: {:?}", err); + //} + + Ok(()) +} + +async fn download_spz( + spz_url: &str, + app_data_root: &AppDataRoot, +) -> AnyhowResult { + info!("Downloading splat from URL: {}", spz_url); + + let response = reqwest::get(spz_url).await?; + let image_bytes = response.bytes().await?; + + let uuid = generate_random_uuid(); + let extension_without_period = "ceramic.spz"; + + let tempdir = app_data_root.temp_dir().path(); + let download_filename = format!("{}.{}", uuid, extension_without_period); + let download_path = tempdir.join(download_filename); + + info!("Writing to path: {:?}", download_path); + + let mut file = File::create(&download_path)?; + file.write_all(&image_bytes)?; + + Ok(download_path) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/mod.rs new file mode 100644 index 0000000000..b4a0713003 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/mod.rs @@ -0,0 +1 @@ +pub mod worldlabs_login_window; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/mod.rs new file mode 100644 index 0000000000..b1a1c8bbc7 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/mod.rs @@ -0,0 +1,4 @@ +pub mod worldlabs_login_window_open; +pub (super) mod worldlabs_login_window_thread; +pub (super) mod worldlabs_login_webview_extract_cookies; +pub (super) mod worldlabs_javascript; diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_javascript.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_javascript.rs new file mode 100644 index 0000000000..097520e099 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_javascript.rs @@ -0,0 +1,91 @@ + + + +pub (super) const WORLDLABS_JAVASCRIPT_EXPORT_BEARER_TOKENS : &str = r#" + (async () => { + if (window.tokensExported) { + return; + } + + const request = indexedDB.open("firebaseLocalStorageDb"); + + request.onerror = (event) => {}; + + request.onsuccess = (event) => { + const db = event.target.result; + + const tx = db.transaction("firebaseLocalStorage", "readonly"); + const store = tx.objectStore("firebaseLocalStorage"); + + const keysRequest = store.getAllKeys(); + + keysRequest.onerror = (event) => {}; + + keysRequest.onsuccess = () => { + const keys = keysRequest.result; // array of keys + if (keys.length < 1) { + return; + } + + const key = keys[0]; + const getKeyRequest = store.getKey(key); + + getKeyRequest.onsuccess = () => {}; + + store.openCursor().onsuccess = async (event) => { + const cursor = event.target.result; + + if (cursor) { + let tokens = cursor.value?.value?.stsTokenManager; + + if (tokens?.accessToken && tokens?.refreshToken) { + // Send to Tauri + let result = await window.__TAURI__.core.invoke("worldlabs_receive_bearer_command", { + request: { + bearer_token: tokens.accessToken, + refresh_token: tokens.refreshToken, + } + }); + window.tokensExported = true; + } + cursor.continue(); + } + }; // end cursor + + }; // end keyRequest + }; // end DB + })(); + "#; + + +/// Take the Bearer tokens and send them to Tauri +pub (super) const WORLDLABS_JAVASCRIPT_SEND_BEARER_TO_TAURI : &str = r#" + (async () => { + if (!(window.tokens?.accessToken && window.tokens?.refreshToken)) { + return; + } + + let result = await window.__TAURI__.core.invoke("worldlabs_receive_bearer_command", { + request: { + bearer_token: window.tokens.accessToken, + refresh_token: window.tokens.refreshToken, + } + }); + + error.log('>>> result', result); + + })(); + "#; + + +/// See if the pricing button exists. +/// This is only for the logged-in state. +/// We can use redirection to detect completion. +pub (super) const WORLDLABS_JAVASCRIPT_SEE_IF_PRICING_EXISTS : &str = r#" + (() => { + let pricing = document.querySelectorAll("a[href='/pricing']"); + if (pricing.length > 0) { + window.location.href = "https://www.worldlabs.ai/about"; + } + })(); + "#; diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs new file mode 100644 index 0000000000..cf2e62f475 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_webview_extract_cookies.rs @@ -0,0 +1,39 @@ +use cookie_store::cookie_store::CookieStore; +use errors::AnyhowResult; +use once_cell::sync::Lazy; +use reqwest::Url; +use tauri::webview::Cookie; +use tauri::WebviewWindow; + +static ROOT_COOKIE_URL: Lazy = Lazy::new(|| { + Url::parse("https://worldlabs.ai").expect("URL should parse") +}); + +pub fn worldlabs_login_webview_extract_cookies(webview: &WebviewWindow) -> AnyhowResult { + let mut cookie_store = CookieStore::empty(); + let cookies = get_all_worldlabs_cookies(webview)?; + for cookie in cookies.iter() { + cookie_store.add_cookie_name_and_value( + cookie.name().to_string(), + cookie.value().to_string(), + ); + } + Ok(cookie_store) +} + +fn get_all_worldlabs_cookies(webview: &WebviewWindow) -> AnyhowResult> { + //let www_cookies = webview.cookies_for_url(WWW_COOKIE_URL.clone())?; + let root_cookies = webview.cookies_for_url(ROOT_COOKIE_URL.clone())?; + + let mut all_cookies = root_cookies; + //let mut cookie_names = HashSet::new(); + + //for cookie in root_cookies.iter() { + // if !cookie_names.contains(cookie.name()) { + // cookie_names.insert(cookie.name().to_string()); + // all_cookies.push(cookie.clone()); + // } + //} + + Ok(all_cookies) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_open.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_open.rs new file mode 100644 index 0000000000..6fb5a3e208 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_open.rs @@ -0,0 +1,108 @@ +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::utils::clear_all_webview_cookies::clear_all_webview_cookies; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use crate::services::worldlabs::windows::worldlabs_login_window::worldlabs_login_window_thread::worldlabs_login_window_thread; +use anyhow::anyhow; +use errors::AnyhowResult; +use log::info; +use once_cell::sync::Lazy; +use reqwest::Url; +use std::time::Duration; +use tauri::webview::NewWindowResponse; +use tauri::{AppHandle, Manager, WebviewUrl, WebviewWindowBuilder}; + +/// Name of the window +pub (super) const WORLDLABS_LOGIN_WINDOW_NAME: &str = "worldlabs_login_window"; + +pub (super) static START_URL: Lazy = Lazy::new(|| { + Url::parse("https://google.com").expect("URL should parse") +}); + +pub (super) static WORLDLABS_HOMEPAGE_URL: Lazy = Lazy::new(|| { + Url::parse("https://marble.worldlabs.ai/").expect("URL should parse") +}); + +pub (super) static WORLDLABS_LOGIN_URL: Lazy = Lazy::new(|| { + Url::parse("https://marble.worldlabs.ai/").expect("URL should parse") +}); + +pub async fn worldlabs_login_window_open( + app: &AppHandle, + app_data_root: &AppDataRoot, + worldlabs_bearer_bridge: &WorldlabsBearerBridge, + worldlabs_creds_manager: &WorldlabsCredentialManager, +) -> AnyhowResult<()> { + + if app.get_window(WORLDLABS_LOGIN_WINDOW_NAME).is_some() { + return Err(anyhow!("Login window already open")); + } + + info!("Clearing bearer bridge..."); + + worldlabs_bearer_bridge.clear()?; + + let url = WebviewUrl::External(START_URL.clone()); + + let window = WebviewWindowBuilder::new(app, WORLDLABS_LOGIN_WINDOW_NAME, url) + //.user_agent(openai_sora_client::credentials::USER_AGENT) + .on_new_window(move |url, features| { + // WorldLabs needs popups as they (1) don't have a distinct login page + // (it's embedded) and (2) they don't redirect with OAuth, but rather + // open a popup to Google. It still works, but it's annoying. + NewWindowResponse::Allow + }) + .always_on_top(false) + .title("Login to WorldLabs") + .center() + .resizable(true) + .visible(true) + .closable(true) + .min_inner_size(800.0, 800.0) + .focused(true) + .devtools(true) + .build()?; + + let webview = window.get_webview(WORLDLABS_LOGIN_WINDOW_NAME) + .ok_or_else(|| anyhow!("no webview found"))?; + + clear_all_webview_cookies(&webview)?; + + webview.navigate(WORLDLABS_HOMEPAGE_URL.clone())?; + + tokio::time::sleep(Duration::from_millis(200)).await; + + info!("Running script..."); + + for _ in 0 .. 10 { + tokio::time::sleep(Duration::from_millis(100)).await; + + // Close the "welcome" modal + webview.eval(r#"document.querySelectorAll("div:has(h2) ~ button[data-slot=dialog-close]").forEach((el) => el.click())"#)?; + } + + for _ in 0 .. 10 { + tokio::time::sleep(Duration::from_millis(100)).await; + + // Open the "login" modal + webview.eval(r#"document.querySelectorAll("button:has(span[data-slot=avatar])").forEach((el) => el.click())"#)?; + } + + info!("Script done..."); + + let app_handle = app.clone(); + let app_data_root = app_data_root.clone(); + let worldlabs_creds_manager = worldlabs_creds_manager.clone(); + let worldlabs_bearer_bridge = worldlabs_bearer_bridge.clone(); + + let _ = tauri::async_runtime::spawn(async move { + worldlabs_login_window_thread( + app_handle, + app_data_root, + worldlabs_bearer_bridge, + worldlabs_creds_manager + ).await; + }); + + Ok(()) +} diff --git a/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_thread.rs b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_thread.rs new file mode 100644 index 0000000000..b8257e2060 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/windows/worldlabs_login_window/worldlabs_login_window_thread.rs @@ -0,0 +1,129 @@ +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::functional_events::refresh_account_state_event::RefreshAccountStateEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::events::sendable_event_trait::SendableEvent; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::utils::window::get_webview_window_hostname::get_webview_window_hostname; +use crate::core::utils::window::get_webview_window_url_path::get_webview_window_url_path; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::sora::events::sora_login_success_event::SoraLoginSuccessEvent; +use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; +use crate::services::sora::windows::sora_login_window::extract_sora_webview_cookies::extract_sora_webview_cookies; +use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use crate::services::worldlabs::windows::worldlabs_login_window::worldlabs_javascript::WORLDLABS_JAVASCRIPT_EXPORT_BEARER_TOKENS; +use crate::services::worldlabs::windows::worldlabs_login_window::worldlabs_login_webview_extract_cookies::worldlabs_login_webview_extract_cookies; +use crate::services::worldlabs::windows::worldlabs_login_window::worldlabs_login_window_open::WORLDLABS_LOGIN_WINDOW_NAME; +use anyhow::anyhow; +use cookie_store::cookie_store::CookieStore; +use enums::common::generation_provider::GenerationProvider; +use errors::AnyhowResult; +use log::{error, info}; +use openai_sora_client::creds::sora_credential_set::SoraCredentialSet; +use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; +use openai_sora_client::utils::has_session_cookie::{has_session_cookie, SessionCookiePresence}; +use tauri::{AppHandle, Manager, WebviewWindow}; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::worldlabs_refresh_token::WorldLabsRefreshToken; + +pub async fn worldlabs_login_window_thread( + app: AppHandle, + app_data_root: AppDataRoot, + worldlabs_bearer_bridge: WorldlabsBearerBridge, + worldlabs_creds_manager: WorldlabsCredentialManager, +) { + let mut visited_login = false; + + loop { + let login_webview_window = match app.get_webview_window(WORLDLABS_LOGIN_WINDOW_NAME) { + Some(webview) => webview, + None => { + info!("Exit WorldLabs login thread."); + return; // NB: Only exit if we don't have the webview. + } + }; + + let result = check_login_window( + &app, + &login_webview_window, + &app_data_root, + &worldlabs_bearer_bridge, + &worldlabs_creds_manager, + &mut visited_login, + ).await; + + match result { + Err(err) => { + error!("Error checking Grok login window: {:?}", err); + } + Ok(false) => {} // Continue iteration and try again... + Ok(true) => { + info!("Successfully saved cookies from login window. Closing."); + if let Err(err) = login_webview_window.close() { + error!("Error closing login window: {:?}", err); + } + return; + } + } + + tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; + } +} + +/// NOTE: THERE IS THIS - https://marble.worldlabs.ai/sign-in?redirect=%2Fworlds + +/// Returns true if we can exit. +async fn check_login_window( + app_handle: &AppHandle, + webview_window: &WebviewWindow, + app_data_root: &AppDataRoot, + worldlabs_bearer_bridge: &WorldlabsBearerBridge, + worldlabs_creds_manager: &WorldlabsCredentialManager, + visited_login: &mut bool, +) -> AnyhowResult { + + // World labs has no distinct login page. Everything is in a single SPA. + + webview_window.eval(WORLDLABS_JAVASCRIPT_EXPORT_BEARER_TOKENS)?; + + let maybe_bearer = worldlabs_bearer_bridge.get()?; + + let bearer_info = match maybe_bearer { + Some(bearer) => bearer, + None => { + return Ok(false); // Not logged in. + } + }; + + let cookie_store = worldlabs_login_webview_extract_cookies(webview_window)?; + + info!("Current cookies (len {}): {:?}", cookie_store.len(), cookie_store.to_cookie_string()); + + worldlabs_creds_manager.replace_cookie_store(cookie_store)?; + + let bearer_token = WorldLabsBearerToken::new(bearer_info.bearer_token); + + worldlabs_creds_manager.replace_bearer_token(bearer_token)?; + + let refresh_token = WorldLabsRefreshToken::new(bearer_info.refresh_token); + + worldlabs_creds_manager.replace_refresh_token(refresh_token)?; + + let result = worldlabs_creds_manager.persist_to_disk(); + + if let Err(err) = result { + error!("Error persisting WorldLabs credentials to disk: {:?}", err); + return Ok(false); + } + + let event = RefreshAccountStateEvent { + provider: Some(GenerationProvider::WorldLabs), + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send RefreshAccountStateEvent: {:?}", err); // Fail open + } + + Ok(true) +} diff --git a/crates/desktop/artcraft/tauri-mac.conf.json b/crates/desktop/artcraft/tauri-mac.conf.json index 78a11a761d..83bbff221c 100644 --- a/crates/desktop/artcraft/tauri-mac.conf.json +++ b/crates/desktop/artcraft/tauri-mac.conf.json @@ -10,6 +10,7 @@ "beforeBuildCommand": "nx run editor3d:build" }, "app": { + "withGlobalTauri": true, "windows": [ ], "security": { diff --git a/crates/desktop/artcraft/tauri.artcraft_3d.no_dev.conf.json b/crates/desktop/artcraft/tauri.artcraft_3d.no_dev.conf.json index 07dbef7f3c..39f3c8b631 100644 --- a/crates/desktop/artcraft/tauri.artcraft_3d.no_dev.conf.json +++ b/crates/desktop/artcraft/tauri.artcraft_3d.no_dev.conf.json @@ -6,6 +6,7 @@ "beforeBuildCommand": "npx nx run editor3d:build" }, "app": { + "withGlobalTauri": true, "windows": [ ] } diff --git a/crates/desktop/artcraft/tauri.conf.json b/crates/desktop/artcraft/tauri.conf.json index 78a11a761d..83bbff221c 100644 --- a/crates/desktop/artcraft/tauri.conf.json +++ b/crates/desktop/artcraft/tauri.conf.json @@ -10,6 +10,7 @@ "beforeBuildCommand": "nx run editor3d:build" }, "app": { + "withGlobalTauri": true, "windows": [ ], "security": { diff --git a/crates/schema/public/enums/src/common/generation_provider.rs b/crates/schema/public/enums/src/common/generation_provider.rs index e365fd603a..ff3f40bf43 100644 --- a/crates/schema/public/enums/src/common/generation_provider.rs +++ b/crates/schema/public/enums/src/common/generation_provider.rs @@ -18,6 +18,7 @@ pub enum GenerationProvider { Grok, Midjourney, Sora, + WorldLabs, } impl_enum_display_and_debug_using_to_str!(GenerationProvider); @@ -37,6 +38,7 @@ impl GenerationProvider { Self::Grok => "grok", Self::Midjourney => "midjourney", Self::Sora => "sora", + Self::WorldLabs => "world_labs", } } @@ -47,6 +49,7 @@ impl GenerationProvider { "grok" => Ok(Self::Grok), "midjourney" => Ok(Self::Midjourney), "sora" => Ok(Self::Sora), + "world_labs" => Ok(Self::WorldLabs), _ => Err(EnumError::CouldNotConvertFromString(value.to_string())), } } @@ -60,6 +63,7 @@ impl GenerationProvider { Self::Grok, Self::Midjourney, Self::Sora, + Self::WorldLabs, ]) } } @@ -80,6 +84,7 @@ mod tests { assert_serialization(GenerationProvider::Grok, "grok"); assert_serialization(GenerationProvider::Midjourney, "midjourney"); assert_serialization(GenerationProvider::Sora, "sora"); + assert_serialization(GenerationProvider::WorldLabs, "world_labs"); } #[test] @@ -89,6 +94,7 @@ mod tests { assert_eq!(GenerationProvider::Grok.to_str(), "grok"); assert_eq!(GenerationProvider::Midjourney.to_str(), "midjourney"); assert_eq!(GenerationProvider::Sora.to_str(), "sora"); + assert_eq!(GenerationProvider::WorldLabs.to_str(), "world_labs"); } #[test] @@ -98,6 +104,7 @@ mod tests { assert_eq!(GenerationProvider::from_str("grok").unwrap(), GenerationProvider::Grok); assert_eq!(GenerationProvider::from_str("midjourney").unwrap(), GenerationProvider::Midjourney); assert_eq!(GenerationProvider::from_str("sora").unwrap(), GenerationProvider::Sora); + assert_eq!(GenerationProvider::from_str("world_labs").unwrap(), GenerationProvider::WorldLabs); } #[test] @@ -114,12 +121,13 @@ mod tests { #[test] fn all_variants() { let mut variants = GenerationProvider::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(GenerationProvider::Artcraft)); assert_eq!(variants.pop_first(), Some(GenerationProvider::Fal)); assert_eq!(variants.pop_first(), Some(GenerationProvider::Grok)); assert_eq!(variants.pop_first(), Some(GenerationProvider::Midjourney)); assert_eq!(variants.pop_first(), Some(GenerationProvider::Sora)); + assert_eq!(variants.pop_first(), Some(GenerationProvider::WorldLabs)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs index b26825be94..f65965b240 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs +++ b/crates/schema/public/enums/src/tauri/tasks/task_model_type.rs @@ -69,6 +69,8 @@ pub enum TaskModelType { Hunyuan3d2_0, #[serde(rename = "hunyuan_3d_2.1")] Hunyuan3d2_1, + #[serde(rename = "worldlabs_marble")] + WorldlabsMarble, } impl_enum_display_and_debug_using_to_str!(TaskModelType); @@ -108,6 +110,7 @@ impl TaskModelType { // 3D Object generation models Self::Hunyuan3d2_0 => "hunyuan_3d_2.0", Self::Hunyuan3d2_1 => "hunyuan_3d_2.1", + Self::WorldlabsMarble => "worldlabs_marble", } } @@ -141,6 +144,7 @@ impl TaskModelType { // 3D Object generation models "hunyuan_3d_2.0" => Ok(Self::Hunyuan3d2_0), "hunyuan_3d_2.1" => Ok(Self::Hunyuan3d2_1), + "worldlabs_marble" => Ok(Self::WorldlabsMarble), _ => Err(EnumError::CouldNotConvertFromString(value.to_string())), } } @@ -177,6 +181,7 @@ impl TaskModelType { // 3D Object generation models Self::Hunyuan3d2_0, Self::Hunyuan3d2_1, + Self::WorldlabsMarble, ]) } } @@ -220,6 +225,7 @@ mod tests { // 3D Object generation models assert_serialization(TaskModelType::Hunyuan3d2_0, "hunyuan_3d_2.0"); assert_serialization(TaskModelType::Hunyuan3d2_1, "hunyuan_3d_2.1"); + assert_serialization(TaskModelType::WorldlabsMarble, "worldlabs_marble"); } #[test] @@ -252,6 +258,7 @@ mod tests { // 3D Object generation models assert_eq!(TaskModelType::Hunyuan3d2_0.to_str(), "hunyuan_3d_2.0"); assert_eq!(TaskModelType::Hunyuan3d2_1.to_str(), "hunyuan_3d_2.1"); + assert_eq!(TaskModelType::WorldlabsMarble.to_str(), "worldlabs_marble"); } #[test] @@ -284,6 +291,7 @@ mod tests { // 3D Object generation models assert_eq!(TaskModelType::from_str("hunyuan_3d_2.0").unwrap(), TaskModelType::Hunyuan3d2_0); assert_eq!(TaskModelType::from_str("hunyuan_3d_2.1").unwrap(), TaskModelType::Hunyuan3d2_1); + assert_eq!(TaskModelType::from_str("worldlabs_marble").unwrap(), TaskModelType::WorldlabsMarble); } #[test] @@ -300,7 +308,7 @@ mod tests { #[test] fn all_variants() { let mut variants = TaskModelType::all_variants(); - assert_eq!(variants.len(), 25); + assert_eq!(variants.len(), 26); // Image models assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Dev)); assert_eq!(variants.pop_first(), Some(TaskModelType::Flux1Schnell)); @@ -329,6 +337,7 @@ mod tests { // 3D Object generation models assert_eq!(variants.pop_first(), Some(TaskModelType::Hunyuan3d2_0)); assert_eq!(variants.pop_first(), Some(TaskModelType::Hunyuan3d2_1)); + assert_eq!(variants.pop_first(), Some(TaskModelType::WorldlabsMarble)); assert_eq!(variants.pop_first(), None); } } diff --git a/crates/schema/public/enums/src/tauri/tasks/task_type.rs b/crates/schema/public/enums/src/tauri/tasks/task_type.rs index 2b2d4edcf6..d049fd27db 100644 --- a/crates/schema/public/enums/src/tauri/tasks/task_type.rs +++ b/crates/schema/public/enums/src/tauri/tasks/task_type.rs @@ -15,6 +15,7 @@ pub enum TaskType { ImageInpaintEdit, VideoGeneration, ObjectGeneration, + GaussianGeneration, BackgroundRemoval, } @@ -31,6 +32,7 @@ impl TaskType { Self::ImageInpaintEdit => "image_inpaint_edit", Self::VideoGeneration => "video_generation", Self::ObjectGeneration => "object_generation", + Self::GaussianGeneration => "gaussian_generation", Self::BackgroundRemoval => "background_removal", } } @@ -41,6 +43,7 @@ impl TaskType { "image_inpaint_edit" => Ok(Self::ImageInpaintEdit), "video_generation" => Ok(Self::VideoGeneration), "object_generation" => Ok(Self::ObjectGeneration), + "gaussian_generation" => Ok(Self::GaussianGeneration), "background_removal" => Ok(Self::BackgroundRemoval), _ => Err(EnumError::CouldNotConvertFromString(value.to_string())), } @@ -54,6 +57,7 @@ impl TaskType { Self::ImageInpaintEdit, Self::VideoGeneration, Self::ObjectGeneration, + Self::GaussianGeneration, Self::BackgroundRemoval, ]) } @@ -74,6 +78,7 @@ mod tests { assert_serialization(TaskType::ImageInpaintEdit, "image_inpaint_edit"); assert_serialization(TaskType::VideoGeneration, "video_generation"); assert_serialization(TaskType::ObjectGeneration, "object_generation"); + assert_serialization(TaskType::GaussianGeneration, "gaussian_generation"); assert_serialization(TaskType::BackgroundRemoval, "background_removal"); } @@ -83,6 +88,7 @@ mod tests { assert_eq!(TaskType::ImageInpaintEdit.to_str(), "image_inpaint_edit"); assert_eq!(TaskType::VideoGeneration.to_str(), "video_generation"); assert_eq!(TaskType::ObjectGeneration.to_str(), "object_generation"); + assert_eq!(TaskType::GaussianGeneration.to_str(), "gaussian_generation"); assert_eq!(TaskType::BackgroundRemoval.to_str(), "background_removal"); } @@ -92,6 +98,7 @@ mod tests { assert_eq!(TaskType::from_str("image_inpaint_edit").unwrap(), TaskType::ImageInpaintEdit); assert_eq!(TaskType::from_str("video_generation").unwrap(), TaskType::VideoGeneration); assert_eq!(TaskType::from_str("object_generation").unwrap(), TaskType::ObjectGeneration); + assert_eq!(TaskType::from_str("gaussian_generation").unwrap(), TaskType::GaussianGeneration); assert_eq!(TaskType::from_str("background_removal").unwrap(), TaskType::BackgroundRemoval); } @@ -109,11 +116,12 @@ mod tests { #[test] fn all_variants() { let mut variants = TaskType::all_variants(); - assert_eq!(variants.len(), 5); + assert_eq!(variants.len(), 6); assert_eq!(variants.pop_first(), Some(TaskType::ImageGeneration)); assert_eq!(variants.pop_first(), Some(TaskType::ImageInpaintEdit)); assert_eq!(variants.pop_first(), Some(TaskType::VideoGeneration)); assert_eq!(variants.pop_first(), Some(TaskType::ObjectGeneration)); + assert_eq!(variants.pop_first(), Some(TaskType::GaussianGeneration)); assert_eq!(variants.pop_first(), Some(TaskType::BackgroundRemoval)); assert_eq!(variants.pop_first(), None); } diff --git a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs index 3c61656bfb..6a9066e395 100644 --- a/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs +++ b/crates/schema/public/enums/src/tauri/ux/tauri_event_name.rs @@ -67,6 +67,11 @@ pub enum TauriEventName { /// Object (3D mesh) generation is complete #[serde(rename = "object_generation_complete_event")] ObjectGenerationCompleteEvent, + + /// Special event: + /// Gaussian generation is complete + #[serde(rename = "gaussian_generation_complete_event")] + GaussianGenerationCompleteEvent, /// Special event: /// Video generation is complete @@ -113,6 +118,7 @@ impl TauriEventName { Self::TextToImageGenerationCompleteEvent => "text_to_image_generation_complete_event", Self::ImageEditCompleteEvent => "image_edit_complete_event", Self::ObjectGenerationCompleteEvent => "object_generation_complete_event", + Self::GaussianGenerationCompleteEvent => "gaussian_generation_complete_event", Self::VideoGenerationCompleteEvent => "video_generation_complete_event", Self::RefreshAccountStateEvent => "refresh_account_state_event", Self::ShowProviderBillingModalEvent => "show_provider_billing_modal_event", @@ -133,6 +139,7 @@ impl TauriEventName { "text_to_image_generation_complete_event" => Ok(Self::TextToImageGenerationCompleteEvent), "image_edit_complete_event" => Ok(Self::ImageEditCompleteEvent), "object_generation_complete_event" => Ok(Self::ObjectGenerationCompleteEvent), + "gaussian_generation_complete_event" => Ok(Self::GaussianGenerationCompleteEvent), "video_generation_complete_event" => Ok(Self::VideoGenerationCompleteEvent), "refresh_account_state_event" => Ok(Self::RefreshAccountStateEvent), "show_provider_billing_modal_event" => Ok(Self::ShowProviderBillingModalEvent), @@ -156,6 +163,7 @@ impl TauriEventName { Self::TextToImageGenerationCompleteEvent, Self::ImageEditCompleteEvent, Self::ObjectGenerationCompleteEvent, + Self::GaussianGenerationCompleteEvent, Self::VideoGenerationCompleteEvent, Self::RefreshAccountStateEvent, Self::ShowProviderBillingModalEvent, @@ -185,6 +193,7 @@ mod tests { assert_serialization(TauriEventName::TextToImageGenerationCompleteEvent, "text_to_image_generation_complete_event"); assert_serialization(TauriEventName::ImageEditCompleteEvent, "image_edit_complete_event"); assert_serialization(TauriEventName::ObjectGenerationCompleteEvent, "object_generation_complete_event"); + assert_serialization(TauriEventName::GaussianGenerationCompleteEvent, "gaussian_generation_complete_event"); assert_serialization(TauriEventName::VideoGenerationCompleteEvent, "video_generation_complete_event"); assert_serialization(TauriEventName::RefreshAccountStateEvent, "refresh_account_state_event"); assert_serialization(TauriEventName::ShowProviderBillingModalEvent, "show_provider_billing_modal_event"); @@ -204,6 +213,7 @@ mod tests { assert_eq!(TauriEventName::TextToImageGenerationCompleteEvent.to_str(), "text_to_image_generation_complete_event"); assert_eq!(TauriEventName::ImageEditCompleteEvent.to_str(), "image_edit_complete_event"); assert_eq!(TauriEventName::ObjectGenerationCompleteEvent.to_str(), "object_generation_complete_event"); + assert_eq!(TauriEventName::GaussianGenerationCompleteEvent.to_str(), "gaussian_generation_complete_event"); assert_eq!(TauriEventName::VideoGenerationCompleteEvent.to_str(), "video_generation_complete_event"); assert_eq!(TauriEventName::RefreshAccountStateEvent.to_str(), "refresh_account_state_event"); assert_eq!(TauriEventName::ShowProviderBillingModalEvent.to_str(), "show_provider_billing_modal_event"); @@ -223,6 +233,7 @@ mod tests { assert_eq!(TauriEventName::from_str("text_to_image_generation_complete_event").unwrap(), TauriEventName::TextToImageGenerationCompleteEvent); assert_eq!(TauriEventName::from_str("image_edit_complete_event").unwrap(), TauriEventName::ImageEditCompleteEvent); assert_eq!(TauriEventName::from_str("object_generation_complete_event").unwrap(), TauriEventName::ObjectGenerationCompleteEvent); + assert_eq!(TauriEventName::from_str("gaussian_generation_complete_event").unwrap(), TauriEventName::GaussianGenerationCompleteEvent); assert_eq!(TauriEventName::from_str("video_generation_complete_event").unwrap(), TauriEventName::VideoGenerationCompleteEvent); assert_eq!(TauriEventName::from_str("refresh_account_state_event").unwrap(), TauriEventName::RefreshAccountStateEvent); assert_eq!(TauriEventName::from_str("show_provider_billing_modal_event").unwrap(), TauriEventName::ShowProviderBillingModalEvent); @@ -233,7 +244,7 @@ mod tests { #[test] fn all_variants() { let mut variants = TauriEventName::all_variants(); - assert_eq!(variants.len(), 15); + assert_eq!(variants.len(), 16); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueSuccessEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationEnqueueFailureEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::GenerationCompleteEvent)); @@ -244,6 +255,7 @@ mod tests { assert_eq!(variants.pop_first(), Some(TauriEventName::TextToImageGenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ImageEditCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ObjectGenerationCompleteEvent)); + assert_eq!(variants.pop_first(), Some(TauriEventName::GaussianGenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::VideoGenerationCompleteEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::RefreshAccountStateEvent)); assert_eq!(variants.pop_first(), Some(TauriEventName::ShowProviderBillingModalEvent)); diff --git a/frontend/apps/editor3d/app/index.html b/frontend/apps/editor3d/app/index.html index f921d41be9..c93313f28a 100644 --- a/frontend/apps/editor3d/app/index.html +++ b/frontend/apps/editor3d/app/index.html @@ -14,7 +14,9 @@ gtag('js', new Date()); gtag('config', 'G-Y4Y8KMLPQ7'); - + + +