worldlabs tauri (#933)

not finished: need refresh tokens and thumbnails
This commit is contained in:
Brandon Thomas
2025-12-18 22:50:36 -05:00
committed by GitHub
parent 6815040cb1
commit 5688d8cef7
84 changed files with 2588 additions and 108 deletions
Generated
+2
View File
@@ -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",
+1
View File
@@ -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
@@ -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
@@ -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)
@@ -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)
@@ -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)
@@ -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)
@@ -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;
@@ -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);
+1
View File
@@ -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))
}
}
@@ -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;
@@ -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 -
@@ -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;
+24 -2
View File
@@ -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.
+2 -1
View File
@@ -1,4 +1,5 @@
pub mod grok;
pub mod midjourney;
pub mod sora;
pub mod storyteller;
pub mod storyteller;
pub mod worldlabs;
@@ -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;
@@ -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(())
}
@@ -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,
})
}
@@ -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(())
}
@@ -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;
@@ -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;
@@ -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";
}
})();
"#;
@@ -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)
}
@@ -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(())
}
@@ -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": [
]
}
+1
View File
@@ -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));
+3 -1
View File
@@ -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,
});
@@ -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 />
@@ -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>
);
};
@@ -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={
@@ -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>
)
}
+2
View File
@@ -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");
}
+4 -2
View File
@@ -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;
};