mirror of
https://github.com/storytold/artcraft.git
synced 2026-10-20 17:02:53 -04:00
worldlabs tauri (#933)
not finished: need refresh tokens and thumbnails
This commit is contained in:
Generated
+2
@@ -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",
|
||||
|
||||
@@ -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
|
||||
|
||||
+143
@@ -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<Path>> {
|
||||
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<GenerationProvider>,
|
||||
|
||||
// /// 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<P: AsRef<Path>>(
|
||||
args: LegacyUploadMediaFileFromFileArgs<'_, P>,
|
||||
) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
|
||||
|
||||
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<P: AsRef<Path>>(
|
||||
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(())
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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
|
||||
|
||||
+1
-1
@@ -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)
|
||||
|
||||
@@ -54,7 +54,7 @@ pub async fn create_run_object(args: CreateRunObjectArgs<'_>) -> Result<CreateRu
|
||||
.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)
|
||||
|
||||
+1
-1
@@ -55,7 +55,7 @@ pub async fn create_upload_object(args: CreateUploadObjectArgs<'_>) -> Result<Cr
|
||||
.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)
|
||||
|
||||
+1
-1
@@ -61,7 +61,7 @@ pub async fn finalize_object_image_upload(args: FinalizeObjectImageUploadArgs<'_
|
||||
.header(ACCEPT_LANGUAGE, "en-US,en;q=0.5")
|
||||
.header(ACCEPT_ENCODING, "gzip, deflate, br, zstd")
|
||||
.header(REFERER, REFERER_VALUE)
|
||||
.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)
|
||||
|
||||
+1
-1
@@ -84,7 +84,7 @@ pub async fn update_run_object_with_upload(args: UpdateRunObjectWithUploadArgs<'
|
||||
.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)
|
||||
|
||||
+1
-1
@@ -91,7 +91,7 @@ pub async fn update_run_object_with_world(args: UpdateRunObjectWithWorldArgs<'_>
|
||||
.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)
|
||||
|
||||
@@ -58,7 +58,7 @@ pub async fn recaption_image(args: RecaptionImageArgs<'_>) -> Result<RecaptionIm
|
||||
.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)
|
||||
|
||||
@@ -60,7 +60,7 @@ pub async fn create_world(args: CreateWorldArgs<'_>) -> Result<CreateWorldRespon
|
||||
.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)
|
||||
|
||||
@@ -62,7 +62,7 @@ pub async fn poll_world_status(args: PollWorldStatusArgs<'_>) -> Result<PollWorl
|
||||
.header(ACCEPT_LANGUAGE, "en-US,en;q=0.5")
|
||||
.header(ACCEPT_ENCODING, "gzip, deflate, br, zstd")
|
||||
.header(REFERER, REFERER_VALUE)
|
||||
.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)
|
||||
|
||||
@@ -0,0 +1,101 @@
|
||||
use base64::Engine;
|
||||
use chrono::{DateTime, Utc};
|
||||
use jwt_light::common_claims::CommonClaims;
|
||||
use jwt_light::error::JwtError;
|
||||
use jwt_light::parse_jwt_claims_trait::ParseJwtClaims;
|
||||
use serde_json::{Map, Value};
|
||||
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct JwtClaims {
|
||||
/// From the "iat" field
|
||||
pub created: DateTime<Utc>,
|
||||
|
||||
/// From the "exp" field
|
||||
pub expiration: DateTime<Utc>,
|
||||
|
||||
/// From the "user_id" field
|
||||
pub user_id: Option<String>,
|
||||
|
||||
/// From the "email" field
|
||||
pub email: Option<String>,
|
||||
}
|
||||
|
||||
impl ParseJwtClaims for JwtClaims {
|
||||
fn extract_claims(common_claims: CommonClaims, extra_claims: Map<String, Value>) -> Result<Self, JwtError> {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@@ -1,2 +1,4 @@
|
||||
pub mod world_labs_bearer_token;
|
||||
pub mod world_labs_cookies;
|
||||
pub mod worldlabs_refresh_token;
|
||||
pub mod jwt_claims;
|
||||
|
||||
@@ -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, WorldLabsError> {
|
||||
JwtClaims::parse_claims(&self.bearer_token)
|
||||
.map_err(|err| {
|
||||
error!("Failed to parse bearer token into JWT claims: {}", err);
|
||||
WorldLabsError::Client(WorldLabsClientError::FailedToParseJwtClaims(err))
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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()
|
||||
}
|
||||
}
|
||||
@@ -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),
|
||||
|
||||
@@ -1 +1 @@
|
||||
mod upload_image_and_create_world_with_retry;
|
||||
pub mod upload_image_and_create_world_with_retry;
|
||||
+31
-4
@@ -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<Duration>,
|
||||
pub file_bytes: Vec<u8>,
|
||||
pub file: FileBytesOrPath,
|
||||
}
|
||||
|
||||
pub enum FileBytesOrPath {
|
||||
Bytes(Vec<u8>),
|
||||
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<UploadImageAndCreateWorldWithRetryResponse, WorldLabsError> {
|
||||
|
||||
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);
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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<StorytellerError> for GenerateError {
|
||||
Self::ProviderFailure(ProviderFailureReason::StorytellerError(value))
|
||||
}
|
||||
}
|
||||
|
||||
impl From<WorldLabsError> for GenerateError {
|
||||
fn from(value: WorldLabsError) -> Self {
|
||||
Self::ProviderFailure(ProviderFailureReason::WorldLabsError(value))
|
||||
}
|
||||
}
|
||||
|
||||
+233
@@ -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<GaussianModel>,
|
||||
|
||||
/// OPTIONAL.
|
||||
/// Text prompt for the gaussian generation. Required.
|
||||
pub prompt: Option<String>,
|
||||
|
||||
/// OPTIONAL.
|
||||
/// Images for the gaussian generation
|
||||
pub image_media_tokens: Option<Vec<MediaFileToken>>,
|
||||
|
||||
/// OPTIONAL.
|
||||
/// Name of the frontend caller.
|
||||
/// We'll use this to selectively trigger events.
|
||||
pub frontend_caller: Option<TauriCommandCaller>,
|
||||
|
||||
/// 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<String>,
|
||||
|
||||
/// 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<String>,
|
||||
}
|
||||
|
||||
|
||||
#[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<EnqueueImageToGaussianSuccessResponse, EnqueueImageToGaussianErrorType, ()> {
|
||||
|
||||
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<TaskEnqueueSuccess, GenerateError> {
|
||||
|
||||
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<TaskEnqueueSuccess, GenerateError> {
|
||||
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1,2 @@
|
||||
pub mod enqueue_image_to_gaussian_command;
|
||||
mod worldlabs;
|
||||
+215
@@ -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<TaskEnqueueSuccess, GenerateError> {
|
||||
|
||||
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<PathBuf, GenerateError> {
|
||||
|
||||
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<Vec<PathBuf>, 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<Vec<TextToImageError>>
|
||||
// ) -> Result<TaskEnqueueSuccess, GenerateError> {
|
||||
// if let Some(errors) = maybe_errors {
|
||||
// if !errors.is_empty() {
|
||||
// let messages: Vec<String> = 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))
|
||||
// }
|
||||
@@ -0,0 +1 @@
|
||||
pub (super) mod handle_worldlabs_marble;
|
||||
@@ -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;
|
||||
pub mod text_to_image;
|
||||
pub mod image_to_gaussian;
|
||||
@@ -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 -
|
||||
|
||||
+26
@@ -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<GeneratedGaussian>,
|
||||
pub maybe_frontend_subscriber_id: Option<String>,
|
||||
pub maybe_frontend_subscriber_payload: Option<String>,
|
||||
}
|
||||
|
||||
#[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<String>,
|
||||
}
|
||||
|
||||
impl BasicSendableEvent for GaussianGenerationCompleteEvent {
|
||||
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GaussianGenerationCompleteEvent;
|
||||
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
|
||||
@@ -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")
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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,
|
||||
}
|
||||
}
|
||||
|
||||
@@ -0,0 +1,7 @@
|
||||
use errors::AnyhowResult;
|
||||
use tauri::WebviewWindow;
|
||||
|
||||
pub fn get_webview_window_url_path(webview: &WebviewWindow) -> AnyhowResult<String> {
|
||||
let url = webview.url()?;
|
||||
Ok(url.path().to_string())
|
||||
}
|
||||
@@ -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;
|
||||
|
||||
@@ -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"))
|
||||
|
||||
@@ -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.
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
pub mod grok;
|
||||
pub mod midjourney;
|
||||
pub mod sora;
|
||||
pub mod storyteller;
|
||||
pub mod storyteller;
|
||||
pub mod worldlabs;
|
||||
+32
-1
@@ -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,
|
||||
|
||||
@@ -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;
|
||||
+44
@@ -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(())
|
||||
}
|
||||
+63
@@ -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<String>,
|
||||
pub can_clear_state: bool,
|
||||
}
|
||||
|
||||
impl SerializeMarker for WorldlabsGetCredentialInfoResponse {}
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn worldlabs_get_credential_info_command(
|
||||
creds_manager: State<'_, WorldlabsCredentialManager>,
|
||||
) -> ResponseOrErrorMessage<WorldlabsGetCredentialInfoResponse> {
|
||||
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<WorldlabsGetCredentialInfoResponse> {
|
||||
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,
|
||||
})
|
||||
}
|
||||
+45
@@ -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<String, String> {
|
||||
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(())
|
||||
}
|
||||
+59
@@ -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(())
|
||||
}
|
||||
@@ -0,0 +1,4 @@
|
||||
pub mod commands;
|
||||
pub mod state;
|
||||
pub mod threads;
|
||||
pub mod windows;
|
||||
@@ -0,0 +1,6 @@
|
||||
pub mod worldlabs_credential_manager;
|
||||
pub mod worldlabs_serializable_state;
|
||||
pub mod worldlabs_bearer_bridge;
|
||||
|
||||
// Private
|
||||
mod worldlabs_credential_holder;
|
||||
@@ -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<RwLock<Option<WorldlabsBearerBridgeInner>>>
|
||||
}
|
||||
|
||||
#[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<Option<WorldlabsBearerBridgeInner>, 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)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<CookieStore>,
|
||||
|
||||
/// Copy of cookies, specifically for the client library.
|
||||
/// NOT PERSISTED TO DISK.
|
||||
pub world_labs_cookies: Option<WorldLabsCookies>,
|
||||
|
||||
/// Bearer token
|
||||
/// This needs to be persisted to disk.
|
||||
pub world_labs_bearer_token: Option<WorldLabsBearerToken>,
|
||||
|
||||
/// Refresh token
|
||||
/// This needs to be persisted to disk.
|
||||
pub world_labs_refresh_token: Option<WorldLabsRefreshToken>,
|
||||
}
|
||||
|
||||
impl WorldlabsCredentialHolder {
|
||||
pub fn empty() -> Self {
|
||||
Self {
|
||||
browser_cookies: None,
|
||||
world_labs_cookies: None,
|
||||
world_labs_bearer_token: None,
|
||||
world_labs_refresh_token: None,
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -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<RwLock<WorldlabsCredentialHolder>>,
|
||||
|
||||
//user_info: Arc<RwLock<Option<MidjourneyUserInfo>>>,
|
||||
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<Option<CookieStore>> {
|
||||
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<Option<String>> {
|
||||
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<Option<WorldLabsCookies>> {
|
||||
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<Option<WorldLabsBearerToken>> {
|
||||
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<Option<WorldLabsRefreshToken>> {
|
||||
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<bool> {
|
||||
self.session_appears_active()
|
||||
}
|
||||
|
||||
// NB: This is just a heuristic. We'll add better checks later.
|
||||
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
|
||||
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(())
|
||||
}
|
||||
}
|
||||
@@ -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<SerializableCookieStore>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) bearer_token: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) refresh_token: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) user_id: Option<String>,
|
||||
|
||||
#[serde(skip_serializing_if = "Option::is_none")]
|
||||
pub (super) user_email: Option<String>,
|
||||
}
|
||||
|
||||
impl WorldlabsSerializableState {
|
||||
pub fn read_from_disk(root: &AppDataRoot) -> anyhow::Result<Option<WorldlabsSerializableState>> {
|
||||
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))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod worldlabs_marble_task_polling;
|
||||
+386
@@ -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::<HashMap<String, Task>>();
|
||||
|
||||
|
||||
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<PathBuf> {
|
||||
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)
|
||||
}
|
||||
@@ -0,0 +1 @@
|
||||
pub mod worldlabs_login_window;
|
||||
@@ -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;
|
||||
+91
@@ -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";
|
||||
}
|
||||
})();
|
||||
"#;
|
||||
+39
@@ -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<Url> = Lazy::new(|| {
|
||||
Url::parse("https://worldlabs.ai").expect("URL should parse")
|
||||
});
|
||||
|
||||
pub fn worldlabs_login_webview_extract_cookies(webview: &WebviewWindow) -> AnyhowResult<CookieStore> {
|
||||
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<Vec<Cookie>> {
|
||||
//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)
|
||||
}
|
||||
+108
@@ -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<Url> = Lazy::new(|| {
|
||||
Url::parse("https://google.com").expect("URL should parse")
|
||||
});
|
||||
|
||||
pub (super) static WORLDLABS_HOMEPAGE_URL: Lazy<Url> = Lazy::new(|| {
|
||||
Url::parse("https://marble.worldlabs.ai/").expect("URL should parse")
|
||||
});
|
||||
|
||||
pub (super) static WORLDLABS_LOGIN_URL: Lazy<Url> = 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(())
|
||||
}
|
||||
+129
@@ -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<bool> {
|
||||
|
||||
// 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)
|
||||
}
|
||||
@@ -10,6 +10,7 @@
|
||||
"beforeBuildCommand": "nx run editor3d:build"
|
||||
},
|
||||
"app": {
|
||||
"withGlobalTauri": true,
|
||||
"windows": [
|
||||
],
|
||||
"security": {
|
||||
|
||||
@@ -6,6 +6,7 @@
|
||||
"beforeBuildCommand": "npx nx run editor3d:build"
|
||||
},
|
||||
"app": {
|
||||
"withGlobalTauri": true,
|
||||
"windows": [
|
||||
]
|
||||
}
|
||||
|
||||
@@ -10,6 +10,7 @@
|
||||
"beforeBuildCommand": "nx run editor3d:build"
|
||||
},
|
||||
"app": {
|
||||
"withGlobalTauri": true,
|
||||
"windows": [
|
||||
],
|
||||
"security": {
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
}
|
||||
|
||||
@@ -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);
|
||||
}
|
||||
|
||||
@@ -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));
|
||||
|
||||
@@ -14,7 +14,9 @@
|
||||
gtag('js', new Date());
|
||||
gtag('config', 'G-Y4Y8KMLPQ7');
|
||||
</script>
|
||||
<script crossOrigin="anonymous" src="//unpkg.com/react-scan/dist/auto.global.js"></script>
|
||||
|
||||
<!---->
|
||||
<!-- <script crossOrigin="anonymous" src="//unpkg.com/react-scan/dist/auto.global.js"></script> -->
|
||||
|
||||
<!-- Microsoft Clarity -->
|
||||
<script type="text/javascript">
|
||||
|
||||
@@ -22,6 +22,7 @@ import {
|
||||
EnqueueImageTo3dObject,
|
||||
EnqueueImageTo3dObjectModel,
|
||||
EnqueueImageTo3dWorld,
|
||||
EnqueueImageToGaussian,
|
||||
EnqueueImageTo3dWorldModel,
|
||||
} from "@storyteller/tauri-api";
|
||||
import { toast } from "react-hot-toast";
|
||||
@@ -31,6 +32,7 @@ import { addObject } from "../../pages/PageEnigma/signals/objectGroup/addObject"
|
||||
import { set3DPageMounted } from "../../pages/PageEnigma/Editor/editor";
|
||||
import { AssetType } from "~/enums";
|
||||
import type { MediaItem } from "../../pages/PageEnigma/models";
|
||||
import { GAUSSIAN_MODELS } from "libs/model-list/src/lib/lists/GaussianModels";
|
||||
|
||||
type Mode = "image" | "text";
|
||||
type Variant = "object" | "world";
|
||||
@@ -206,7 +208,9 @@ export const ImageTo3DExperience = ({
|
||||
if (activeMode === "image" && !uploadedMediaToken) return;
|
||||
if (activeMode === "text" && prompt.trim().length <= 3) return;
|
||||
|
||||
console.log("Set is generating...");
|
||||
setIsGenerating(true);
|
||||
|
||||
const snapshotPrompt = prompt.trim();
|
||||
const snapshotPreview = uploadedPreview || undefined;
|
||||
const snapshotName = uploadedName;
|
||||
@@ -239,9 +243,9 @@ export const ImageTo3DExperience = ({
|
||||
frontend_caller: "mini_app",
|
||||
frontend_subscriber_id: subscriberId,
|
||||
})
|
||||
: await EnqueueImageTo3dWorld({
|
||||
image_media_token: uploadedMediaToken || undefined,
|
||||
model: EnqueueImageTo3dWorldModel.Hunyuan3d2,
|
||||
: await EnqueueImageToGaussian({
|
||||
image_media_tokens: uploadedMediaToken ? [uploadedMediaToken] : undefined,
|
||||
model: GAUSSIAN_MODELS[0],
|
||||
frontend_caller: "mini_app",
|
||||
frontend_subscriber_id: subscriberId,
|
||||
});
|
||||
|
||||
+3
-3
@@ -1,8 +1,8 @@
|
||||
import { SoraAccountBlock } from "./SoraAccountBlock";
|
||||
import { ArtcraftAccountBlock } from "./ArtcraftAccountBlock";
|
||||
import { FalApiKeyBlock } from "./FalApiKeyBlock";
|
||||
import { MidjourneyAccountBlock } from "./MidjourneyAccountBlock";
|
||||
import { GrokAccountBlock } from "./GrokAccountBlock";
|
||||
import { WorldLabsAccountBlock } from "./WorldLabsAccountBlock";
|
||||
|
||||
interface AccountSettingsPaneProps {
|
||||
globalAccountLogoutCallback: () => void;
|
||||
@@ -42,10 +42,10 @@ export const AccountSettingsPane = ({
|
||||
<ArtcraftAccountBlock
|
||||
globalAccountLogoutCallback={globalAccountLogoutCallback}
|
||||
/>
|
||||
<MidjourneyAccountBlock />
|
||||
<GrokAccountBlock />
|
||||
<MidjourneyAccountBlock />
|
||||
<SoraAccountBlock />
|
||||
<FalApiKeyBlock />
|
||||
<WorldLabsAccountBlock />
|
||||
{/*
|
||||
<OpenAIApiKeyBlock/>
|
||||
<hr />
|
||||
|
||||
-71
@@ -1,71 +0,0 @@
|
||||
import { useState, useEffect } from "react";
|
||||
import { GetFalApiKey, SetFalApiKey } from "@storyteller/tauri-api";
|
||||
import { Input } from "@storyteller/ui-input";
|
||||
|
||||
|
||||
export const FalApiKeyBlock = () => {
|
||||
const [falApiKey, setFalApiKey] = useState<string | undefined>(
|
||||
undefined
|
||||
);
|
||||
const [isCheckingFalApiKey, setIsCheckingFalApiKey] =
|
||||
useState(false);
|
||||
|
||||
console.log(isCheckingFalApiKey);
|
||||
|
||||
useEffect(() => {
|
||||
const fetchKey = async () => {
|
||||
setIsCheckingFalApiKey(true);
|
||||
try {
|
||||
const result = await GetFalApiKey();
|
||||
if ("payload" in result) {
|
||||
setFalApiKey(result.payload.key);
|
||||
}
|
||||
} catch (e) {
|
||||
console.error("Error fetching Fal API key", e);
|
||||
setFalApiKey(undefined);
|
||||
} finally {
|
||||
setIsCheckingFalApiKey(false);
|
||||
}
|
||||
};
|
||||
fetchKey();
|
||||
}, []);
|
||||
|
||||
const syncFalApiKey = async (key: string) => {
|
||||
await SetFalApiKey(key);
|
||||
setFalApiKey(key);
|
||||
};
|
||||
|
||||
return (
|
||||
<div className="flex justify-between items-center">
|
||||
<span>Fal API Key:</span>
|
||||
<Input
|
||||
value={falApiKey}
|
||||
onChange={(e) => syncFalApiKey(e.target.value)}
|
||||
placeholder="Enter Fal API key"
|
||||
/>
|
||||
{/*<Button
|
||||
variant={
|
||||
isCheckingArtcraftSession
|
||||
? "secondary"
|
||||
: isLoggedIn
|
||||
? "destructive"
|
||||
: "primary"
|
||||
}
|
||||
className="h-[30px]"
|
||||
onClick={handleArtcraftButton}
|
||||
disabled={isCheckingArtcraftSession}
|
||||
>
|
||||
{isCheckingArtcraftSession ? (
|
||||
<FontAwesomeIcon
|
||||
icon={faSpinnerThird}
|
||||
className="animate-spin text-sm"
|
||||
/>
|
||||
) : isLoggedIn ? (
|
||||
"Log Out"
|
||||
) : (
|
||||
"Log In"
|
||||
)}
|
||||
</Button>*/}
|
||||
</div>
|
||||
);
|
||||
};
|
||||
+1
-1
@@ -63,7 +63,7 @@ export const SoraAccountBlock = () => {
|
||||
|
||||
return(
|
||||
<div className="flex justify-between items-center">
|
||||
<span>OpenAI / Sora Account:</span>
|
||||
<span>Sora (OpenAI) Account:</span>
|
||||
<pre>{soraSession?.payload?.maybe_email || "Not logged in"}</pre>
|
||||
<Button
|
||||
variant={
|
||||
|
||||
+91
@@ -0,0 +1,91 @@
|
||||
import { Button } from "@storyteller/ui-button";
|
||||
import { invoke } from "@tauri-apps/api/core";
|
||||
import { useEffect, useState } from "react";
|
||||
import { faSpinnerThird } from "@fortawesome/pro-solid-svg-icons";
|
||||
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
|
||||
import { useRefreshAccountStateEvent } from "@storyteller/tauri-events";
|
||||
import { RefreshAccountStateEvent } from "@storyteller/tauri-events";
|
||||
import { WorldLabsGetCredentialInfo, WorldLabsGetCredentialInfoSuccess } from "@storyteller/tauri-api";
|
||||
|
||||
export const WorldLabsAccountBlock = () => {
|
||||
const [worldlabsSession, setWorldlabsSession] = useState<WorldLabsGetCredentialInfoSuccess| undefined>(undefined);
|
||||
const [isCheckingWorldlabsSession, setIsCheckingWorldlabsSession] = useState(false);
|
||||
|
||||
const fetchSession = async () => {
|
||||
setIsCheckingWorldlabsSession(true);
|
||||
try {
|
||||
const result = await WorldLabsGetCredentialInfo();
|
||||
setWorldlabsSession(result);
|
||||
} catch (e) {
|
||||
console.error("Error fetching WorldLabs session", e);
|
||||
setWorldlabsSession(undefined);
|
||||
} finally {
|
||||
setIsCheckingWorldlabsSession(false);
|
||||
}
|
||||
};
|
||||
|
||||
useEffect(() => {
|
||||
fetchSession();
|
||||
}, []);
|
||||
|
||||
useRefreshAccountStateEvent(async (event: RefreshAccountStateEvent) => {
|
||||
fetchSession();
|
||||
});
|
||||
|
||||
const clearState = async() => {
|
||||
try {
|
||||
await invoke("worldlabs_clear_credentials_command");
|
||||
} catch (e) {
|
||||
console.error("Error clearing WorldLabs credentials", e);
|
||||
}
|
||||
}
|
||||
|
||||
const openLogin = async() => {
|
||||
try {
|
||||
await invoke("worldlabs_open_login_command");
|
||||
} catch (e) {
|
||||
console.error("Error opening WorldLabs login", e);
|
||||
}
|
||||
}
|
||||
|
||||
const handleWorldlabsButton = async () => {
|
||||
if (worldlabsSession?.payload?.can_clear_state) {
|
||||
await clearState();
|
||||
setWorldlabsSession(undefined);
|
||||
} else {
|
||||
await openLogin();
|
||||
}
|
||||
};
|
||||
|
||||
console.log("worldlabs session", worldlabsSession);
|
||||
|
||||
return(
|
||||
<div className="flex justify-between items-center">
|
||||
<span>WorldLabs Account:</span>
|
||||
<pre>{worldlabsSession?.payload?.maybe_email || "Not logged in"}</pre>
|
||||
<Button
|
||||
variant={
|
||||
worldlabsSession?.payload?.can_clear_state && !isCheckingWorldlabsSession
|
||||
? "destructive"
|
||||
: worldlabsSession?.payload?.can_clear_state
|
||||
? "primary"
|
||||
: "secondary"
|
||||
}
|
||||
className="h-[30px]"
|
||||
onClick={handleWorldlabsButton}
|
||||
disabled={isCheckingWorldlabsSession}
|
||||
>
|
||||
{isCheckingWorldlabsSession ? (
|
||||
<FontAwesomeIcon
|
||||
icon={faSpinnerThird}
|
||||
className="animate-spin text-sm"
|
||||
/>
|
||||
) : worldlabsSession?.payload?.can_clear_state ? (
|
||||
"Disconnect"
|
||||
) : (
|
||||
"Connect"
|
||||
)}
|
||||
</Button>
|
||||
</div>
|
||||
)
|
||||
}
|
||||
@@ -14,4 +14,6 @@ export * from "./lib/lists/AllModels.js";
|
||||
export * from "./lib/lists/AllModelsById.js";
|
||||
export * from "./lib/lists/ImageModels.js";
|
||||
export * from "./lib/lists/VideoModels.js";
|
||||
export * from "./lib/lists/GaussianModels.js";
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,27 @@
|
||||
import { ModelCreator } from "src/index.js";
|
||||
import { Model, ModelKind } from "./Model.js";
|
||||
import { ModelCategory } from "../legacy/ModelConfig.js";
|
||||
import { ModelTag } from "./metadata/ModelTag.js";
|
||||
|
||||
export class GaussianModel extends Model {
|
||||
// Typescript type discriminator property
|
||||
// Since Vite minification and class name mangling can break instanceof checks,
|
||||
// we have a type discriminator property to check against.
|
||||
override readonly kind: ModelKind = "gaussian_model";
|
||||
|
||||
|
||||
constructor(args: {
|
||||
id: string;
|
||||
tauriId: string;
|
||||
fullName: string;
|
||||
category: ModelCategory;
|
||||
creator: ModelCreator;
|
||||
selectorName: string;
|
||||
selectorDescription: string;
|
||||
selectorBadges: string[];
|
||||
tags?: ModelTag[];
|
||||
progressBarTime?: number;
|
||||
}) {
|
||||
super(args);
|
||||
}
|
||||
}
|
||||
@@ -2,7 +2,7 @@ import { ModelCreator } from "./metadata/ModelCreator.js";
|
||||
import { ModelCategory, ModelConfig } from "../legacy/ModelConfig.js";
|
||||
import { ModelTag } from "./metadata/ModelTag.js";
|
||||
|
||||
export type ModelKind = "model" | "image_model" | "video_model";
|
||||
export type ModelKind = "model" | "image_model" | "video_model" | "gaussian_model";
|
||||
|
||||
// NB: Do not create instances of this class directly, use subclasses.
|
||||
export class Model {
|
||||
|
||||
@@ -20,4 +20,5 @@ export enum ModelCreator {
|
||||
Vidu = "Vidu",
|
||||
ArtCraft = "ArtCraft",
|
||||
Grok = "Grok",
|
||||
WorldLabs = "WorldLabs",
|
||||
}
|
||||
|
||||
@@ -1,7 +1,7 @@
|
||||
import { ModelInfo } from "./ModelInfo.js";
|
||||
import { ModelTag } from "../classes/metadata/ModelTag.js";
|
||||
|
||||
export type ModelCategory = "image" | "video";
|
||||
export type ModelCategory = "image" | "video" | "gaussian";
|
||||
|
||||
// Centralized model capability definition
|
||||
// For setting options for the model
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
import { IMAGE_MODELS } from "./ImageModels.js";
|
||||
import { VIDEO_MODELS } from "./VideoModels.js";
|
||||
import { Model } from "../classes/Model.js";
|
||||
import { GAUSSIAN_MODELS } from "./GaussianModels.js";
|
||||
|
||||
export const ALL_MODELS_LIST: Model[] = [...IMAGE_MODELS, ...VIDEO_MODELS];
|
||||
export const ALL_MODELS_LIST: Model[] = [...IMAGE_MODELS, ...VIDEO_MODELS, ...GAUSSIAN_MODELS];
|
||||
|
||||
@@ -0,0 +1,26 @@
|
||||
import { GaussianModel } from "../classes/GaussianModel.js";
|
||||
import { ModelCreator } from "../classes/metadata/ModelCreator.js";
|
||||
|
||||
// TODO: Some of the model configs, such as generation counts, are authoritatively controlled in `legacy/Models.ts`
|
||||
|
||||
export const GAUSSIAN_MODELS : GaussianModel [] = [
|
||||
new GaussianModel({
|
||||
id: "world_labs_marble",
|
||||
tauriId: "world_labs_marble",
|
||||
fullName: "WorldLabs Marble",
|
||||
category: "gaussian",
|
||||
creator: ModelCreator.WorldLabs,
|
||||
selectorName: "WorldLabs Marble",
|
||||
selectorDescription: "Amazing",
|
||||
selectorBadges: ["10 sec."],
|
||||
progressBarTime: 10000,
|
||||
}),
|
||||
];
|
||||
|
||||
export const GAUSSIAN_MODELS_BY_ID: Map<string, GaussianModel> = new Map(
|
||||
GAUSSIAN_MODELS.map((model) => [model.id, model]),
|
||||
);
|
||||
|
||||
if (GAUSSIAN_MODELS_BY_ID.size !== GAUSSIAN_MODELS.length) {
|
||||
throw new Error("All gaussian models must have unique IDs");
|
||||
}
|
||||
@@ -1,11 +1,14 @@
|
||||
export * from "./lib/artcraft/ArtcraftGetCredits.js";
|
||||
export * from "./lib/artcraft/ArtcraftGetSubscription.js";
|
||||
export * from "./lib/artcraft/media_files/MediaFileDelete.js";
|
||||
export * from "./lib/artcraft/task_queue/GetTaskQueue.js";
|
||||
export * from "./lib/artcraft/task_queue/MarkTaskAsDismissed.js";
|
||||
export * from "./lib/common/CommandStatus.js";
|
||||
export * from "./lib/enqueue/EnqueueEditImage.js";
|
||||
export * from "./lib/enqueue/EnqueueImageBgRemovalCommand.js";
|
||||
export * from "./lib/enqueue/EnqueueImageInpaint.js";
|
||||
export * from "./lib/enqueue/EnqueueImageTo3dObject.js";
|
||||
export * from "./lib/enqueue/EnqueueImageToGaussian.js";
|
||||
export * from "./lib/enqueue/EnqueueImageTo3dWorld.js";
|
||||
export * from "./lib/enqueue/EnqueueImageToVideo.js";
|
||||
export * from "./lib/enqueue/EnqueueTextToImage.js";
|
||||
@@ -22,8 +25,6 @@ export * from "./lib/openai/GetOpenAIApiKey.js";
|
||||
export * from "./lib/openai/SetOpenAIApiKey.js";
|
||||
export * from "./lib/providers/GetProviderOrder.js";
|
||||
export * from "./lib/providers/Provider.js";
|
||||
export * from "./lib/artcraft/task_queue/GetTaskQueue.js";
|
||||
export * from "./lib/artcraft/task_queue/MarkTaskAsDismissed.js";
|
||||
export * from "./lib/providers/SetProviderOrder.js";
|
||||
export * from "./lib/settings/GetAppPreferences.js";
|
||||
export * from "./lib/settings/UpdateAppPreference.js";
|
||||
@@ -35,3 +36,4 @@ export * from "./lib/sora/useSoraLoginListener.js";
|
||||
export * from "./lib/sora/waitForSoraLogin.js";
|
||||
export * from "./lib/system/GetAppInfo.js";
|
||||
export * from "./lib/util/LoadWithoutCors.js";
|
||||
export * from "./lib/worldlabs/WorldLabsGetCredentialInfo.js";
|
||||
|
||||
@@ -0,0 +1,81 @@
|
||||
import { invoke } from "@tauri-apps/api/core";
|
||||
import { CommandResult } from "../common/CommandStatus";
|
||||
import { Model } from "@storyteller/model-list";
|
||||
|
||||
export enum EnqueueImageToGaussianErrorType {
|
||||
/// Caller didn't specify a model
|
||||
ModelNotSpecified = "model_not_specified",
|
||||
/// Generic server error
|
||||
ServerError = "server_error",
|
||||
/// No Storyteller creds
|
||||
NeedsStorytellerCredentials = "needs_storyteller_credentials",
|
||||
/// No provider available
|
||||
NoProviderAvailable = "no_provider_available",
|
||||
}
|
||||
|
||||
export interface EnqueueImageToGaussianRequest {
|
||||
// Required. The model to use.
|
||||
model?: Model;
|
||||
|
||||
// Optional. Text prompt to direct the gaussian.
|
||||
prompt?: string;
|
||||
|
||||
// Required. Starting frame.
|
||||
image_media_tokens?: string[];
|
||||
|
||||
// Optional frontend state to return later.
|
||||
frontend_caller?: string;
|
||||
|
||||
// Optional frontend state to return later.
|
||||
frontend_subscriber_id?: string;
|
||||
}
|
||||
|
||||
interface RawEnqueueImageToGaussianRequest {
|
||||
model?: string;
|
||||
prompt?: string;
|
||||
image_media_tokens?: string[];
|
||||
frontend_caller?: string;
|
||||
frontend_subscriber_id?: string;
|
||||
}
|
||||
|
||||
export interface EnqueueImageToGaussianError extends CommandResult {
|
||||
error_type: EnqueueImageToGaussianErrorType;
|
||||
error_message?: string;
|
||||
}
|
||||
|
||||
export type EnqueueImageToGaussianPayload = Record<string, never>;
|
||||
|
||||
export interface EnqueueImageToGaussianSuccess extends CommandResult {
|
||||
payload: EnqueueImageToGaussianPayload;
|
||||
}
|
||||
|
||||
export type EnqueueImageToGaussianResult =
|
||||
| EnqueueImageToGaussianSuccess
|
||||
| EnqueueImageToGaussianError;
|
||||
|
||||
export const EnqueueImageToGaussian = async (
|
||||
request: EnqueueImageToGaussianRequest
|
||||
): Promise<EnqueueImageToGaussianResult> => {
|
||||
const mutableRequest: RawEnqueueImageToGaussianRequest = {
|
||||
model: request.model?.tauriId,
|
||||
image_media_tokens: request.image_media_tokens,
|
||||
};
|
||||
|
||||
if (request.prompt) {
|
||||
mutableRequest.prompt = request.prompt;
|
||||
}
|
||||
|
||||
if (request.frontend_caller) {
|
||||
mutableRequest.frontend_caller = request.frontend_caller;
|
||||
}
|
||||
|
||||
if (request.frontend_subscriber_id) {
|
||||
mutableRequest.frontend_subscriber_id = request.frontend_subscriber_id;
|
||||
}
|
||||
|
||||
const result = await invoke("enqueue_image_to_gaussian_command", {
|
||||
request: mutableRequest,
|
||||
});
|
||||
|
||||
return result as EnqueueImageToGaussianResult;
|
||||
};
|
||||
@@ -0,0 +1,16 @@
|
||||
import { invoke } from "@tauri-apps/api/core";
|
||||
import { CommandResult } from "../common/CommandStatus";
|
||||
|
||||
export interface WorldLabsGetCredentialInfoPayload {
|
||||
maybe_email?: string;
|
||||
can_clear_state: boolean;
|
||||
}
|
||||
|
||||
export interface WorldLabsGetCredentialInfoSuccess extends CommandResult {
|
||||
payload: WorldLabsGetCredentialInfoPayload;
|
||||
}
|
||||
|
||||
export const WorldLabsGetCredentialInfo = async (): Promise<WorldLabsGetCredentialInfoSuccess> => {
|
||||
const result = await invoke("worldlabs_get_credential_info_command");
|
||||
return result as WorldLabsGetCredentialInfoSuccess;
|
||||
};
|
||||
Reference in New Issue
Block a user