From 31b56fc910a31439dce8b2e8a55d329eea597732 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 31 Oct 2025 19:42:21 -0400 Subject: [PATCH] wip refactors --- .../artcraft/src/core/artcraft_error.rs | 2 + .../enqueue_text_to_image_command.rs | 7 + .../enqueue/text_to_image/grok/handle_grok.rs | 60 ++--- .../lifecycle/startup/handle_tauri_startup.rs | 13 + crates/desktop/artcraft/src/lib.rs | 6 + .../grok/state/grok_image_prompt_queue.rs | 66 +++++ .../artcraft/src/services/grok/state/mod.rs | 1 + .../grok_image_websocket_thread.rs | 234 ++++++++++++++++++ .../grok_image_websocket_thread/mod.rs | 1 + .../artcraft/src/services/grok/threads/mod.rs | 3 +- 10 files changed, 352 insertions(+), 41 deletions(-) create mode 100644 crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs create mode 100644 crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs create mode 100644 crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs diff --git a/crates/desktop/artcraft/src/core/artcraft_error.rs b/crates/desktop/artcraft/src/core/artcraft_error.rs index 293b2b6169..72ef8847b2 100644 --- a/crates/desktop/artcraft/src/core/artcraft_error.rs +++ b/crates/desktop/artcraft/src/core/artcraft_error.rs @@ -14,6 +14,7 @@ pub enum ArtcraftError { // Lock errors RwLockReadError, RwLockWriteError, + MutexLockError, } impl Error for ArtcraftError {} @@ -28,6 +29,7 @@ impl Display for ArtcraftError { Self::StorytellerError(e) => write!(f, "StorytellerError: {:?}", e), Self::RwLockReadError => write!(f, "RwLockReadError"), Self::RwLockWriteError => write!(f, "RwLockWriteError"), + Self::MutexLockError => write!(f, "MutexLockError"), } } } diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs index 58c1a324a0..c19a7b4f23 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/enqueue_text_to_image_command.rs @@ -21,6 +21,7 @@ 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::grok::state::grok_websocket_manager::GrokWebsocketManager; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::sora::state::sora_credential_manager::SoraCredentialManager; @@ -151,6 +152,7 @@ pub async fn enqueue_text_to_image_command( mj_creds_manager: State<'_, MidjourneyCredentialManager>, grok_creds_manager: State<'_, GrokCredentialManager>, grok_websocket_manager: State<'_, GrokWebsocketManager>, + grok_image_prompt_queue: State<'_, GrokImagePromptQueue>, storyteller_creds_manager: State<'_, StorytellerCredentialManager>, sora_creds_manager: State<'_, SoraCredentialManager>, sora_task_queue: State<'_, SoraTaskQueue>, @@ -167,6 +169,7 @@ pub async fn enqueue_text_to_image_command( &mj_creds_manager, &grok_creds_manager, &grok_websocket_manager, + &grok_image_prompt_queue, &storyteller_creds_manager, &app_env_configs, &sora_creds_manager, @@ -237,6 +240,7 @@ pub async fn handle_request( mj_creds_manager: &MidjourneyCredentialManager, grok_creds_manager: &GrokCredentialManager, grok_websocket_manager: &GrokWebsocketManager, + grok_image_prompt_queue: &GrokImagePromptQueue, storyteller_creds_manager: &StorytellerCredentialManager, app_env_configs: &AppEnvConfigs, sora_creds_manager: &SoraCredentialManager, @@ -253,6 +257,7 @@ pub async fn handle_request( &mj_creds_manager, &grok_creds_manager, &grok_websocket_manager, + &grok_image_prompt_queue, &sora_creds_manager, &sora_task_queue, ).await; @@ -289,6 +294,7 @@ pub async fn dispatch_request( mj_creds_manager: &MidjourneyCredentialManager, grok_creds_manager: &GrokCredentialManager, grok_websocket_manager: &GrokWebsocketManager, + grok_image_prompt_queue: &GrokImagePromptQueue, sora_creds_manager: &SoraCredentialManager, sora_task_queue: &SoraTaskQueue, ) -> Result { @@ -328,6 +334,7 @@ pub async fn dispatch_request( app_env_configs, grok_creds_manager, grok_websocket_manager, + grok_image_prompt_queue, ).await; } Some(ImageModel::Midjourney) => { diff --git a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs index 3ca29d56f6..58dd1ff4d7 100644 --- a/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs +++ b/crates/desktop/artcraft/src/core/commands/enqueue/text_to_image/grok/handle_grok.rs @@ -1,44 +1,15 @@ -use crate::core::commands::enqueue::generate_error::{GenerateError, MissingCredentialsReason, ProviderFailureReason}; -use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize}; +use crate::core::commands::enqueue::generate_error::GenerateError; 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::basic_sendable_event_trait::BasicSendableEvent; -use crate::core::events::functional_events::canvas_background_removal_complete_event::CanvasBackgroundRemovalCompleteEvent; -use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent; -use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider}; -use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent; -use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent; +use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest; +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::provider_priority::ProviderPriorityStore; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; -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 artcraft_api_defs::generate::image::generate_gpt_image_1_text_to_image::{GenerateGptImage1TextToImageImageQuality, GenerateGptImage1TextToImageImageSize, GenerateGptImage1TextToImageNumImages, GenerateGptImage1TextToImageRequest}; use enums::common::generation_provider::GenerationProvider; use enums::tauri::tasks::task_type::TaskType; -use grok_client::error::grok_error::GrokError; -use grok_client::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; -use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; -use grok_client::requests::image_websocket::listen_for_websocket_request_id::{listen_for_websocket_request_id, ListenForWebsocketRequestIdArgs}; -use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; use idempotency::uuid::generate_random_uuid; -use log::{error, info, warn}; -use midjourney_client::client::midjourney_hostname::MidjourneyHostname; -use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest}; -use midjourney_client::error::midjourney_api_error::MidjourneyApiError; -use midjourney_client::recipes::channel_id::ChannelId; -use midjourney_client::recipes::text_to_image::{text_to_image, TextToImageError, TextToImageRequest}; -use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session; -use openai_sora_client::recipes::simple_image_gen_with_session_auto_renew::{simple_image_gen_with_session_auto_renew, SimpleImageGenAutoRenewRequest}; -use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages}; -use std::time::Duration; -use storyteller_client::endpoints::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image; -use storyteller_client::endpoints::generate::image::generate_gpt_image_1_text_to_image::generate_gpt_image_1_text_to_image; use tauri::AppHandle; -use tokens::tokens::media_files::MediaFileToken; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; pub async fn handle_grok( app: &AppHandle, @@ -46,9 +17,9 @@ pub async fn handle_grok( app_env_configs: &AppEnvConfigs, creds_manager: &GrokCredentialManager, websocket_manager: &GrokWebsocketManager, + grok_image_prompt_queue: &GrokImagePromptQueue, ) -> Result { - unimplemented!() //// TODO: We can request population of the user info if absent or expired. //let websocket = get_websocket( @@ -76,13 +47,22 @@ pub async fn handle_grok( //}; // //info!("Successfully enqueued MidJourney. Job token: {}", request_id); + + let job_id = generate_random_uuid(); + + let prompt = request.prompt + .as_deref() + .map(|prompt| prompt.trim().to_string()) + .unwrap_or_else(|| "".to_string()); + + grok_image_prompt_queue.enqueue(&prompt)?; - //Ok(TaskEnqueueSuccess { - // provider: GenerationProvider::Grok, - // model: Some(GenerationModel::GrokImage), - // provider_job_id: Some(request_id), - // task_type: TaskType::ImageGeneration, - //}) + Ok(TaskEnqueueSuccess { + provider: GenerationProvider::Grok, + model: Some(GenerationModel::GrokImage), + provider_job_id: Some(job_id), + task_type: TaskType::ImageGeneration, + }) } // TODO: 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 5d68d6d878..791b7ed6b1 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 @@ -10,7 +10,9 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs; use crate::core::state::artcraft_platform_info::ArtcraftPlatformInfo; use crate::core::state::data_dir::app_data_root::AppDataRoot; use crate::services::grok::state::grok_credential_manager::GrokCredentialManager; +use crate::services::grok::state::grok_image_prompt_queue::GrokImagePromptQueue; use crate::services::grok::state::grok_websocket_manager::GrokWebsocketManager; +use crate::services::grok::threads::grok_image_websocket_thread::grok_image_websocket_thread::grok_image_websocket_thread; use crate::services::grok::threads::grok_video_task_polling::grok_video_task_polling_thread::grok_video_task_polling_thread; use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager; use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread; @@ -31,6 +33,7 @@ pub async fn handle_tauri_startup( mj_creds_manager: MidjourneyCredentialManager, grok_creds_manager: GrokCredentialManager, grok_websocket_manager: GrokWebsocketManager, + grok_image_prompt_queue: GrokImagePromptQueue, ) -> AnyhowResult<()> { set_app_log_level( @@ -79,6 +82,16 @@ pub async fn handle_tauri_startup( storyteller_creds_manager.clone(), )); + tauri::async_runtime::spawn(grok_image_websocket_thread( + app.clone(), + app_env_configs.clone(), + root.clone(), + task_database.clone(), + grok_creds_manager.clone(), + grok_image_prompt_queue.clone(), + storyteller_creds_manager.clone(), + )); + tauri::async_runtime::spawn(midjourney_long_polling_thread( app.clone(), app_env_configs.clone(), diff --git a/crates/desktop/artcraft/src/lib.rs b/crates/desktop/artcraft/src/lib.rs index 1cd7a5ef4d..21b07b501c 100644 --- a/crates/desktop/artcraft/src/lib.rs +++ b/crates/desktop/artcraft/src/lib.rs @@ -32,6 +32,7 @@ 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_websocket_manager::GrokWebsocketManager; use crate::services::midjourney::commands::midjourney_clear_credentials_command::midjourney_clear_credentials_command; @@ -108,6 +109,9 @@ pub fn run() { let grok_websocket_manager = GrokWebsocketManager::new(); let grok_websocket_manager_2= grok_websocket_manager.clone(); + let grok_prompt_queue = GrokImagePromptQueue::new(); + let grok_prompt_queue_2 = grok_image_prompt_queue.clone(); + println!("Initializing backend runtime..."); let builder = tauri::Builder::default() @@ -147,6 +151,7 @@ pub fn run() { midjourney_creds_manager_2, grok_creds_manager_2, grok_websocket_manager_2, + grok_prompt_queue_2, ).await; if let Err(err) = result { @@ -163,6 +168,7 @@ pub fn run() { .manage(artcraft_platform_info) .manage(grok_creds_manager) .manage(grok_websocket_manager) + .manage(grok_prompt_queue) .manage(midjourney_creds_manager) .manage(sora_creds_manager) .manage(sora_task_queue) diff --git a/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs new file mode 100644 index 0000000000..2ffc071b0c --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/state/grok_image_prompt_queue.rs @@ -0,0 +1,66 @@ +use crate::core::artcraft_error::ArtcraftError; +use log::error; +use std::collections::VecDeque; +use std::sync::{Arc, Mutex}; + +#[derive(Clone)] +pub struct GrokImagePromptQueue { + pub prompt_queue: Arc>>, +} + +impl GrokImagePromptQueue { + pub fn new() -> Self { + Self { + prompt_queue: Arc::new(Mutex::new(VecDeque::new())), + } + } + + pub fn enqueue(&self, prompt: &str) -> Result<(), ArtcraftError> { + match self.prompt_queue.lock() { + Ok(mut queue) => { + queue.push_back(prompt.to_string()); + Ok(()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn dequeue(&self) -> Result, ArtcraftError> { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.pop_front()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn is_empty(&self) -> Result { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.is_empty()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } + + pub fn len(&self) -> Result { + match self.prompt_queue.lock() { + Ok(mut queue) => { + Ok(queue.len()) + }, + Err(err) => { + error!("Error locking prompt queue: {:?}", err); + Err(ArtcraftError::MutexLockError) + }, + } + } +} diff --git a/crates/desktop/artcraft/src/services/grok/state/mod.rs b/crates/desktop/artcraft/src/services/grok/state/mod.rs index f467802043..5b557e4567 100644 --- a/crates/desktop/artcraft/src/services/grok/state/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/state/mod.rs @@ -1,4 +1,5 @@ pub mod grok_credential_manager; +pub mod grok_image_prompt_queue; pub mod grok_serializable_state; pub mod grok_websocket_manager; diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs new file mode 100644 index 0000000000..5c5a7f4f60 --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/grok_image_websocket_thread.rs @@ -0,0 +1,234 @@ +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::task_database::TaskDatabase; +use crate::core::utils::download_url_to_temp_dir::download_url_to_temp_dir; +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::state::grok_image_prompt_queue::GrokImagePromptQueue; +use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager; +use artcraft_api_defs::prompts::create_prompt::CreatePromptRequest; +use enums::common::generation_provider::GenerationProvider; +use enums::common::model_type::ModelType; +use enums::tauri::tasks::task_model_type::TaskModelType; +use errors::AnyhowResult; +use grok_client::requests::image_websocket::create_listen_websocket::{create_listen_websocket, CreateListenWebsocketArgs}; +use grok_client::requests::image_websocket::grok_websocket::GrokWebsocket; +use grok_client::requests::image_websocket::listen_for_websocket_images::{listen_for_websocket_images, ImageResults, ListenForWebsocketImagesArgs}; +use grok_client::requests::image_websocket::prompt_websocket_image::{prompt_websocket_image, PromptWebsocketImageArgs}; +use idempotency::uuid::generate_random_uuid; +use log::{error, info}; +use sqlite_tasks::queries::list_tasks_by_provider_and_status::{list_tasks_by_provider_and_status, ListTasksByProviderAndStatusArgs}; +use std::time::Duration; +use storyteller_client::endpoints::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs, UploadImageMediaFileSuccessResponse}; +use storyteller_client::endpoints::prompts::create_prompt::create_prompt; +use storyteller_client::error::api_error::ApiError; +use storyteller_client::error::storyteller_error::StorytellerError; +use tauri::AppHandle; +use tokens::tokens::batch_generations::BatchGenerationToken; + +pub async fn grok_image_websocket_thread( + app_handle: AppHandle, + app_env_configs: AppEnvConfigs, + app_data_root: AppDataRoot, + task_database: TaskDatabase, + creds: GrokCredentialManager, + prompt_queue: GrokImagePromptQueue, + storyteller_creds_manager: StorytellerCredentialManager, +) -> ! { + loop { + let res = inner_loop( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &creds, + &prompt_queue, + &storyteller_creds_manager, + ).await; + if let Err(err) = res { + error!("An error occurred: {:?}", err); + } + // NB: Only sleep if an error occurs. + tokio::time::sleep(std::time::Duration::from_millis(30_000)).await; + } +} + +async fn inner_loop( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + grok_creds: &GrokCredentialManager, + prompt_queue: &GrokImagePromptQueue, + storyteller_creds_manager: &StorytellerCredentialManager, +) -> AnyhowResult<()> { + + loop { + if !grok_creds.do_task_polling()? { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + + let maybe_cookie_header = grok_creds.maybe_copy_cookie_header_string()?; + + let cookie_header = match maybe_cookie_header { + Some(cookie_header) => cookie_header, + None => { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + }; + + let websocket = create_listen_websocket(CreateListenWebsocketArgs { + cookies: &cookie_header, + }).await?; + + let websocket = GrokWebsocket::new(websocket); + + //poll_task_loop( + // &app_handle, + // &app_env_configs, + // &app_data_root, + // &task_database, + // &grok_creds, + // websocket, + // prompt_queue, + // &storyteller_creds_manager, + //).await?; + + loop { + // let local_tasks = list_tasks_by_provider_and_status(ListTasksByProviderAndStatusArgs { + // db: task_database.get_connection(), + // provider: GenerationProvider::Grok, + // task_statuses: &TASK_DATABASE_PENDING_STATUSES, + // }).await?; + // // Only Grok images + // let local_tasks = local_tasks.tasks.into_iter() + // .filter(|task| match task.model_type { + // Some(TaskModelType::GrokImage) => true, + // _ => false, + // }) + // .collect::>(); + + let maybe_prompt = prompt_queue.dequeue()?; + + let prompt = match maybe_prompt { + Some(prompt) => prompt, + None => { + tokio::time::sleep(std::time::Duration::from_millis(1_000)).await; + continue; + } + }; + + info!("Prompting Grok websocket: {}", prompt); + + let _result = prompt_websocket_image(PromptWebsocketImageArgs { + websocket_wrapped: &websocket, + prompt: &prompt, + }).await?; + + let images = listen_for_websocket_images(ListenForWebsocketImagesArgs { + websocket: &websocket, + timeout: Duration::from_millis(30_000), + }).await?; + + upload_images_to_storyteller( + &app_handle, + &app_env_configs, + &app_data_root, + &task_database, + &storyteller_creds_manager, + &prompt, + &images, + ).await? + } + + } +} + +//async fn poll_task_loop( +// app_handle: &AppHandle, +// app_env_configs: &AppEnvConfigs, +// app_data_root: &AppDataRoot, +// task_database: &TaskDatabase, +// grok_creds: &GrokCredentialManager, +// grok_websocket: GrokWebsocket, +// prompt_queue: &GrokImagePromptQueue, +// storyteller_creds_manager: &StorytellerCredentialManager, +//) -> AnyhowResult<()> { +// +//} + +async fn upload_images_to_storyteller( + app_handle: &AppHandle, + app_env_configs: &AppEnvConfigs, + app_data_root: &AppDataRoot, + task_database: &TaskDatabase, + storyteller_creds_manager: &StorytellerCredentialManager, + prompt: &str, + images: &ImageResults, +) -> AnyhowResult<()> { + + loop { + let storyteller_creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + error!("No Storyteller credentials found. Cannot proceed with Grok polling."); + tokio::time::sleep(std::time::Duration::from_millis(5_000)).await; + continue; + } + }; + + let request = CreatePromptRequest { + uuid_idempotency_token: generate_random_uuid(), + positive_prompt: Some(prompt.trim().to_string()), + negative_prompt: None, + model_type: Some(ModelType::GrokImage), + generation_provider: Some(GenerationProvider::Grok), + }; + + let prompt_response = create_prompt( + &app_env_configs.storyteller_host, + Some(&storyteller_creds), + request + ).await?; + + info!("Created prompt: {:?}", &prompt_response.prompt_token); + + // TODO: Move this from clientside to the backend. + // The first upload should produce a batch token that we can reuse. + let batch_token = BatchGenerationToken::generate(); + + for image in images.images.iter() { + loop { + let url = image.url.to_string(); + let file = download_url_to_temp_dir(&url, app_data_root).await?; + + let result = upload_image_media_file_from_file(UploadImageFromFileArgs { + api_host: &app_env_configs.storyteller_host, + maybe_creds: Some(&storyteller_creds), + path: file.path(), + is_intermediate_system_file: false, + maybe_prompt_token: Some(&prompt_response.prompt_token), + maybe_batch_token: Some(&batch_token), + }).await; + + match result { + Ok(_) => { + break; + } + Err(StorytellerError::Api(ApiError::TooManyRequests(_))) => { + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + Err(err) => { + error!("Error uploading image: {}", err); + break; + } + } + } // Retry on 429s, etc. + } + + return Ok(()) + } +} diff --git a/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs new file mode 100644 index 0000000000..5021bcae7c --- /dev/null +++ b/crates/desktop/artcraft/src/services/grok/threads/grok_image_websocket_thread/mod.rs @@ -0,0 +1 @@ +pub mod grok_image_websocket_thread; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/grok/threads/mod.rs b/crates/desktop/artcraft/src/services/grok/threads/mod.rs index 5239c8592a..07dab82921 100644 --- a/crates/desktop/artcraft/src/services/grok/threads/mod.rs +++ b/crates/desktop/artcraft/src/services/grok/threads/mod.rs @@ -1 +1,2 @@ -pub mod grok_video_task_polling; \ No newline at end of file +pub mod grok_video_task_polling; +pub mod grok_image_websocket_thread; \ No newline at end of file