diff --git a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs new file mode 100644 index 0000000000..f66c4d3aa3 --- /dev/null +++ b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/legacy_upload_media_file_from_file.rs @@ -0,0 +1,143 @@ +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +use crate::error::api_error::ApiError; +use crate::error::client_error::ClientError; +use crate::error::storyteller_error::StorytellerError; +use crate::utils::api_host::ApiHost; +use crate::utils::constants::{APPLICATION_JSON, USER_AGENT}; +use crate::utils::filter_bad_response::filter_bad_response; +use crate::utils::http_get_anonymous::http_get_anonymous; +use chrono::{DateTime, Utc}; +use enums::by_table::media_files::media_file_animation_type::MediaFileAnimationType; +use enums::by_table::media_files::media_file_class::MediaFileClass; +use enums::by_table::media_files::media_file_engine_category::MediaFileEngineCategory; +use enums::by_table::media_files::media_file_type::MediaFileType; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::common::generation_provider::GenerationProvider; +use enums::common::visibility::Visibility; +use enums::no_table::style_transfer::style_transfer_name::StyleTransferName; +use enums_public::by_table::model_weights::public_weights_types::PublicWeightsType; +use idempotency::uuid::generate_random_uuid; +use log::debug; +use reqwest::multipart::Form; +use reqwest::Client; +use serde_derive::{Deserialize, Serialize}; +use std::fs::File; +use std::io::Read; +use std::path::Path; +use tokens::tokens::batch_generations::BatchGenerationToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::model_weights::ModelWeightToken; +use tokens::tokens::prompts::PromptToken; +use uuid::uuid; + +pub struct LegacyUploadMediaFileFromFileArgs<'a, P: AsRef> { + pub api_host: &'a ApiHost, + pub maybe_creds: Option<&'a StorytellerCredentialSet>, + + // NB: Path needs to be owned for the request. + pub path: P, + + // /// If true, we should hide the image from the user's gallery. + // pub is_intermediate_system_file: bool, + + // /// If provided, this is the prompt that this image is associated with. + // /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set. + // pub maybe_prompt_token: Option<&'a PromptToken>, + + // /// If provided, this is the service provider that created the image. + // /// NOTE: Cannot set `is_intermediate_system_file = true` if this is set. + // pub maybe_generation_provider: Option, + + // /// If provided, this groups the file into a batch + // /// TODO: This shouldn't be set clientside without the backend generating the token + // /// and cryptographically securing it. But we need to go fast here. + // pub maybe_batch_token: Option<&'a BatchGenerationToken>, +} + + +/// Upload using the legacy endpoint (spz, etc.) +pub async fn legacy_upload_media_file_from_file>( + args: LegacyUploadMediaFileFromFileArgs<'_, P>, +) -> Result { + + validate_args(&args)?; + + let url = get_route(args.api_host); + + debug!("Requesting {:?}", &url); + + //let mut file = File::open(path)?; + + let client = Client::builder() + .gzip(true) + .build() + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + let mut form = Form::new() + .text("uuid_idempotency_token", generate_random_uuid()) + .file("file", args.path) + .await + .map_err(|err| StorytellerError::Client(ClientError::from(err)))?; + + //if args.is_intermediate_system_file { + // form = form.text("is_intermediate_system_file", "true"); + //} + + //if let Some(prompt_token) = &args.maybe_prompt_token { + // form = form.text("maybe_prompt_token", prompt_token.to_string()); + //} + + //if let Some(batch_token) = &args.maybe_batch_token { + // form = form.text("maybe_batch_token", batch_token.to_string()); + //} + + let mut request_builder = client.post(url) + .header("User-Agent", USER_AGENT) + .header("Accept", APPLICATION_JSON); + + if let Some(creds) = args.maybe_creds { + if let Some(header) = &creds.maybe_as_cookie_header() { + request_builder = request_builder.header("Cookie", header); + } + } + + let response = request_builder + .multipart(form) + .send() + .await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let response = filter_bad_response(response).await?; + let response_body = &response.text().await + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + let media_file = serde_json::from_str(&response_body) + .map_err(|err| StorytellerError::Api(ApiError::from(err)))?; + + Ok(media_file) +} + +fn get_route(api_host: &ApiHost) -> String { + let api_hostname_and_scheme = api_host.to_api_hostname_and_scheme(); + format!("{}/v1/media_files/upload", api_hostname_and_scheme) +} + +// TODO(bt,2025-04-22): Share API definitions between client and server in common crate. + +#[derive(Deserialize, Debug)] +pub struct UploadImageMediaFileSuccessResponse { + pub success: bool, + pub media_file_token: MediaFileToken, +} + +fn validate_args>( + args: &LegacyUploadMediaFileFromFileArgs<'_, P>, +) -> Result<(), StorytellerError> { + // if args.is_intermediate_system_file && args.maybe_prompt_token.is_some() { + // return Err(StorytellerError::Client(ClientError::InvalidPreflightRequest( + // "Cannot set `is_intermediate_system_file` to true if `maybe_prompt_token` is provided." + // .to_string()))); + // } + + Ok(()) +} diff --git a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs index d7cf4c1378..99cc13aa5f 100644 --- a/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs +++ b/crates/api_clients/artcraft/storyteller_client/src/endpoints/media_files/mod.rs @@ -5,3 +5,4 @@ pub mod upload_image_media_file_from_bytes; pub mod upload_image_media_file_from_file; pub mod upload_new_engine_asset_from_file; pub mod upload_video_media_file_from_file; +pub mod legacy_upload_media_file_from_file; diff --git a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs index 2fc3c49cf0..bbdb7b3f06 100644 --- a/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs +++ b/crates/desktop/artcraft/src/core/lifecycle/startup/handle_tauri_startup.rs @@ -18,10 +18,11 @@ 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}; -use crate::services::worldlabs::state::worldlabs_bearer_bridge::WorldlabsBearerBridge; pub async fn handle_tauri_startup( app: AppHandle, @@ -103,6 +104,15 @@ pub async fn handle_tauri_startup( storyteller_creds_manager.clone(), )); + tauri::async_runtime::spawn(worldlabs_marble_task_polling( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + worldlabs_creds_manager.clone(), + storyteller_creds_manager.clone(), + )); + spawn_discord_presence_thread()?; initially_size_and_position_windows(&app, &root); diff --git a/crates/desktop/artcraft/src/services/worldlabs/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/mod.rs index 132c03a929..2393fb390d 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/mod.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/mod.rs @@ -1,3 +1,4 @@ pub mod commands; pub mod state; +pub mod threads; pub mod windows; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs index 2ab7e3d246..8f23128774 100644 --- a/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs +++ b/crates/desktop/artcraft/src/services/worldlabs/state/worldlabs_credential_manager.rs @@ -100,6 +100,13 @@ impl WorldlabsCredentialManager { } } + pub fn maybe_copy_refresh_token(&self) -> anyhow::Result> { + match self.credential_data.read() { + Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)), + Ok(holder) => Ok(holder.world_labs_refresh_token.clone()), + } + } + pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> { match self.credential_data.write() { Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)), diff --git a/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs b/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs new file mode 100644 index 0000000000..e03acdf874 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/threads/mod.rs @@ -0,0 +1 @@ +pub mod worldlabs_marble_task_polling; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs new file mode 100644 index 0000000000..ef24dde5f0 --- /dev/null +++ b/crates/desktop/artcraft/src/services/worldlabs/threads/worldlabs_marble_task_polling.rs @@ -0,0 +1,386 @@ +use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason}; +use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; +use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; +use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent; +use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; +use crate::core::state::data_dir::app_data_root::AppDataRoot; +use crate::core::state::data_dir::trait_data_subdir::DataSubdir; +use crate::core::state::task_database::TaskDatabase; +use crate::core::utils::task_database_pending_statuses::TASK_DATABASE_PENDING_STATUSES; +use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::util::get_or_upgrade_grok_full_credentials::get_or_update_grok_full_credentials; +use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; +use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use crate::services::worldlabs::state::worldlabs_credential_manager::WorldlabsCredentialManager; +use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; +use artcraft_api_defs::utils::media_links_to_thumbnail_template::media_links_to_thumbnail_template; +use cookie_store::cookie_store::CookieStore; +use enums::by_table::prompts::prompt_type::PromptType; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::tauri::tasks::task_media_file_class::TaskMediaFileClass; +use enums::tauri::tasks::task_status::TaskStatus; +use errors::AnyhowResult; +use futures::poll; +use grok_client::credentials::grok_full_credentials::GrokFullCredentials; +use grok_client::error::grok_error::GrokError; +use grok_client::requests::download_video_file::download_video_file::{download_video_file, DownloadVideoFileArgs}; +use grok_client::requests::download_video_file::grok_download_video::GrokDownloadVideo; +use grok_client::requests::media_posts::list_media_posts::grok_list_media_posts::{GrokMediaPostList, GrokMediaPostListRequest, VideoData}; +use idempotency::uuid::generate_random_uuid; +use log::{error, info}; +use midjourney_client::client::midjourney_hostname::MidjourneyHostname; +use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId; +use midjourney_client::endpoints::imagine::{imagine, ImagineItem, ImagineRequest, MidjourneyJobType}; +use midjourney_client::utils::get_image_url::get_image_url; +use midjourney_client::utils::image_downloader_client::ImageDownloaderClient; +use once_cell::sync::Lazy; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs, TaskList}; +use sqlite_tasks::queries::task::Task; +use sqlite_tasks::queries::update_successful_task_status_with_metadata::{update_successful_task_status_with_metadata, UpdateSuccessfulTaskArgs}; +use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs}; +use std::collections::{HashMap, HashSet}; +use std::fs::File; +use std::io::Write; +use std::path::PathBuf; +use std::time::Duration; +use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet; +use storyteller_client::endpoints::media_files::get_media_file::get_media_file; +use storyteller_client::endpoints::media_files::legacy_upload_media_file_from_file::{legacy_upload_media_file_from_file, LegacyUploadMediaFileFromFileArgs}; +use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs}; +use storyteller_client::endpoints::media_files::upload_video_media_file_from_file::{upload_video_media_file_from_file, UploadVideoFromFileArgs}; +use storyteller_client::endpoints::prompts::create_prompt::create_prompt; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use storyteller_client::recipes::upload_media_file_from_file::upload_media_file_from_file; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; +use url::Url; +use url_utils::extension::extract_download_extension_from_url::extract_download_extension_from_url_str; +use world_labs_client::api::api_types::world_id::WorldObjectId; +use world_labs_client::api::requests::worlds::poll_world_status::{poll_world_status, PollWorldStatusArgs}; +use world_labs_client::credentials::world_labs_bearer_token::WorldLabsBearerToken; +use world_labs_client::credentials::world_labs_cookies::WorldLabsCookies; +use world_labs_client::credentials::worldlabs_refresh_token::WorldLabsRefreshToken; + +pub async fn worldlabs_marble_task_polling( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + creds: WorldlabsCredentialManager, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = polling_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &creds, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(Duration::from_millis(30_000)).await; + } +} + +async fn polling_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + worldlabs_creds: &WorldlabsCredentialManager, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + loop { + if !worldlabs_creds.do_task_polling()? { + tokio::time::sleep(Duration::from_millis(10_000)).await; + continue; + } + + // TODO: Graceful wait, fix this long function body + let storyteller_creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + error!("No Storyteller credentials found. Cannot proceed with WorldLabs polling."); + tokio::time::sleep(Duration::from_millis(5_000)).await; + continue; + } + }; + + let world_labs_cookies = match worldlabs_creds.maybe_copy_typed_cookies()? { + Some(cookies) => cookies, + None => { + info!("No full WorldLabs cookies"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let world_labs_bearer = match worldlabs_creds.maybe_copy_bearer_token()? { + Some(bearer) => bearer, + None => { + info!("No full WorldLabs bearer"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let world_labs_refresh = match worldlabs_creds.maybe_copy_refresh_token()? { + Some(bearer) => bearer, + None => { + info!("No full WorldLabs refresh"); + tokio::time::sleep(Duration::from_millis(30_000)).await; + continue; + } + }; + + let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { + db: task_database.get_connection(), + provider: GenerationProvider::WorldLabs, + task_statuses: &TASK_DATABASE_PENDING_STATUSES, + }).await?; + + poll_grok_tasks( + app_handle, + app_env_configs, + app_data_root, + task_database, + &world_labs_cookies, + &world_labs_bearer, + &world_labs_refresh, + &storyteller_creds, + local_tasks, + ).await?; + + tokio::time::sleep(Duration::from_millis(2_000)).await; + } +} + +async fn poll_grok_tasks( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + world_labs_cookies: &WorldLabsCookies, + world_labs_bearer: &WorldLabsBearerToken, + world_labs_refresh: &WorldLabsRefreshToken, + storyteller_creds: &StorytellerCredentialSet, + local_tasks: TaskList, +) -> AnyhowResult<()> { + let local_tasks = local_tasks.tasks; + + if local_tasks.is_empty() { + return Ok(()) + } + + info!("WorldLabs tasks waiting: {:?}", local_tasks.len()); + + // Map of WorldLabs World ID to Local Task. + let local_tasks_by_world_labs_world_id = local_tasks.iter() + .filter_map(|task| { + if let Some(provider_job_id) = &task.provider_job_id { + Some((provider_job_id.clone(), task.clone())) + } else { + None + } + }) + .collect::>(); + + + for (world_id, local_task) in local_tasks_by_world_labs_world_id.iter() { + let world_id = WorldObjectId(world_id.to_string()); + + let poll_world_response = poll_world_status(PollWorldStatusArgs { + cookies: &world_labs_cookies, + bearer_token: &world_labs_bearer, + world_id: &world_id, + request_timeout: None, + }).await?; + + if !poll_world_response.is_complete { + tokio::time::sleep(Duration::from_millis(2_000)).await; + continue; + } + + let spz_url = match poll_world_response.spz_splat_url { + Some(url) => url, + None => { + error!("No spz splat URL despite being marked complete"); + continue; + } + }; + + + upload_spz_splat( + &app_handle, + &app_env_configs, + app_data_root, + task_database, + &storyteller_creds, + &local_task, + &spz_url, + ).await?; + } + + + tokio::time::sleep(Duration::from_millis(5_000)).await; + + Ok(()) +} + +async fn upload_spz_splat( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds: &StorytellerCredentialSet, + local_task: &Task, + spz_url: &str, +) -> AnyhowResult<()> { + + let mut maybe_primary_media_file_token = None; + + info!("Downloading generated spz splat ..."); + + let spz_download_filename = download_spz(spz_url, app_data_root).await?; + + let mut wait_delay = 0; + + loop { + info!("Uploading to backend..."); + + // TODO: media_files.origin_category + // TODO: media_files.maybe_prompt_token + // TODO: media_files.maybe_generation_provider + // TODO: media_files.maybe_origin_model_type + // TODO: media_files.maybe_origin_model_token (sref?) + // TODO: media_files.is_batch_generated + // TODO: media_files.maybe_batch_token + // TODO: media_files.is_user_upload + + // TODO: batch_generations.token + // TODO: batch_generations.entity_type + // TODO: batch_generations.entity_token + + let result = legacy_upload_media_file_from_file(LegacyUploadMediaFileFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: &spz_download_filename, + }).await; + + match result { + Ok(result) => { + info!("Successfully uploaded to backend: {:?}", result.media_file_token); + if maybe_primary_media_file_token.is_none() { + maybe_primary_media_file_token = Some(result.media_file_token); + } + break; + }, + Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => { + error!("Too many requests, retrying upload after delay..."); + // If we hit a rate limit, we can retry after a short delay. + wait_delay += 10; + if wait_delay > 60 { + wait_delay = 60; + } + tokio::time::sleep(Duration::from_secs(wait_delay)).await; + continue; // Retry the upload. + } + Err(err) => { + error!("Failed to upload to backend: {:?}", err); + return Err(err.into()) + }, + } + } // End loop + + let mut maybe_cdn_url = None; + let mut maybe_thumbnail_url_template = None; + + if let Some(media_file_token) = maybe_primary_media_file_token.as_ref() { + info!("Looking up file to grab CDN and thumbnail URLs: {:?} ...", media_file_token); + + let lookup_result = get_media_file( + &app_env_configs.storyteller_host, + media_file_token, + ).await; + match lookup_result { + Ok(response) => { + maybe_cdn_url = Some(response.media_file.media_links.cdn_url.to_string()); + maybe_thumbnail_url_template = media_links_to_thumbnail_template(&response.media_file.media_links) + .map(|s| s.to_string()); + } + Err(err) => { + error!("Failed to look up media file after upload: {:?} (failing open)", err); + } + } + } + + let updated = update_successful_task_status_with_metadata(UpdateSuccessfulTaskArgs { + db: task_database.get_connection(), + task_id: &local_task.id, + maybe_batch_token: None, + maybe_primary_media_file_token: maybe_primary_media_file_token.as_ref(), + maybe_primary_media_file_class: Some(TaskMediaFileClass::Dimensional), + maybe_primary_media_file_thumbnail_url_template: maybe_thumbnail_url_template.as_deref(), + maybe_primary_media_file_cdn_url: maybe_cdn_url.as_deref(), + }).await?; + + if !updated { + return Ok(()); // If anything breaks with queries, don't spam events. + } + + let event = GenerationCompleteEvent { + //media_file_token: result.media_file_token, + action: Some(GenerationAction::GenerateGaussian), + service: GenerationServiceProvider::WorldLabs, + model: None, + }; + + if let Err(err) = event.send(&app_handle) { + error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open + } + + //let result = maybe_handle_text_to_image_complete_event( + // app_handle, + // app_env_configs, + // Some(storyteller_creds), + // local_task, + // &batch_token, + //).await; + + //if let Err(err) = result { + // error!("Failed to send text-to-image complete event: {:?}", err); + //} + + Ok(()) +} + +async fn download_spz( + spz_url: &str, + app_data_root: &AppDataRoot, +) -> AnyhowResult { + info!("Downloading splat from URL: {}", spz_url); + + let response = reqwest::get(spz_url).await?; + let image_bytes = response.bytes().await?; + + let uuid = generate_random_uuid(); + let extension_without_period = "ceramic.spz"; + + let tempdir = app_data_root.temp_dir().path(); + let download_filename = format!("{}.{}", uuid, extension_without_period); + let download_path = tempdir.join(download_filename); + + info!("Writing to path: {:?}", download_path); + + let mut file = File::create(&download_path)?; + file.write_all(&image_bytes)?; + + Ok(download_path) +}