From 0388c3bbb8118e1a8e6edfa09efdaf91912fa8ef Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Thu, 14 Aug 2025 23:56:41 -0400 Subject: [PATCH] midjourney image upload works! --- .../src/error/midjourney_api_error.rs | 9 +++ .../src/utils/image_downloader_client.rs | 70 +++++++++++++++++++ .../midjourney_client/src/utils/mod.rs | 3 +- .../threads/midjourney_long_polling_thread.rs | 20 +++++- .../utils/download_midjourney_image.rs | 14 +++- .../state/storyteller_credential_holder.rs | 2 +- 6 files changed, 111 insertions(+), 7 deletions(-) create mode 100644 crates/api_clients/midjourney_client/src/utils/image_downloader_client.rs diff --git a/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs b/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs index 2c417d5712..45b5de4d9a 100644 --- a/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs +++ b/crates/api_clients/midjourney_client/src/error/midjourney_api_error.rs @@ -41,6 +41,12 @@ pub enum MidjourneyApiError { backend_hostname: Option, }, + /// Eg. when downloading images + UnknownHttpFailure { + status_code: u16, + body: String, + }, + /// Cloudflare errors. CloudflareError(CloudflareError), @@ -81,6 +87,9 @@ impl Display for MidjourneyApiError { Self::TooManyRequests(msg) => write!(f, "Too many requests: {}", msg), Self::InternalServerError {body, backend_hostname} => write!(f, "Internal Server Error; backend hostname: {:?} ; body: {}; ", backend_hostname, body), + Self::UnknownHttpFailure {status_code, body} => + write!(f, "Unknown HTTP failure; status code: {}; body: {}", status_code, body), + // Deserialization errors // Server response handling errors Self::DeserializationError(error) => write!(f, "Deserialization error: {}", error), // Network errors diff --git a/crates/api_clients/midjourney_client/src/utils/image_downloader_client.rs b/crates/api_clients/midjourney_client/src/utils/image_downloader_client.rs new file mode 100644 index 0000000000..c5e3720353 --- /dev/null +++ b/crates/api_clients/midjourney_client/src/utils/image_downloader_client.rs @@ -0,0 +1,70 @@ +use crate::error::midjourney_api_error::MidjourneyApiError; +use crate::error::midjourney_client_error::MidjourneyClientError; +use crate::error::midjourney_error::MidjourneyError; +use crate::utils::get_image_url::get_image_url; +use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors; +use wreq::Client; +use wreq_util::Emulation; + +#[derive(Clone)] +pub struct ImageDownloaderClient { + client: Client, +} + +impl ImageDownloaderClient { + pub fn create() -> Result { + Ok(Self { + client: Client::builder() + .emulation(Emulation::Firefox139) + .build() + .map_err(|err| MidjourneyClientError::WreqError(err))? + }) + } + + pub async fn download_image(&self, job_id: &str, image_index: u8) -> anyhow::Result, MidjourneyError> { + let url = get_image_url(job_id, image_index)?; + + // TODO: Cookies + // TODO: Cache control headers? + let mut http_request = self.client.get(url) + //.header("Referrer-Policy", "origin-when-cross-origin") + //.header("content-type", "application/json") + //.header("sec-fetch-user", "?1") + //.header("upgrade-insecure-requests", "1"); + .header("Referrer", "https://www.midjourney.com/") + .header("accept", "image/avif,image/webp,image/apng,image/svg+xml,image/*,*/*;q=0.8") + .header("accept-language", "en-US,en;q=0.8") + .header("priority", "i") + .header("sec-ch-ua-mobile", "?0") + .header("sec-fetch-dest", "image") + .header("sec-fetch-mode", "no-cors") + .header("sec-fetch-site", "same-site"); + + let http_request = http_request + .build() + .map_err(|err| MidjourneyClientError::WreqError(err))?; + + let response = self.client.execute(http_request) + .await + .map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + let status = response.status(); + + let response_bytes = response.bytes().await + .map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?; + + if !status.is_success() { + let response_body = String::from_utf8_lossy(&response_bytes).to_string(); + if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) { + return Err(MidjourneyApiError::CloudflareError(err).into()); + } + + return Err(MidjourneyApiError::UnknownHttpFailure { + status_code: status.as_u16(), + body: response_body, + }.into()); + } + + Ok(response_bytes.to_vec()) + } +} diff --git a/crates/api_clients/midjourney_client/src/utils/mod.rs b/crates/api_clients/midjourney_client/src/utils/mod.rs index 27d59c832f..9f61b7f98a 100644 --- a/crates/api_clients/midjourney_client/src/utils/mod.rs +++ b/crates/api_clients/midjourney_client/src/utils/mod.rs @@ -1 +1,2 @@ -pub mod get_image_url; \ No newline at end of file +pub mod get_image_url; +pub mod image_downloader_client; \ No newline at end of file diff --git a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs index 9ec26a5e4a..6e20c7fd35 100644 --- a/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs +++ b/crates/desktop/artcraft/src/services/midjourney/threads/midjourney_long_polling_thread.rs @@ -18,6 +18,7 @@ use std::io::Write; use std::path::PathBuf; use tauri::AppHandle; use url::Url; +use midjourney_client::utils::image_downloader_client::ImageDownloaderClient; use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs}; use crate::core::events::basic_sendable_event_trait::BasicSendableEvent; use crate::core::events::generation_events::common::{GenerationAction, GenerationServiceProvider}; @@ -66,7 +67,14 @@ async fn polling_loop( } // TODO: Graceful wait, fix this long function body - let storyteller_creds = storyteller_creds_manager.get_credentials_required()?; + let storyteller_creds = match storyteller_creds_manager.get_credentials()? { + Some(creds) => creds, + None => { + error!("No Storyteller credentials found. Cannot proceed with Midjourney polling."); + tokio::time::sleep(std::time::Duration::from_millis(10_000)).await; + continue; + } + }; let cookies = creds.maybe_copy_cookie_store()?; @@ -150,6 +158,8 @@ async fn polling_loop( // TODO: If we introduce another job polling mechanism, we may need to handle concurrency. + let image_downloader = ImageDownloaderClient::create()?; + for (midjourney_job_id, local_task) in local_tasks_by_midjourney_job_id.iter() { let midjourney_item = match midjourney_items_by_id.get(midjourney_job_id) { Some(item) => item, @@ -158,7 +168,13 @@ async fn polling_loop( for index in 0..4 { info!("Downloading generated Midjourney file..."); - let download_path = download_midjourney_image(midjourney_job_id, index, app_data_root).await?; + + let download_path = download_midjourney_image( + &image_downloader, + midjourney_job_id, + index, + app_data_root + ).await?; info!("Uploading to backend..."); diff --git a/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs b/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs index 6bcebd851e..fb63021c15 100644 --- a/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs +++ b/crates/desktop/artcraft/src/services/midjourney/utils/download_midjourney_image.rs @@ -6,13 +6,21 @@ use std::fs::File; use std::io::Write; use std::path::PathBuf; use url::Url; +use midjourney_client::utils::image_downloader_client::ImageDownloaderClient; -pub async fn download_midjourney_image(job_id: &str, index: u8, app_data_root: &AppDataRoot) -> AnyhowResult { +pub async fn download_midjourney_image( + image_downloader: &ImageDownloaderClient, + job_id: &str, + index: u8, + app_data_root: &AppDataRoot +) -> AnyhowResult { let url = get_image_url(job_id, index)?; let parsed_url = Url::parse(&url)?; - let response = reqwest::get(&url).await?; - let image_bytes = response.bytes().await?; + //let response = reqwest::get(&url).await?; + //let image_bytes = response.bytes().await?; + + let image_bytes = image_downloader.download_image(job_id, index).await?; let ext = parsed_url.path().split(".").last().unwrap_or("png"); diff --git a/crates/desktop/artcraft/src/services/storyteller/state/storyteller_credential_holder.rs b/crates/desktop/artcraft/src/services/storyteller/state/storyteller_credential_holder.rs index 8add3f37e3..26f98017e3 100644 --- a/crates/desktop/artcraft/src/services/storyteller/state/storyteller_credential_holder.rs +++ b/crates/desktop/artcraft/src/services/storyteller/state/storyteller_credential_holder.rs @@ -49,7 +49,7 @@ impl StorytellerCredentialHolder { match self.credentials.read() { Err(err) => Err(anyhow!("Failed to acquire read lock: {:?}", err)), Ok(creds) => match &*creds { - None => Err(anyhow!("Credentials not set")), + None => Err(anyhow!("Storyteller credentials not set")), Some(creds) => Ok(creds.clone()), } }