Compare commits

..
Author SHA1 Message Date
bflatastic 58accbcf56 repositioned outliner above model selector and copy paste fixed? 2025-07-31 20:44:38 +07:00
237 changed files with 1011 additions and 13384 deletions
+3 -14
View File
@@ -26,23 +26,14 @@ jobs:
steps:
- uses: actions/checkout@v4
- name: install Rust stable
uses: dtolnay/rust-toolchain@stable
- name: install wreq dependencies
run: |
choco install cmake -y
choco install strawberryperl -y
choco install pkgconfiglite -y
choco install llvm -y
choco install nasm -y
shell: cmd
- name: setup node
uses: actions/setup-node@v4
with:
node-version: lts/*
- name: install Rust stable
uses: dtolnay/rust-toolchain@stable
- name: install frontend dependencies
run: |
# Navigate to the frontend directory to install dependencies
@@ -92,8 +83,6 @@ jobs:
- name: build and publish
uses: tauri-apps/tauri-action@v0
env:
# NB(bt): This is needed by wreq. Hopefully it doesn't kill Tauri
RUSTFLAGS: "-C target-feature=+crt-static"
GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
# Use the SQLx query cache for SQLite
SQLX_OFFLINE: true
@@ -1,12 +0,0 @@
{
"db_name": "SQLite",
"query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_caller,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 9
},
"nullable": []
},
"hash": "01a0bf33e8331c19dc73f064c54eaa863bec158ef7713450a06c75e246f5f08c"
}
@@ -0,0 +1,12 @@
{
"db_name": "SQLite",
"query": "\n INSERT INTO tasks (\n id,\n task_status,\n task_type,\n model_type,\n provider,\n provider_job_id,\n frontend_subscriber_id,\n frontend_subscriber_payload\n )\n VALUES (?, ?, ?, ?, ?, ?, ?, ?)\n ",
"describe": {
"columns": [],
"parameters": {
"Right": 8
},
"nullable": []
},
"hash": "22b91a21cbb88abc9bc17b338cf2f9418e08e414667d44476623eb141ab17b87"
}
Generated
+107 -579
View File
File diff suppressed because it is too large Load Diff
+1 -7
View File
@@ -10,7 +10,6 @@ members = [
"crates/api_clients/fal_client",
"crates/api_clients/google_drive_common",
"crates/api_clients/google_sign_in",
"crates/api_clients/midjourney_client",
"crates/api_clients/openai_sora_client",
"crates/api_clients/openart_client",
"crates/api_clients/runwayml_client",
@@ -43,7 +42,6 @@ members = [
"crates/lib/caching/redis_caching",
"crates/lib/cloudflare_errors",
"crates/lib/concurrency",
"crates/lib/cookie_store",
"crates/lib/cookies",
"crates/lib/data/collections",
"crates/lib/data/crockford",
@@ -201,8 +199,7 @@ log = "0.4.27"
once_cell = "1.21.3"
rand = "0.9.1"
regex = "1.11.1"
reqwest = { version = "0.12.22", features = ["json", "cookies"] }
reqwest-websocket = { version = "0.5.0" }
reqwest = { version = "0.12.15", features = ["json", "cookies"] }
serde = { version = "1.0.219", features = ["derive"] }
serde_derive = "1.0.219"
serde_json = "1.0.140"
@@ -221,7 +218,6 @@ artcraft_api_defs = { path = "crates/schema/public/api/artcraft_api_defs" }
build_metadata = { path = "crates/lib/build_metadata" }
cloud_storage = { path = "crates/api_clients/cloud_storage" }
cloudflare_errors = { path = "crates/lib/cloudflare_errors" }
cookie_store = { path = "crates/lib/cookie_store" }
crockford = { path = "crates/lib/data/crockford" }
enums = { path = "crates/schema/public/enums" }
enums_public = { path = "crates/schema/public/enums_public" }
@@ -234,7 +230,6 @@ images = { path = "crates/lib/files/images" }
jwt_light = { path = "crates/lib/jwt_light" }
memory_caching = { path = "crates/lib/caching/memory_caching" }
memory_store = { path = "crates/lib/caching/memory_store" }
midjourney_client = { path = "crates/api_clients/midjourney_client" }
mimetypes = { path = "crates/lib/files/mimetypes" }
openai_sora_client = { path = "crates/api_clients/openai_sora_client" }
primitives = { path = "crates/lib/data/primitives" }
@@ -243,6 +238,5 @@ sqlite_tasks = { path = "crates/schema/database/sqlite_tasks" }
storyteller_client = { path = "crates/api_clients/storyteller_client" }
subprocess_common = { path = "crates/lib/subprocess_common" }
test_data = { path = "crates/lib/test_data" }
testing = { path = "crates/lib/testing" }
tokens = { path = "crates/schema/public/tokens" }
web_base64 = { path = "crates/lib/data/web_base64" }
@@ -1,53 +1,20 @@
-- noinspection SqlDialectInspectionForFile
--
-- NB: This version needs to be bumped even if just adding or changing comments!
-- Failure to do so will cause Windows to segfault at start. (This is horrible!)
--
-- NB: `app_state_dir.rs` contains the file version number.
--
-- Migration history:
-- tasks_v1.sqlite - initial version
-- tasks_v2.sqlite - added model_type (nullable)
-- tasks_v3.sqlite - added frontend_caller, comments
CREATE TABLE tasks (
-- Task auto-incrementing primary key.
-- We don't use this.
id TEXT NOT NULL PRIMARY KEY,
-- TaskStatus enum
-- This is how the job system manages job states.
-- e.g. pending, started, complete_success, dead, etc.
task_status TEXT NOT NULL,
-- TaskType enum
-- e.g. image_generation, video_generation, etc.
task_type TEXT NOT NULL,
-- TaskModelType enum
-- e.g. flux_1_dev, veo_2, etc.
model_type TEXT,
-- GenerationProvider enum
-- Together with `provider_job_id`, this creates a foreign key lookup to the job.
-- e.g. 'artcraft', 'fal', 'sora', etc.
provider TEXT,
-- The primary key for the job in the provider's system.
provider_job_id TEXT,
-- OPTIONAL.
-- Tell the job system which caller initiated the task.
-- e.g. 'canvas'
frontend_caller TEXT,
-- OPTIONAL.
-- An arbitrary opaque identifier that frontend can set.
frontend_subscriber_id TEXT,
-- OPTIONAL.
-- An opaque JSON payload set by the frontend subscriber that
-- will be re-emitted back to the frontend.
frontend_subscriber_payload TEXT,
created_at INTEGER NOT NULL DEFAULT (unixepoch('now')),
@@ -1,89 +0,0 @@
use crate::creds::fal_api_key::FalApiKey;
use crate::error::classify_fal_error::classify_fal_error;
use crate::error::fal_error_plus::FalErrorPlus;
use fal::prelude::fal_ai::flux_pro::kontext::max::{max, FluxProKontextMaxInput};
use fal::webhook::WebhookResponse;
use reqwest::IntoUrl;
pub struct FluxProKontextMaxArgs<'a, U: IntoUrl, R: IntoUrl> {
// Request
pub prompt: &'a str,
pub image_url: U,
pub num_images: FluxProKontextMaxNumImages,
// Fulfillment
pub webhook_url: R,
pub api_key: &'a FalApiKey,
}
#[derive(Copy, Clone, Debug)]
pub enum FluxProKontextMaxNumImages {
One, // Default
Two,
Three,
Four,
}
pub async fn enqueue_flux_pro_kontext_max_edit_webhook<U: IntoUrl, R: IntoUrl>(
args: FluxProKontextMaxArgs<'_, U, R>
) -> Result<WebhookResponse, FalErrorPlus> {
let num_images = match args.num_images {
FluxProKontextMaxNumImages::One => 1,
FluxProKontextMaxNumImages::Two => 2,
FluxProKontextMaxNumImages::Three => 3,
FluxProKontextMaxNumImages::Four => 4,
};
let request = FluxProKontextMaxInput {
prompt: args.prompt.to_string(),
image_url: args.image_url.as_str().to_string(),
num_images: Some(num_images),
// Maybe expose
aspect_ratio: None,
safety_tolerance: Some("5".to_string()), // NB: 5 is most tolerant
output_format: Some("png".to_string()), // png or jpeg
seed: None,
// Constants
sync_mode: None, // Synchronous / slow
};
let result = max(request)
.with_api_key(&args.api_key.0)
.queue_webhook(args.webhook_url)
.await;
result.map_err(|err| classify_fal_error(err))
}
#[cfg(test)]
mod tests {
use crate::creds::fal_api_key::FalApiKey;
use crate::requests::webhook::image::edit::enqueue_flux_pro_kontext_max_edit_webhook::{enqueue_flux_pro_kontext_max_edit_webhook, FluxProKontextMaxArgs, FluxProKontextMaxNumImages};
use errors::AnyhowResult;
use std::fs::read_to_string;
use test_data::web::image_urls::TALL_MOCHI_WITH_GLASSES_IMAGE_URL;
#[tokio::test]
#[ignore]
async fn test() -> AnyhowResult<()> {
// XXX: Don't commit secrets!
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
let api_key = FalApiKey::from_str(&secret);
let args = FluxProKontextMaxArgs {
image_url: TALL_MOCHI_WITH_GLASSES_IMAGE_URL,
prompt: "turn the glasses into sunglasses, make them sleek sunglasses with black rims, square shaped",
api_key: &api_key,
webhook_url: "https://example.com/webhook",
num_images: FluxProKontextMaxNumImages::One,
};
let result = enqueue_flux_pro_kontext_max_edit_webhook(args).await?;
Ok(())
}
}
@@ -1 +0,0 @@
pub mod enqueue_flux_pro_kontext_max_edit_webhook;
@@ -1,92 +0,0 @@
use crate::creds::fal_api_key::FalApiKey;
use crate::error::classify_fal_error::classify_fal_error;
use crate::error::fal_error_plus::FalErrorPlus;
use fal::prelude::fal_ai::flux_pro::v1::fill::fill;
use fal::prelude::fal_ai::flux_pro::v1::fill::FluxProFillInput;
use fal::webhook::WebhookResponse;
use reqwest::IntoUrl;
pub struct FluxPro1InfillArgs<'a, U: IntoUrl, R: IntoUrl, L: IntoUrl> {
// Request
pub prompt: &'a str,
pub image_url: U,
pub mask_url: R,
pub num_images: FluxPro1InfillNumImages,
// Fulfillment
pub webhook_url: L,
pub api_key: &'a FalApiKey,
}
#[derive(Copy, Clone, Debug)]
pub enum FluxPro1InfillNumImages {
One, // Default
Two,
Three,
Four,
}
pub async fn enqueue_flux_pro_1_infill_webhook<U: IntoUrl, R: IntoUrl, L: IntoUrl>(
args: FluxPro1InfillArgs<'_, U, R, L>
) -> Result<WebhookResponse, FalErrorPlus> {
let num_images = match args.num_images {
FluxPro1InfillNumImages::One => 1,
FluxPro1InfillNumImages::Two => 2,
FluxPro1InfillNumImages::Three => 3,
FluxPro1InfillNumImages::Four => 4,
};
let request = FluxProFillInput {
prompt: args.prompt.to_string(),
image_url: args.image_url.as_str().to_string(),
mask_url: args.mask_url.as_str().to_string(),
num_images: Some(num_images),
// Maybe expose
safety_tolerance: Some("5".to_string()), // NB: 5 is most tolerant
output_format: Some("png".to_string()), // png or jpeg
seed: None,
// Constants
sync_mode: None, // Synchronous / slow
};
let result = fill(request)
.with_api_key(&args.api_key.0)
.queue_webhook(args.webhook_url)
.await;
result.map_err(|err| classify_fal_error(err))
}
#[cfg(test)]
mod tests {
use crate::creds::fal_api_key::FalApiKey;
use crate::requests::webhook::image::infill::enqueue_flux_pro_1_infill_webhook::{enqueue_flux_pro_1_infill_webhook, FluxPro1InfillArgs, FluxPro1InfillNumImages};
use errors::AnyhowResult;
use std::fs::read_to_string;
use test_data::web::image_urls::{TALL_MOCHI_WITH_GLASSES_GLASSES_MASK_IMAGE_URL, TALL_MOCHI_WITH_GLASSES_IMAGE_URL};
#[tokio::test]
#[ignore]
async fn test() -> AnyhowResult<()> {
// XXX: Don't commit secrets!
let secret = read_to_string("/Users/bt/Artcraft/credentials/fal_api_key.txt")?;
let api_key = FalApiKey::from_str(&secret);
let args = FluxPro1InfillArgs {
image_url: TALL_MOCHI_WITH_GLASSES_IMAGE_URL,
mask_url: TALL_MOCHI_WITH_GLASSES_GLASSES_MASK_IMAGE_URL,
prompt: "slick sunglasses, cool glasses, reflection in glasses lenses",
api_key: &api_key,
webhook_url: "https://example.com/webhook",
num_images: FluxPro1InfillNumImages::One,
};
let result = enqueue_flux_pro_1_infill_webhook(args).await?;
Ok(())
}
}
@@ -1 +0,0 @@
pub mod enqueue_flux_pro_1_infill_webhook;
@@ -1,9 +1,7 @@
pub mod edit;
pub mod enqueue_flux_1_dev_text_to_image_webhook;
pub mod enqueue_flux_1_schnell_text_to_image_webhook;
pub mod enqueue_flux_pro_11_text_to_image_webhook;
pub mod enqueue_flux_pro_11_ultra_text_to_image_webhook;
pub mod enqueue_gpt_image_1_edit_image_webhook;
pub mod enqueue_gpt_image_1_text_to_image_webhook;
pub mod infill;
pub mod remove_background_rembg_webhook;
@@ -1,55 +0,0 @@
[package]
name = "midjourney_client"
edition = "2024"
version = "0.0.1"
publish = false
[package.metadata.cargo-udeps.ignore]
normal = ["workspace-build-acceleration"]
[lib]
name = "midjourney_client"
path = "src/lib.rs"
[dependencies]
# Workspace hack (faster compile times)
workspace-build-acceleration.workspace = true
# Internal
cloudflare_errors.workspace = true
cookie_store.workspace = true
errors.workspace = true
# External
anyhow.workspace = true
futures.workspace = true
log.workspace = true
#reqwest-websocket.workspace = true
#reqwest.workspace = true
#reqwest-impersonate = { version = "0", features = ["websocket"] }
serde.workspace = true
serde_json.workspace = true
#tokio.workspace = true
url.workspace = true
# NB: wreq gets past fingerprinting, but we have to deal with platform-specific build issues.
#
# Mac:
# brew install --formula cmake pkg-config llvm
# Windows:
# choco install cmake -y
# choco install strawberryperl -y
# choco install pkgconfiglite -y
# choco install llvm -y
# choco install nasm -y
wreq = { version = "5", features = ["json"] }
wreq-util = "2"
tokio = { version = "1", features = ["full"] }
# NB: We need to parse HTML to get out the websocket token
dom_query = "0.20.2"
[dev-dependencies]
filesys.workspace = true
test_data.workspace = true
testing.workspace = true
@@ -1,17 +0,0 @@
/// Multiple third party Midjourney clients make this configurable.
/// Not sure why, I'm just cargo culting it.
#[derive(Clone)]
pub enum MidjourneyHostname {
Standard,
Custom(String),
}
impl MidjourneyHostname {
pub fn as_str(&self) -> &str {
match self {
Self::Standard => "www.midjourney.com",
Self::Custom(domain) => domain,
}
}
}
@@ -1,2 +0,0 @@
pub mod midjourney_hostname;
pub mod midjourney_websocket_client;
@@ -1,6 +0,0 @@
/// One of the apparent authentication cookies used by Midjourney.
pub (crate) const AUTH_COOKIE_NAME_R : &str = "__Host-Midjourney.AuthUserTokenV3_r";
/// One of the apparent authentication cookies used by Midjourney.
pub (crate) const AUTH_COOKIE_NAME_I : &str = "__Host-Midjourney.AuthUserTokenV3_i";
@@ -1,8 +0,0 @@
use crate::credentials::cookie_names::{AUTH_COOKIE_NAME_I, AUTH_COOKIE_NAME_R};
use cookie_store::cookie_store::CookieStore;
/// Returns true if the given cookie store has the necessary authentication cookies.
pub fn cookie_store_has_auth_cookies(cookie_store: &CookieStore) -> bool {
cookie_store.has_cookie(AUTH_COOKIE_NAME_I)
&& cookie_store.has_cookie(AUTH_COOKIE_NAME_R)
}
@@ -1,31 +0,0 @@
/*
{
"name": "u3869671173",
"picture": "https://lh3.googleusercontent.com/a/ACg8ocKw0aaXKYSqy-Xkkn69ydIJFGnJA37S13XvjNY4d9xCzUzKUw=s96-c",
"midjourney_id": "f8a57ac3-e416-4dd4-9be8-2c4223691b01",
"iss": "https://securetoken.google.com/authjourney",
"aud": "authjourney",
"auth_time": 1754970849,
"user_id": "bS9idl73DlSbE61KnCyKpDfazCX2",
"sub": "bS9idl73DlSbE61KnCyKpDfazCX2",
"iat": 1754970849,
"exp": 1754974449,
"email": "victorvoice89@gmail.com",
"email_verified": true,
"firebase": {
"identities": {
"google.com": [
"116459300198619888836"
],
"email": [
"victorvoice89@gmail.com"
]
},
"sign_in_provider": "google.com"
}
}
*/
/*
*/
@@ -1,22 +0,0 @@
/// The user ID used for several midjourney requests.
#[derive(Clone, Debug, Eq, PartialEq, Hash)]
pub struct MidjourneyUserId(pub String);
impl MidjourneyUserId {
pub fn from_str(user_id: &str) -> Self {
Self(user_id.to_string())
}
pub fn from_string(user_id: String) -> Self {
Self(user_id)
}
pub fn as_str(&self) -> &str {
&self.0
}
pub fn to_string(&self) -> String {
self.0.to_string()
}
}
@@ -1,4 +0,0 @@
pub mod cookie_names;
pub mod cookie_store_has_auth_cookies;
pub mod jwt_decode;
pub mod midjourney_user_id;
@@ -1,95 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
use log::error;
use serde::Deserialize;
use wreq::Client;
use wreq_util::Emulation;
/// This grabs the pre-rendered HTML page, which contains the websocket token and other details.
pub struct GetIndexPageRequest {
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
pub async fn get_index_page_html(req: GetIndexPageRequest) -> Result<String, MidjourneyError> {
let url = format!("https://{}/", req.hostname.as_str());
let client = Client::builder()
.emulation(Emulation::Firefox139)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let cookie_header = req.cookie_header.trim();
if cookie_header.len() < 20 {
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
return Err(MidjourneyClientError::CookieTooShort.into());
}
// NB: missing headers that were in the browser request:
// -H 'sec-ch-ua-platform: "macOS"' \
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
let mut http_request = client.get(url)
.header("cookie", cookie_header)
.header("Referrer-Policy", "origin-when-cross-origin")
.header("accept", "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,image/apng,*/*;q=0.8,application/signed-exchange;v=b3;q=0.7")
.header("accept-language", "en-US,en;q=0.8")
.header("content-type", "application/json")
.header("priority", "u=0, i")
.header("sec-ch-ua-mobile", "?0")
.header("sec-fetch-dest", "document")
.header("sec-fetch-mode", "navigate")
.header("sec-fetch-site", "none")
.header("sec-fetch-user", "?1")
.header("upgrade-insecure-requests", "1");
let http_request = http_request
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let response = client.execute(http_request)
.await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
let status = response.status();
let response_body = response.text().await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
if !status.is_success() {
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
return Err(MidjourneyApiError::CloudflareError(err).into());
}
}
Ok(response_body)
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest};
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let result = get_index_page_html(GetIndexPageRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}\n\n", result);
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,245 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::credentials::midjourney_user_id::MidjourneyUserId;
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
use log::error;
use serde::Deserialize;
use wreq::Client;
use wreq_util::Emulation;
const DEFAULT_PAGE_SIZE : u64 = 1000;
/// This returns Midjourney's job/media list
pub struct ImagineRequest{
pub hostname: MidjourneyHostname,
pub cookie_header: String,
pub user_id: MidjourneyUserId,
pub page_size: Option<u64>,
}
#[derive(Debug, Clone)]
pub struct ImagineResponse {
pub items: Vec<ImagineItem>,
pub cursor: Option<String>,
}
#[derive(Debug, Clone)]
pub struct ImagineItem {
pub id: Option<String>,
/// This is the prompt, e.g.
/// `cute puppy wearing a polar bear suit --ar 4:3 --raw --v 7.0`
pub full_command: Option<String>,
pub job_type: Option<JobType>,
}
#[derive(Debug, Clone)]
pub enum JobType {
V6Diffusion,
V7Diffusion,
V7DraftRawDiffusion,
V7RawDiffusion,
Vid11I2vRenderAJointVideo,
Vid11I2vRenderBJointVideo,
Vid11I2vStartEndAVideo,
Vid11I2vStartEndBVideo,
Other(String),
}
impl JobType {
pub fn from_str(s: &str) -> JobType {
match s {
"v6_diffusion" => JobType::V6Diffusion,
"v7_diffusion" => JobType::V7Diffusion,
"v7_draft_raw_diffusion" => JobType::V7DraftRawDiffusion,
"v7_raw_diffusion" => JobType::V7RawDiffusion,
"vid_1.1_i2v_render_a_joint_video" => JobType::Vid11I2vRenderAJointVideo,
"vid_1.1_i2v_render_b_joint_video" => JobType::Vid11I2vRenderBJointVideo,
"vid_1.1_i2v_start_end_a_video" => JobType::Vid11I2vStartEndAVideo,
"vid_1.1_i2v_start_end_b_video" => JobType::Vid11I2vStartEndBVideo,
other => JobType::Other(other.to_string()),
}
}
}
pub async fn imagine(req: ImagineRequest) -> Result<ImagineResponse, MidjourneyError> {
let cookie_header = req.cookie_header.trim();
if cookie_header.len() < 20 {
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
return Err(MidjourneyClientError::CookieTooShort.into());
}
let page_size = req.page_size.unwrap_or(DEFAULT_PAGE_SIZE);
let referer = format!("https://{}/imagine", req.hostname.as_str());
let url = format!("https://{}/api/imagine?user_id={}&page_size={}",
req.hostname.as_str(),
req.user_id.as_str(),
page_size,
);
let client = Client::builder()
.emulation(Emulation::Firefox139)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
// NB: missing headers that were in the browser request:
// -H 'sec-ch-ua-platform: "macOS"' \
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
let mut http_request = client.get(url)
.header("cookie", cookie_header)
.header("Referer", &referer)
.header("Referrer-Policy", "origin-when-cross-origin")
.header("accept", "*/*")
.header("accept-encoding", "gzip, deflate, br, zstd")
.header("accept-language", "en-US,en;q=0.8")
.header("content-type", "application/json")
.header("priority", "u=1, i")
.header("sec-ch-ua-mobile", "?0")
.header("sec-fetch-dest", "empty")
.header("sec-fetch-mode", "cors")
.header("sec-fetch-site", "same-origin")
.header("x-csrf-protection", "1");
let http_request = http_request
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let response = client.execute(http_request)
.await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
let status = response.status();
let response_body = response.text().await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
if !status.is_success() {
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
return Err(MidjourneyApiError::CloudflareError(err).into());
}
}
/*
{
"data": [
{
"id": "a2339f14-ff71-43e4-a69f-8da15e7ba871",
"enqueue_time": "2025-08-14T17:21:28.403560+00:00",
"parent_id": null,
"video_segments": null,
"rating": null,
"job_type": "v7_diffusion",
"event_type": "diffusion",
"parent_grid": null,
"full_command": "a modern n64 console",
"batch_size": 4,
"width": 1024,
"height": 1024,
"published": true,
"shown": true,
"user_hidden": false,
"template": null,
"vid_framecount": null
},
],
"cursor": "encrypted cursor",
"checkpoint": "encrypted something ???"
}
*/
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawImagineResponse {
data: Vec<RawImagineItem>,
cursor: Option<String>,
//pub checkpoint: Option<String>,
}
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawImagineItem {
// The job id
id: Option<String>,
// The prompt plus `--` args.
full_command: Option<String>,
// Number of images, typically 4
batch_size: Option<u8>,
// Some possible values:
// "v6_diffusion"
// "v7_diffusion"
// "v7_draft_raw_diffusion"
// "v7_raw_diffusion"
// "vid_1.1_i2v_render_b_joint_video",
// "vid_1.1_i2v_start_end_a_video"
// "vid_1.1_i2v_start_end_b_video"
job_type: Option<String>,
// Some possible values:
// "diffusion"
// "draft"
// "variation"
// "video_diffusion"
// "video_start_end"
event_type: Option<String>,
}
let response : RawImagineResponse = serde_json::from_str(&response_body)
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
Ok(ImagineResponse {
cursor: response.cursor.clone(),
items: response.data
.into_iter()
.map(|item| {
ImagineItem {
id: item.id,
full_command: item.full_command,
job_type: item.job_type.as_ref()
.map(|jt| JobType::from_str(jt))
}
}).collect(),
})
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::credentials::midjourney_user_id::MidjourneyUserId;
use crate::endpoints::imagine::{imagine, ImagineRequest};
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let user_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/user_id.txt")?;
let user_id = MidjourneyUserId::from_string(user_id);
let result = imagine(ImagineRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
user_id,
page_size: None,
}).await?;
println!("Response: {:?}\n\n", result);
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,5 +0,0 @@
pub mod get_index_page_html;
pub mod imagine;
pub mod storage_list;
pub mod submit_job;
pub mod user_connections;
@@ -1,147 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
use log::error;
use serde::Deserialize;
use wreq::Client;
use wreq_util::Emulation;
/// This endpoint returns the user id
pub struct GetStorageListRequest {
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
#[derive(Debug)]
pub struct StorageItem {
// NB: The path contains the user ID.
pub bucket_pathname: Option<String>,
// NB: Other fields aren't necessary for us right now.
}
pub async fn storage_list(req: GetStorageListRequest) -> Result<Vec<StorageItem>, MidjourneyError> {
let referer = format!("https://{}", req.hostname.as_str());
let url = format!("https://{}/api/storage-list", req.hostname.as_str());
let client = Client::builder()
.emulation(Emulation::Firefox139)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let cookie_header = req.cookie_header.trim();
if cookie_header.len() < 20 {
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
return Err(MidjourneyClientError::CookieTooShort.into());
}
// NB: missing headers that were in the browser request:
// -H 'sec-ch-ua-platform: "macOS"' \
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
let mut http_request = client.get(url)
.header("cookie", cookie_header)
.header("Referer", &referer)
.header("Referrer-Policy", "origin-when-cross-origin")
.header("accept", "*/*")
.header("accept-language", "en-US,en;q=0.8")
.header("content-type", "application/json")
.header("priority", "u=1, i")
.header("sec-ch-ua-mobile", "?0")
.header("sec-fetch-dest", "empty")
.header("sec-fetch-mode", "cors")
.header("sec-fetch-site", "same-origin")
.header("x-csrf-protection", "1");
let http_request = http_request
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let response = client.execute(http_request)
.await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
let status = response.status();
// if status == 301 {
// for (name, value) in response.headers().iter() {
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
// }
// }
let response_body = &response.text().await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
if !status.is_success() {
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
return Err(MidjourneyApiError::CloudflareError(err).into());
}
}
/*
[
{
"bucketPathname": "UUID/OTHER_HEX_ENTROPY.png",
"shortUrl": null,
"type": "image-prompt",
"origin": "web",
"timeCreated": TIMESTAMP,
"hidden": false,
"contentType": "image/png",
"state": "loaded"
}
]
*/
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawStorageItem {
bucketPathname: Option<String>,
}
let response : Vec<RawStorageItem> = serde_json::from_str(response_body)
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
let items = response
.into_iter()
.map(|r| StorageItem {
bucket_pathname: r.bucketPathname,
})
.collect::<Vec<_>>();
Ok(items)
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
use crate::endpoints::storage_list::{storage_list, GetStorageListRequest};
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let result = storage_list(GetStorageListRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}\n\n", result);
for provider in result {
println!("Provider: {:?}", provider);
}
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,200 +0,0 @@
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
use log::error;
use serde::{Deserialize, Serialize};
use crate::client::midjourney_hostname::MidjourneyHostname;
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
use wreq::Client;
use wreq_util::Emulation;
pub struct SubmitJobRequest<'a> {
pub prompt: &'a str,
pub channel_id: &'a str,
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
#[derive(Debug, Clone)]
pub struct SubmitJobResponse {
pub job_id: String,
}
pub async fn submit_job(req: SubmitJobRequest<'_>) -> Result<SubmitJobResponse, MidjourneyError> {
let client = Client::builder()
.emulation(Emulation::Firefox139)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let referer = format!("https://{}", req.hostname.as_str());
// NB: Other recent clients use /api/app/submit-jobs, but this appears invalid.
let url = format!("https://{}/api/submit-jobs", req.hostname.as_str());
let cookie_header = req.cookie_header.trim();
if cookie_header.len() < 20 {
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
return Err(MidjourneyClientError::CookieTooShort.into());
}
let mut http_request = client.post(url)
.header("cookie", cookie_header)
.header("Referer", &referer)
.header("Referrer-Policy", "origin-when-cross-origin")
.header("accept", "*/*")
.header("accept-language", "en-US,en;q=0.8")
.header("content-type", "application/json")
.header("priority", "u=1, i")
.header("sec-ch-ua-mobile", "?0")
.header("sec-fetch-dest", "empty")
.header("sec-fetch-mode", "cors")
.header("sec-fetch-site", "same-origin")
.header("x-csrf-protection", "1");
#[derive(Serialize)]
struct F {
mode: String,
private: bool,
}
#[derive(Serialize)]
#[allow(non_snake_case)]
struct Metadata {
imagePrompts: u8,
imageReferences: u8,
characterReferences: u8,
depthReferences: u8,
lightboxOpen: String,
}
#[derive(Serialize)]
#[allow(non_snake_case)]
struct RawRequest {
f: F,
channelId: String,
metadata: Metadata,
t: String,
prompt: String,
}
let body = RawRequest {
f: F {
mode: "fast".to_string(),
private: false,
},
channelId: req.channel_id.to_string(),
metadata: Metadata {
imagePrompts: 0,
imageReferences: 0,
characterReferences: 0,
depthReferences: 0,
lightboxOpen: "".to_string(),
},
t: "imagine".to_string(),
prompt: req.prompt.to_string(),
};
let http_request = http_request
.json(&body)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let response = client.execute(http_request)
.await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
let status = response.status();
// if status == 301 {
// for (name, value) in response.headers().iter() {
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
// }
// }
let response_body = &response.text().await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
if !status.is_success() {
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
return Err(MidjourneyApiError::CloudflareError(err).into());
}
}
/*
{"success":[
{
"job_id":"UUID-like-string",
"prompt":"prompt",
"is_queued":false,
"event_type":"diffusion",
"flags":{
"mode":"fast","visibility":"public"
},
"meta":{
"height":1024,"width":1024,"batch_size":4,"parent_id":null,"parent_grid":null
},
"optimisticJobIndex":0,
"personalization_codes":null
}],
"failure":[]}
*/
#[derive(Deserialize)]
struct SuccessPayload {
job_id: String,
}
#[derive(Deserialize)]
#[allow(non_snake_case)]
struct RawResponse {
success: Vec<SuccessPayload>,
}
let response = serde_json::from_str::<RawResponse>(response_body)
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
let job_id = match response.success.get(0) {
None => {
error!("No job id found in body: {:?}", response_body);
return Err(MidjourneyApiError::NoJobId.into());
}
Some(item) => item.job_id.clone(),
};
Ok(SubmitJobResponse {
job_id,
})
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::endpoints::submit_job::{submit_job, SubmitJobRequest};
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
// Get channel id via:
// https://identitytoolkit.googleapis.com/v1/accounts:lookup?key=[TOKEN]
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let channel_id = read_to_trimmed_string("/Users/bt/secrets/midjourney/channel_id.txt")?;
let result = submit_job(SubmitJobRequest {
prompt: "a modern n64 console",
channel_id: &channel_id,
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}", result);
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,178 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
use cloudflare_errors::filter_cloudflare_errors::filter_cloudflare_errors;
use log::error;
use serde::Deserialize;
use wreq::Client;
use wreq_util::Emulation;
const USER_CONNECTIONS_URL :&str = "https://www.midjourney.com/api/user-connections";
pub struct GetUserConnectionsRequest {
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
#[derive(Debug)]
pub struct UserData {
pub provider_id: String,
pub display_name: Option<String>,
pub email: Option<String>,
pub user_id: Option<String>,
}
#[derive(Debug)]
pub struct Provider {
pub provider_id: String,
pub user_data: Option<UserData>,
pub is_linked: bool,
}
pub async fn user_connections(req: GetUserConnectionsRequest) -> Result<Vec<Provider>, MidjourneyError> {
let referer = format!("https://{}", req.hostname.as_str());
let url = format!("https://{}/api/user-connections", req.hostname.as_str());
let client = Client::builder()
.emulation(Emulation::Firefox139)
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let cookie_header = req.cookie_header.trim();
if cookie_header.len() < 20 {
error!("Cookie header is too short (len: {}): {}", cookie_header.len(), cookie_header);
return Err(MidjourneyClientError::CookieTooShort.into());
}
// NB: missing headers that were in the browser request:
// -H 'sec-ch-ua-platform: "macOS"' \
// -H 'sec-ch-ua: "Not)A;Brand";v="8", "Chromium";v="138", "Google Chrome";v="138"' \
let mut http_request = client.get(url)
.header("cookie", cookie_header)
.header("Referer", &referer)
.header("Referrer-Policy", "origin-when-cross-origin")
.header("accept", "*/*")
.header("accept-language", "en-US,en;q=0.8")
.header("content-type", "application/json")
.header("priority", "u=1, i")
.header("sec-ch-ua-mobile", "?0")
.header("sec-fetch-dest", "empty")
.header("sec-fetch-mode", "cors")
.header("sec-fetch-site", "same-origin")
.header("x-csrf-protection", "1");
let http_request = http_request
.build()
.map_err(|err| MidjourneyClientError::WreqError(err))?;
let response = client.execute(http_request)
.await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
let status = response.status();
// if status == 301 {
// for (name, value) in response.headers().iter() {
// println!("- {}: {}", name.as_str(), value.to_str().unwrap());
// }
// }
let response_body = &response.text().await
.map_err(|e| MidjourneyApiError::NetworkError(e.to_string()))?;
if !status.is_success() {
if let Err(err) = filter_cloudflare_errors(status.as_u16(), &response_body) {
return Err(MidjourneyApiError::CloudflareError(err).into());
}
}
/*
[
{
"providerId": "google.com",
"userData": {
"providerId": "google.com",
"displayName": "John Doe",
"email": "john.doe@gmail.com",
"userId": "123456789000000001",
},
"isLinked": true
},
{
"providerId": "discord.com",
"isLinked": false
}
]
*/
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawProviderUserData {
providerId: String,
displayName: Option<String>,
email: Option<String>,
userId: Option<String>,
}
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawProvider {
providerId: String,
userData: Option<RawProviderUserData>,
isLinked: bool,
}
let response : Vec<RawProvider> = serde_json::from_str(response_body)
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
let providers = response
.into_iter()
.map(|r| Provider {
provider_id: r.providerId,
user_data: r.userData.map(|u| UserData {
provider_id: u.providerId,
display_name: u.displayName,
email: u.email,
user_id: u.userId,
}),
is_linked: r.isLinked,
})
.collect::<Vec<_>>();
Ok(providers)
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::endpoints::user_connections::{user_connections, GetUserConnectionsRequest};
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let result = user_connections(GetUserConnectionsRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}\n\n", result);
for provider in result {
println!("Provider: {:?}", provider);
}
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,137 +0,0 @@
use cloudflare_errors::cloudflare_error::CloudflareError;
use errors::AnyhowError;
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::io;
#[derive(Debug)]
pub enum MidjourneyApiError {
/// There was no job ID in the otherwise valid-looking response.
NoJobId,
/// There was no user ID in an otherwise valid-looking response.
NoUserId,
/// No user props in index HTML payload.
/// We need this for the user ID and the websocket token.
NoUserProps,
/// No initialAuthUser in the index HTML payload.
/// We need this for the user ID and the websocket token.
NoInitialAuthUser,
/// 400. The request was invalid.
InvalidRequest(String),
/// 401. The request was not authorized.
Unauthorized(String),
/// 403. The request was forbidden.
Forbidden(String),
/// 404. The requested resource was not found.
NotFound(String),
/// 429. Too many requests.
TooManyRequests(String),
/// 500. An internal server error occurred.
InternalServerError {
body: String,
backend_hostname: Option<String>,
},
/// Eg. when downloading images
UnknownHttpFailure {
status_code: u16,
body: String,
},
/// Cloudflare errors.
CloudflareError(CloudflareError),
/// A deserialization error with the response.
DeserializationError(serde_json::Error),
/// The request timed out.
Timeout(String),
/// A network error occurred.
NetworkError(String),
// /// Uncategorized reqwest error.
// OtherReqwestError(reqwest::Error),
/// An error doing file I/O (on our side)
IoError(io::Error),
/// Another type of error.
Other(AnyhowError),
}
impl Error for MidjourneyApiError {}
impl Display for MidjourneyApiError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
// Response body errors
Self::NoJobId => write!(f, "No job ID found in the response body."),
Self::NoUserId => write!(f, "No user ID found in the response body."),
Self::NoUserProps => write!(f, "No user properties found in the index HTML payload."),
Self::NoInitialAuthUser => write!(f, "No initialAuthUser in the index HTML payload."),
// Server response code errors
Self::InvalidRequest(msg) => write!(f, "Invalid request: {}", msg),
Self::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
Self::Forbidden(msg) => write!(f, "Forbidden: {}", msg),
Self::NotFound(msg) => write!(f, "Not found: {}", msg),
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
Self::Timeout(msg) => write!(f, "Timeout: {}", msg),
Self::NetworkError(msg) => write!(f, "Network error: {}", msg),
// Cloudflare errors
Self::CloudflareError(error) => write!(f, "Cloudflare Error: {}", error),
// I/O errors
Self::IoError(error) => write!(f, "IO error: {}", error),
// Other
// Self::OtherReqwestError(error) => write!(f, "Reqwest error: {}", error),
Self::Other(error) => write!(f, "Other error: {}", error),
}
}
}
/*impl From<reqwest::Error> for MidjourneyApiError {
fn from(error: reqwest::Error) -> Self {
if error.is_timeout() {
Self::Timeout(error.to_string())
} else if error.is_connect() {
Self::NetworkError(error.to_string())
} else {
Self::OtherReqwestError(error)
}
}
}*/
impl From<serde_json::Error> for MidjourneyApiError {
fn from(error: serde_json::Error) -> Self {
Self::DeserializationError(error)
}
}
impl From<io::Error> for MidjourneyApiError {
fn from(error: io::Error) -> Self {
Self::IoError(error)
}
}
impl From<CloudflareError> for MidjourneyApiError {
fn from(error: CloudflareError) -> Self {
Self::CloudflareError(error)
}
}
@@ -1,71 +0,0 @@
use std::error::Error;
use std::fmt::{Display, Formatter};
use std::io;
use std::path::PathBuf;
/// An error arising from the client, either before or after the request is sent.
#[derive(Debug)]
pub enum MidjourneyClientError {
/// The cookie was too short to be valid.
CookieTooShort,
/// Tried to look up an image out of range [0,4)
InvalidImageIndex,
/// Could not determine the filetype for the file.
FileTypeNotKnown(PathBuf),
/// Could not handle the file type.
FileTypeNotHandled(PathBuf),
/// An error doing file I/O (on our side)
IoError(io::Error),
// /// An error setting up the client
// ReqwestError(reqwest::Error),
WreqError(wreq::Error),
/// A serialization error with the request.
SerializationError(serde_json::Error),
}
impl Error for MidjourneyClientError {}
impl Display for MidjourneyClientError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
Self::CookieTooShort => write!(f, "The cookie header is too short to be valid."),
Self::InvalidImageIndex => write!(f, "The image index is out of range. Valid range is [0, 4)."),
Self::FileTypeNotKnown(path) => write!(f, "Could not determine the filetype for the file: {:?}", path),
Self::FileTypeNotHandled(path) => write!(f, "Could not handle the file type for the file: {:?}", path),
Self::IoError(err) => write!(f, "IO error: {}", err),
//Self::ReqwestError(err) => write!(f, "Reqwest client error: {}", err),
Self::WreqError(err) => write!(f, "Wreq client error: {}", err),
Self::SerializationError(err) => write!(f, "Request serialization error: {}", err),
}
}
}
impl From<io::Error> for MidjourneyClientError {
fn from(error: io::Error) -> Self {
Self::IoError(error)
}
}
impl From<serde_json::Error> for MidjourneyClientError {
fn from(error: serde_json::Error) -> Self {
Self::SerializationError(error)
}
}
/*impl From<reqwest::Error> for MidjourneyClientError {
fn from(error: reqwest::Error) -> Self {
Self::ReqwestError(error)
}
}*/
impl From<wreq::Error> for MidjourneyClientError {
fn from(error: wreq::Error) -> Self {
Self::WreqError(error)
}
}
@@ -1,32 +0,0 @@
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_client_error::MidjourneyClientError;
use std::error::Error;
#[derive(Debug)]
pub enum MidjourneyError {
Api(MidjourneyApiError),
Client(MidjourneyClientError),
}
impl Error for MidjourneyError {}
impl std::fmt::Display for MidjourneyError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Api(e) => write!(f, "Midjourney API error: {:?}", e),
Self::Client(e) => write!(f, "Midjourney Client error: {:?}", e),
}
}
}
impl From<MidjourneyApiError> for MidjourneyError {
fn from(error: MidjourneyApiError) -> Self {
Self::Api(error)
}
}
impl From<MidjourneyClientError> for MidjourneyError {
fn from(error: MidjourneyClientError) -> Self {
Self::Client(error)
}
}
@@ -1,3 +0,0 @@
pub mod midjourney_error;
pub mod midjourney_api_error;
pub mod midjourney_client_error;
@@ -1,8 +0,0 @@
//! A midjourney API client.
pub mod client;
pub mod credentials;
pub mod endpoints;
pub mod error;
pub mod recipes;
pub mod utils;
@@ -1,20 +0,0 @@
use crate::credentials::midjourney_user_id::MidjourneyUserId;
#[derive(Debug, Clone)]
pub enum ChannelId {
/// Channel based on the user ID.
/// This is what Midjourney uses by default.
UserId(MidjourneyUserId),
// Bare string ID
Raw(String),
}
impl ChannelId {
pub fn to_string(&self) -> String {
match self {
Self::UserId(user_id) => format!("singleplayer_{}", user_id.0),
Self::Raw(id) => id.clone(),
}
}
}
@@ -1,52 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::credentials::midjourney_user_id::MidjourneyUserId;
use crate::endpoints::storage_list::{storage_list, GetStorageListRequest};
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_error::MidjourneyError;
pub struct GetMidjourneyUserIdRequest {
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
pub async fn get_midjourney_user_id(req: GetMidjourneyUserIdRequest) -> Result<MidjourneyUserId, MidjourneyError> {
let items = storage_list(GetStorageListRequest {
hostname: req.hostname,
cookie_header: req.cookie_header,
}).await?;
let user_id = items
.into_iter()
.find_map(|item| item.bucket_pathname)
.and_then(|path| path.split('/').next().map(|id| id.to_owned()))
.map(|id| MidjourneyUserId(id.to_string()))
.ok_or(MidjourneyApiError::NoUserId)?;
Ok(user_id)
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::recipes::get_midjourney_user_id::{get_midjourney_user_id, GetMidjourneyUserIdRequest};
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let result = get_midjourney_user_id(GetMidjourneyUserIdRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}\n\n", result);
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,153 +0,0 @@
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::credentials::midjourney_user_id::MidjourneyUserId;
use crate::endpoints::get_index_page_html::{get_index_page_html, GetIndexPageRequest};
use crate::error::midjourney_api_error::MidjourneyApiError;
use crate::error::midjourney_error::MidjourneyError;
use dom_query::Document;
use serde::Deserialize;
pub struct GetUserInfoRequest {
pub hostname: MidjourneyHostname,
pub cookie_header: String,
}
#[derive(Debug, Clone)]
pub struct GetUserInfoResponse {
pub user_id: Option<MidjourneyUserId>,
pub email: Option<String>,
pub websocket_token: Option<String>,
}
pub async fn get_user_info(req: GetUserInfoRequest) -> Result<GetUserInfoResponse, MidjourneyError> {
let index_html = get_index_page_html(GetIndexPageRequest {
hostname: req.hostname,
cookie_header: req.cookie_header,
}).await?;
/*
<script id="initialProps" type="application/json">
{
"initialAuthUser": {
"idpID": "ENTROPY",
"id": "UUID",
"displayName": "uNUMBER",
"email": "USER@gmail.com",
"emailVerified": true,
"photoURL": "https://lh3.googleusercontent.com/a/ENTROPY",
"abilities": {
"admin": false,
"developer": false,
"accepted_tos": true,
"moderator": false,
"guide": false,
"vip": false,
"employee": false,
"content_moderator": false,
"gdpr_preferences_set": false,
"gdpr_allow_advertising": false,
"gdpr_allow_analytics": false,
"editor_whitelist": false,
"allow_nsfw": false,
"cooldowns_removed": false,
"blocked": false,
"community": false,
"billing": false,
"system_verified": false,
"simplified_flow": false,
"total_jobs": 57,
"is_sub_at_least_one_year_old": false,
"can_private": false,
"can_relax": false,
"has_discord": false,
"delete_at": null,
"subscription": {
"type": "subscription",
"plan_type": "basic",
"recurring": "month"
},
"relax_video_job_concurrency": 0
},
"websocketToken": "WEBSOCKET_TOKEN_ENTROPY",
"providers": [
"google.com"
],
"system": "firebase",
"streamToken": "STREAM_TOKEN_ENTROPY"
},
"urlFromInitialServerRequest": "/"
}
</script>
*/
let document = Document::from(index_html);
let maybe_nodes = document.select("script[id=initialProps]");
let node = match maybe_nodes.get(0) {
Some(node) => node,
None => {
return Err(MidjourneyApiError::NoUserProps.into());
}
};
let inner_json = node.inner_html().to_string();
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawAuthUser {
// NB: The Midjourney user id.
id: Option<String>,
email: Option<String>,
websocketToken: Option<String>,
}
#[derive(Deserialize, Debug)]
#[allow(non_snake_case)]
struct RawBody {
initialAuthUser: Option<RawAuthUser>,
}
let response : RawBody = serde_json::from_str(&inner_json)
.map_err(|err| MidjourneyApiError::DeserializationError(err))?;
let props = match response.initialAuthUser {
Some(props) => props,
None => {
return Err(MidjourneyApiError::NoInitialAuthUser.into());
}
};
Ok(GetUserInfoResponse {
user_id: props.id
.map(|id| MidjourneyUserId::from_string(id)),
email: props.email
.map(|email| email.to_string()),
websocket_token: props.websocketToken
.map(|email| email.to_string()),
})
}
#[cfg(test)]
mod tests {
use crate::client::midjourney_hostname::MidjourneyHostname;
use crate::recipes::get_user_info::{get_user_info, GetUserInfoRequest};
use errors::AnyhowResult;
use filesys::read_to_trimmed_string::read_to_trimmed_string;
#[ignore]
#[tokio::test]
async fn test() -> AnyhowResult<()> {
let cookie_header = read_to_trimmed_string("/Users/bt/secrets/midjourney/cookie.txt")?;
let result = get_user_info(GetUserInfoRequest {
cookie_header,
hostname: MidjourneyHostname::Standard,
}).await?;
println!("Response: {:?}\n\n", result);
assert_eq!(1, 2);
Ok(())
}
}
@@ -1,3 +0,0 @@
mod get_midjourney_user_id;
pub mod channel_id;
pub mod get_user_info;
@@ -1,9 +0,0 @@
use crate::error::midjourney_client_error::MidjourneyClientError;
use crate::error::midjourney_error::MidjourneyError;
pub fn get_image_url(job_id: &str, image_index: u8) -> Result<String, MidjourneyError> {
if image_index > 3 {
return Err(MidjourneyClientError::InvalidImageIndex.into())
}
Ok(format!("https://cdn.midjourney.com/{}/0_{}.png", job_id, image_index))
}
@@ -1,70 +0,0 @@
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<Self, MidjourneyClientError> {
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<Vec<u8>, 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())
}
}
@@ -1,2 +0,0 @@
pub mod get_image_url;
pub mod image_downloader_client;
@@ -59,7 +59,7 @@ impl Display for ApiError {
ApiError::Unauthorized(msg) => write!(f, "Unauthorized: {}", msg),
ApiError::Forbidden(msg) => write!(f, "Forbidden: {}", msg),
ApiError::NotFound(msg) => write!(f, "Not found: {}", msg),
ApiError::TooManyRequests(msg) => write!(f, "Too many requests to Storyteller backend: {}", msg),
ApiError::TooManyRequests(msg) => write!(f, "Too many requests: {}", msg),
ApiError::InternalServerError {body, backend_hostname} =>
write!(f, "Internal Server Error; backend hostname: {:?} ; body: {}; ", backend_hostname, body),
// Server response handling errors
@@ -13,8 +13,8 @@ impl Error for StorytellerError {}
impl std::fmt::Display for StorytellerError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
StorytellerError::Api(e) => write!(f, "Storyteller API error: {:?}", e),
StorytellerError::Client(e) => write!(f, "Storyteller Client error: {:?}", e),
StorytellerError::Api(e) => write!(f, "API error: {:?}", e),
StorytellerError::Client(e) => write!(f, "Client error: {:?}", e),
}
}
}
@@ -1,19 +0,0 @@
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::error::storyteller_error::StorytellerError;
use crate::utils::api_host::ApiHost;
use crate::utils::basic_json_post_request::basic_json_post_request;
use artcraft_api_defs::generate::image::edit::flux_pro_kontext_max_edit_image::{FluxProKontextMaxEditImageRequest, FluxProKontextMaxEditImageResponse, FLUX_PRO_KONTEXT_MAX_EDIT_IMAGE_PATH};
pub async fn flux_pro_kontext_max_edit_image(
api_host: &ApiHost,
maybe_creds: Option<&StorytellerCredentialSet>,
request: FluxProKontextMaxEditImageRequest,
) -> Result<FluxProKontextMaxEditImageResponse, StorytellerError> {
Ok(basic_json_post_request(
api_host,
FLUX_PRO_KONTEXT_MAX_EDIT_IMAGE_PATH,
maybe_creds,
request,
).await?)
}
@@ -1,2 +1 @@
pub mod flux_pro_kontext_max_edit_image;
pub mod gpt_image_1_edit_image;
@@ -1,18 +0,0 @@
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::error::storyteller_error::StorytellerError;
use crate::utils::api_host::ApiHost;
use crate::utils::basic_json_post_request::basic_json_post_request;
use artcraft_api_defs::generate::image::inpaint::flux_pro_1_inpaint_image::{FluxPro1InpaintImageRequest, FluxPro1InpaintImageResponse, FLUX_PRO_1_INPAINT_PATH};
pub async fn flux_pro_1_inpaint_image(
api_host: &ApiHost,
maybe_creds: Option<&StorytellerCredentialSet>,
request: FluxPro1InpaintImageRequest,
) -> Result<FluxPro1InpaintImageResponse, StorytellerError> {
Ok(basic_json_post_request(
api_host,
FLUX_PRO_1_INPAINT_PATH,
maybe_creds,
request,
).await?)
}
@@ -1 +0,0 @@
pub mod flux_pro_1_inpaint_image;
@@ -4,5 +4,4 @@ pub mod generate_flux_1_schnell_text_to_image;
pub mod generate_flux_pro_11_text_to_image;
pub mod generate_flux_pro_11_ultra_text_to_image;
pub mod generate_gpt_image_1_text_to_image;
pub mod inpaint;
pub mod remove_image_background;
@@ -1,5 +1,4 @@
pub mod get_media_file;
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;
@@ -1,138 +0,0 @@
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::shared_response_types::media_file_cover_image_details::MediaFileCoverImageDetails;
use crate::shared_response_types::media_links::MediaLinks;
use crate::shared_response_types::simple_entity_stats::SimpleEntityStats;
use crate::shared_response_types::user_details_light::UserDetailsLight;
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::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, Part};
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 UploadImageBytesArgs<'a> {
pub api_host: &'a ApiHost,
pub maybe_creds: Option<&'a StorytellerCredentialSet>,
// NB: Bytes need to be owned for the request.
pub image_bytes: Vec<u8>,
pub image_type: ImageType,
/// If true, we should hide the image from the user's gallery.
pub is_intermediate_system_file: bool,
}
#[derive(Clone, Copy, Debug)]
pub enum ImageType {
Png,
Jpeg,
Gif,
}
impl ImageType {
pub fn mime_type(&self) -> &str {
match self {
ImageType::Png => "image/png",
ImageType::Jpeg => "image/jpeg",
ImageType::Gif => "image/gif",
}
}
pub fn file_extension(&self) -> &str {
match self {
ImageType::Png => "png",
ImageType::Jpeg => "jpg",
ImageType::Gif => "gif",
}
}
}
/// Upload an image media file from a file.
/// NB: File type is not checked, so caller needs to enforce.
/// NB: We need owned bytes for the request.
pub async fn upload_image_media_file_from_bytes(
args: UploadImageBytesArgs<'_>
) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
let url = get_route(args.api_host);
debug!("Requesting {:?}", &url);
let client = Client::builder()
.gzip(true)
.build()
.map_err(|err| StorytellerError::Client(ClientError::from(err)))?;
let part = Part::bytes(args.image_bytes)
.file_name(format!("image.{}", args.image_type.file_extension()))
.mime_str(args.image_type.mime_type())
.map_err(|err| StorytellerError::Client(ClientError::from(err)))?;
let mut form = Form::new()
.text("uuid_idempotency_token", generate_random_uuid())
.part("file", part);
if args.is_intermediate_system_file {
form = form.text("is_intermediate_system_file", "true");
}
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/image", 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,
}
@@ -33,24 +33,14 @@ use tokens::tokens::model_weights::ModelWeightToken;
use tokens::tokens::prompts::PromptToken;
use uuid::uuid;
pub struct UploadImageFromFileArgs<'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,
}
/// Upload an image media file from a file.
pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
args: UploadImageFromFileArgs<'_, P>,
api_host: &ApiHost,
maybe_creds: Option<&StorytellerCredentialSet>,
path: P,
) -> Result<UploadImageMediaFileSuccessResponse, StorytellerError> {
let url = get_route(args.api_host);
let url = get_route(api_host);
debug!("Requesting {:?}", &url);
@@ -61,21 +51,17 @@ pub async fn upload_image_media_file_from_file<P: AsRef<Path>>(
.build()
.map_err(|err| StorytellerError::Client(ClientError::from(err)))?;
let mut form = Form::new()
let form = Form::new()
.text("uuid_idempotency_token", generate_random_uuid())
.file("file", args.path)
.file("file", path)
.await
.map_err(|err| StorytellerError::Client(ClientError::from(err)))?;
if args.is_intermediate_system_file {
form = form.text("is_intermediate_system_file", "true");
}
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(creds) = maybe_creds {
if let Some(header) = &creds.maybe_as_cookie_header() {
request_builder = request_builder.header("Cookie", header);
}
@@ -1,7 +1,7 @@
use crate::credentials::storyteller_credential_set::StorytellerCredentialSet;
use crate::error::client_error::ClientError;
use crate::error::storyteller_error::StorytellerError;
use crate::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
use crate::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
use crate::media_files::upload_new_engine_asset_from_file::upload_new_engine_asset_from_file;
use crate::media_files::upload_video_media_file_from_file::upload_video_media_file_from_file;
use crate::utils::api_host::ApiHost;
@@ -49,13 +49,7 @@ pub async fn upload_media_file_from_file<P: AsRef<Path>>(
}
}
Some(FileExtension::Png) | Some(FileExtension::Jpg) | Some(FileExtension::Gif) | Some(FileExtension::Webp) => {
let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
api_host,
maybe_creds,
path,
is_intermediate_system_file: false,
}).await;
match result {
match upload_image_media_file_from_file(api_host, maybe_creds, path).await {
Ok(result) => Ok(UploadMediaFileSuccessResponse {
success: result.success,
media_file_token: result.media_file_token,
+15 -18
View File
@@ -23,27 +23,24 @@ tauri-build = { version = "2.1.1", features = [] }
[dependencies]
# Workspace hack (faster compile times)
workspace-build-acceleration.workspace = true
workspace-build-acceleration = { workspace = true }
# Internal
artcraft_api_defs.workspace = true
build_metadata.workspace = true
cookie_store.workspace = true
enums.workspace = true
errors.workspace = true
fal_client.workspace = true
filesys.workspace = true
idempotency.workspace = true
images.workspace = true
memory_store.workspace = true
midjourney_client.workspace = true
mimetypes.workspace = true
openai_sora_client.workspace = true
primitives.workspace = true
artcraft_api_defs = { workspace = true }
build_metadata = { workspace = true }
enums = { workspace = true }
errors = { workspace = true }
fal_client = { workspace = true }
filesys = { workspace = true }
idempotency = { workspace = true }
memory_store = { workspace = true }
mimetypes = { workspace = true }
openai_sora_client = { workspace = true }
primitives = { workspace = true }
sqlite_tasks.workspace = true
storyteller_client.workspace = true
tokens.workspace = true
web_base64.workspace = true
storyteller_client = { workspace = true }
tokens = { workspace = true }
web_base64 = { workspace = true }
# External
anyhow = { workspace = true }
@@ -3,7 +3,6 @@ use crate::core::commands::enqueue::image::generic::handle_image_fal::handle_ima
use crate::core::commands::enqueue::image::gpt_image_1::handle_gpt_image_1::handle_gpt_image_1;
use crate::core::commands::enqueue::image::gpt_image_1::handle_gpt_image_1_sora::handle_gpt_image_1_sora;
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
use crate::core::commands::enqueue::image::midjourney::handle_midjourney::handle_midjourney;
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;
@@ -19,7 +18,6 @@ use crate::core::state::provider_priority::{Provider, ProviderPriorityStore};
use crate::core::state::task_database::TaskDatabase;
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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;
@@ -95,7 +93,6 @@ pub async fn enqueue_text_to_image_command(
provider_priority_store: State<'_, ProviderPriorityStore>,
task_database: State<'_, TaskDatabase>,
fal_creds_manager: State<'_, FalCredentialManager>,
mj_creds_manager: State<'_, MidjourneyCredentialManager>,
storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
fal_task_queue: State<'_, FalTaskQueue>,
sora_creds_manager: State<'_, SoraCredentialManager>,
@@ -111,7 +108,6 @@ pub async fn enqueue_text_to_image_command(
&provider_priority_store,
&task_database,
&fal_creds_manager,
&mj_creds_manager,
&storyteller_creds_manager,
&app_env_configs,
&fal_task_queue,
@@ -177,7 +173,6 @@ pub async fn handle_request(
provider_priority_store: &ProviderPriorityStore,
task_database: &TaskDatabase,
fal_creds_manager: &FalCredentialManager,
mj_creds_manager: &MidjourneyCredentialManager,
storyteller_creds_manager: &StorytellerCredentialManager,
app_env_configs: &AppEnvConfigs,
fal_task_queue: &FalTaskQueue,
@@ -194,7 +189,6 @@ pub async fn handle_request(
&storyteller_creds_manager,
&app_env_configs,
&fal_task_queue,
&mj_creds_manager,
&sora_creds_manager,
&sora_task_queue,
).await;
@@ -225,7 +219,6 @@ pub async fn dispatch_request(
storyteller_creds_manager: &StorytellerCredentialManager,
app_env_configs: &AppEnvConfigs,
fal_task_queue: &FalTaskQueue,
mj_creds_manager: &MidjourneyCredentialManager,
sora_creds_manager: &SoraCredentialManager,
sora_task_queue: &SoraTaskQueue,
) -> Result<TaskEnqueueSuccess, InternalImageError> {
@@ -246,13 +239,6 @@ pub async fn dispatch_request(
sora_task_queue,
).await;
}
Some(ImageModel::Midjourney) => {
return handle_midjourney(
request,
app_env_configs,
mj_creds_manager,
).await;
}
_ => {
// Fall-through
}
@@ -57,11 +57,8 @@ pub async fn handle_image_artcraft(
None => {
return Err(InternalImageError::NoModelSpecified);
}
Some(
ImageModel::GptImage1 |
ImageModel::Midjourney
) => {
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model)));
Some(ImageModel::GptImage1) => {
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: an earlier branch should have handled gpt-image-1 images for Artcraft")));
}
Some(ImageModel::Recraft3) => {
return Err(InternalImageError::AnyhowError(anyhow!("not yet implemented in Artcraft")));
@@ -51,11 +51,8 @@ pub async fn handle_image_fal(
None => {
return Err(InternalImageError::NoModelSpecified);
}
Some(
ImageModel::GptImage1 |
ImageModel::Midjourney
) => {
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: another branch should handle this: {:?}", request.model)));
Some(ImageModel::GptImage1) => {
return Err(InternalImageError::AnyhowError(anyhow!("wrong logic: fal is handling sora images")));
}
Some(
ImageModel::Flux1Dev |
@@ -2,7 +2,6 @@ use crate::core::artcraft_error::ArtcraftError;
use base64::DecodeError;
use errors::AnyhowError;
use fal_client::error::fal_error_plus::FalErrorPlus;
use midjourney_client::error::midjourney_error::MidjourneyError;
use openai_sora_client::sora_error::SoraError;
use storyteller_client::error::storyteller_error::StorytellerError;
@@ -11,11 +10,9 @@ pub enum InternalImageError {
NoModelSpecified,
NoProviderAvailable,
NeedsFalApiKey,
NeedsMidjourneyCredentials,
NeedsSoraCredentials,
NeedsStorytellerCredentials,
FalError(FalErrorPlus),
MidjourneyError(MidjourneyError),
SoraError(SoraError),
AnyhowError(AnyhowError),
StorytellerError(StorytellerError),
@@ -1,84 +0,0 @@
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize};
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
use crate::core::model::image_models::ImageModel;
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::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 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 fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs};
use idempotency::uuid::generate_random_uuid;
use log::{error, info};
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
use midjourney_client::endpoints::submit_job::{submit_job, SubmitJobRequest};
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::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image;
use storyteller_client::generate::image::generate_gpt_image_1_text_to_image::generate_gpt_image_1_text_to_image;
use tauri::AppHandle;
pub async fn handle_midjourney(
request: EnqueueTextToImageRequest,
app_env_configs: &AppEnvConfigs,
mj_creds_manager: &MidjourneyCredentialManager,
) -> Result<TaskEnqueueSuccess, InternalImageError> {
let creds = match mj_creds_manager.maybe_copy_cookie_store() {
Ok(Some(creds)) => creds,
Ok(None) => {
return Err(InternalImageError::NeedsMidjourneyCredentials);
}
Err(err) => {
error!("Error reading Midjourney credentials: {:?}", err);
return Err(InternalImageError::NeedsMidjourneyCredentials);
},
};
info!("Calling midjourney ...");
let cookie_header = creds.to_cookie_string();
let prompt = request.prompt
.as_deref()
.unwrap_or("");
let result = submit_job(SubmitJobRequest {
prompt,
channel_id: "singleplayer_f8a57ac3-e416-4dd4-9be8-2c4223691b01", // TODO: DO NOT COMMIT
hostname: MidjourneyHostname::Standard,
cookie_header,
}).await;
let job_id = match result {
Ok(result) => {
// TODO(bt,2025-07-05): Enqueue job token?
info!("Successfully enqueued MidJourney. Job token: {}",result.job_id);
result.job_id
}
Err(err) => {
error!("Failed to use MidJourney: {:?}", err);
return Err(InternalImageError::MidjourneyError(err));
}
};
Ok(TaskEnqueueSuccess {
provider: GenerationProvider::Midjourney,
model: Some(GenerationModel::Midjourney),
provider_job_id: Some(job_id),
task_type: TaskType::ImageGeneration,
})
}
@@ -1 +0,0 @@
pub (super) mod handle_midjourney;
@@ -1,5 +1,4 @@
pub mod enqueue_text_to_image_command;
pub mod generic;
pub mod internal_image_error;
mod gpt_image_1;
mod midjourney;
mod gpt_image_1;
@@ -43,7 +43,6 @@ use storyteller_client::error::storyteller_error::StorytellerError;
use storyteller_client::media_files::get_media_file::get_media_file;
use storyteller_client::utils::api_host::ApiHost;
use tauri::{AppHandle, Manager, State};
use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
use tokens::tokens::media_files::MediaFileToken;
#[derive(Deserialize, Debug)]
@@ -56,21 +55,6 @@ pub struct EnqueueImageBgRemovalCommand {
/// Base64-encoded image
pub base64_image: Option<String>,
/// 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, Debug)]
@@ -185,12 +169,7 @@ pub async fn handle_request(
).await?;
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()
)
.insert_into_task_database(task_database)
.await;
if let Err(err) = result {
@@ -39,7 +39,7 @@ use storyteller_client::generate::image::generate_flux_1_schnell_text_to_image::
use storyteller_client::generate::image::generate_flux_pro_11_text_to_image::generate_flux_pro_11_text_to_image;
use storyteller_client::generate::image::generate_flux_pro_11_ultra_text_to_image::generate_flux_pro_11_ultra_text_to_image;
use storyteller_client::generate::image::remove_image_background::remove_image_background;
use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
use tauri::AppHandle;
use tokens::tokens::media_files::MediaFileToken;
@@ -123,13 +123,9 @@ async fn upload_image_from_base64_bytes(
info!("Uploading image media file from temp file: {:?}", temp_file.path());
let result =
upload_image_media_file_from_file(UploadImageFromFileArgs {
api_host: &app_env_configs.storyteller_host,
maybe_creds: Some(&creds),
path: temp_file,
is_intermediate_system_file: true, // NB: Probably not essential to keep this.
}).await
let result =
upload_image_media_file_from_file(&app_env_configs.storyteller_host, Some(&creds), temp_file)
.await
.map_err(|err| {
error!("Failed to upload image media file: {:?}", err);
InternalBgRemovalError::StorytellerError(err)
@@ -43,7 +43,6 @@ use storyteller_client::media_files::get_media_file::get_media_file;
use storyteller_client::utils::api_host::ApiHost;
use tauri::{AppHandle, Manager, State};
use tokens::tokens::media_files::MediaFileToken;
use crate::core::commands::enqueue::image_edit::flux_kontext::handle_flux_kontext_edit::handle_flux_kontext_edit;
#[derive(Deserialize, Debug)]
pub struct EnqueueContextualEditImageCommand {
@@ -208,20 +207,6 @@ pub async fn handle_request(
None => {
return Err(InternalContextualEditImageError::NoModelSpecified)
}
Some(ContextualImageEditModel::FluxProKontextMax) => {
handle_flux_kontext_edit(
request,
app,
app_data_root,
app_env_configs,
provider_priority_store,
storyteller_creds_manager,
fal_creds_manager,
fal_task_queue,
sora_creds_manager,
sora_task_queue,
).await?
}
Some(ContextualImageEditModel::GptImage1) => {
handle_gpt_image_1_edit(
request,
@@ -236,6 +221,7 @@ pub async fn handle_request(
sora_task_queue,
).await?
}
// TODO(bt,2025-07-05): Flux Kontext, etc.
};
let result = success_event
@@ -30,12 +30,6 @@ pub enum InternalContextualEditImageError {
requested: u32,
},
// TODO: Maybe just reuse the other error enum variant
/// An invalid number of input images was requested for Flux Kontext.
InvalidNumberOfInputImagesForFluxKontext {
message: String,
},
/// No sora credentials are set.
NeedsSoraCredentials,
@@ -90,11 +84,6 @@ impl InternalContextualEditImageError {
error_type = EnqueueContextualEditImageErrorType::BadRequest;
error_message = format!("Invalid number of input images ({}). Must be between {} and {}", requested, min, max);
}
InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext { message } => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueContextualEditImageErrorType::BadRequest;
error_message = message.clone();
}
InternalContextualEditImageError::SoraError(_) => {}
InternalContextualEditImageError::AnyhowError(_) => {}
InternalContextualEditImageError::StorytellerError(_) => {}
@@ -1,59 +0,0 @@
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::EnqueueContextualEditImageCommand;
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
use crate::core::commands::enqueue::image_edit::flux_kontext::handle_flux_kontext_edit_artcraft::handle_flux_kontext_edit_artcraft;
use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_artcraft::handle_gpt_image_1_edit_artcraft;
use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_sora::handle_gpt_image_1_edit_sora;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
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::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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 log::info;
use tauri::AppHandle;
pub async fn handle_flux_kontext_edit(
request: &EnqueueContextualEditImageCommand,
app: &AppHandle,
app_data_root: &AppDataRoot,
app_env_configs: &AppEnvConfigs,
provider_priority_store: &ProviderPriorityStore,
storyteller_creds_manager: &StorytellerCredentialManager,
fal_creds_manager: &FalCredentialManager,
fal_task_queue: &FalTaskQueue,
sora_creds_manager: &SoraCredentialManager,
sora_task_queue: &SoraTaskQueue,
) -> Result<TaskEnqueueSuccess, InternalContextualEditImageError> {
let priority = provider_priority_store.get_priority()?;
// TODO: Check if providers are available before proceeding.
info!("Providers by priority: {:?}", priority);
for provider in priority.iter() {
match provider {
Provider::Fal => {
// Fallthrough (TODO)
}
Provider::Sora => {
// Fallthrough
}
Provider::Artcraft => {
info!("Dispatching gpt-image-1 (edit) via Artcraft...");
return handle_flux_kontext_edit_artcraft(
request,
app,
app_data_root,
app_env_configs,
storyteller_creds_manager
).await;
}
}
}
Err(InternalContextualEditImageError::NoProviderAvailable)
}
@@ -1,148 +0,0 @@
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::EnqueueTextToImageRequest;
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize, EnqueueContextualEditImageCommand};
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
use crate::core::model::contextual_image_edit_models::ContextualImageEditModel;
use crate::core::model::image_models::ImageModel;
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::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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 anyhow::anyhow;
use artcraft_api_defs::generate::image::edit::flux_pro_kontext_max_edit_image::{FluxProKontextMaxEditImageNumImages, FluxProKontextMaxEditImageRequest};
use artcraft_api_defs::generate::image::edit::gpt_image_1_edit_image::{GptImage1EditImageImageQuality, GptImage1EditImageImageSize, GptImage1EditImageNumImages, GptImage1EditImageRequest};
use artcraft_api_defs::generate::image::generate_flux_1_dev_text_to_image::GenerateFlux1DevTextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::GenerateFlux1SchnellTextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_pro_11_text_to_image::GenerateFluxPro11TextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_pro_11_ultra_text_to_image::GenerateFluxPro11UltraTextToImageRequest;
use enums::common::generation_provider::GenerationProvider;
use enums::tauri::tasks::task_type::TaskType;
use fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs};
use fal_client::requests::webhook::image::edit::enqueue_flux_pro_kontext_max_edit_webhook::FluxProKontextMaxNumImages;
use fal_client::requests::webhook::image::enqueue_gpt_image_1_edit_image_webhook::GptEditImageQuality;
use idempotency::uuid::generate_random_uuid;
use log::{error, info};
use storyteller_client::generate::image::edit::flux_pro_kontext_max_edit_image::flux_pro_kontext_max_edit_image;
use storyteller_client::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image;
use storyteller_client::generate::image::generate_flux_1_dev_text_to_image::generate_flux_1_dev_text_to_image;
use storyteller_client::generate::image::generate_flux_1_schnell_text_to_image::generate_flux_1_schnell_text_to_image;
use storyteller_client::generate::image::generate_flux_pro_11_text_to_image::generate_flux_pro_11_text_to_image;
use storyteller_client::generate::image::generate_flux_pro_11_ultra_text_to_image::generate_flux_pro_11_ultra_text_to_image;
use tauri::AppHandle;
pub async fn handle_flux_kontext_edit_artcraft(
request: &EnqueueContextualEditImageCommand,
app: &AppHandle,
app_data_root: &AppDataRoot,
app_env_configs: &AppEnvConfigs,
storyteller_creds_manager: &StorytellerCredentialManager,
) -> Result<TaskEnqueueSuccess, InternalContextualEditImageError> {
let creds = match storyteller_creds_manager.get_credentials()? {
Some(creds) => creds,
None => {
return Err(InternalContextualEditImageError::NeedsStorytellerCredentials);
},
};
info!("Calling Artcraft flux kontext max edit (edit) ...");
let uuid_idempotency_token = generate_random_uuid();
let num_images = match request.image_count {
None => None,
Some(1) => Some(FluxProKontextMaxEditImageNumImages::One),
Some(2) => Some(FluxProKontextMaxEditImageNumImages::Two),
Some(3) => Some(FluxProKontextMaxEditImageNumImages::Three),
Some(4) => Some(FluxProKontextMaxEditImageNumImages::Four),
Some(other) => {
return Err(InternalContextualEditImageError::InvalidNumberOfRequestedImages {
min: 1,
max: 4,
requested: other,
});
},
};
let mut maybe_media_token = None;
let list_count = request.image_media_tokens
.as_ref()
.map(|tokens| tokens.len())
.unwrap_or(0);
if request.scene_image_media_token.is_some() && list_count > 0 {
return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext {
message: "Cannot specify both scene_image_media_token and image_media_tokens for Flux Kontext Max".to_string()
});
}
if list_count > 1 {
return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext {
message: "Cannot specify more than one image in image_media_tokens for Flux Kontext Max".to_string()
});
}
if let Some(media_token) = &request.scene_image_media_token {
maybe_media_token = Some(media_token.clone());
}
let maybe_list_media_token = request.image_media_tokens
.as_ref()
.map(|tokens| tokens.first())
.flatten();
if let Some(media_token) = maybe_list_media_token {
maybe_media_token = Some(media_token.clone());
}
let media_token = match maybe_media_token {
Some(token) => token,
None => {
return Err(InternalContextualEditImageError::InvalidNumberOfInputImagesForFluxKontext {
message: "No input image specified for Flux Kontext Max edit".to_string(),
});
},
};
let request = FluxProKontextMaxEditImageRequest {
uuid_idempotency_token,
prompt: Some(request.prompt.clone()),
image_media_token: media_token,
num_images,
};
let result = flux_pro_kontext_max_edit_image(
&app_env_configs.storyteller_host,
Some(&creds),
request,
).await;
let job_id = match result {
Ok(enqueued) => {
// TODO(bt,2025-07-05): Enqueue job token?
info!("Successfully enqueued Artcraft gpt-image-1. Job token: {}",
enqueued.inference_job_token);
enqueued.inference_job_token
}
Err(err) => {
error!("Failed to use Artcraft gpt-image-1: {:?}", err);
return Err(InternalContextualEditImageError::StorytellerError(err));
}
};
Ok(TaskEnqueueSuccess {
provider: GenerationProvider::Artcraft,
model: Some(GenerationModel::FluxProKontextMax),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageGeneration,
})
}
@@ -1,2 +0,0 @@
pub (super) mod handle_flux_kontext_edit;
mod handle_flux_kontext_edit_artcraft;
@@ -1,4 +1,3 @@
pub mod enqueue_contextual_edit_image_command;
pub mod errors;
pub mod gpt_image_1;
mod flux_kontext;
pub mod gpt_image_1;
@@ -1,249 +0,0 @@
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, EnqueueTextToImageSuccessResponse};
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit::handle_gpt_image_1_edit;
use crate::core::commands::enqueue::image_inpaint::errors::InternalImageInpaintError;
use crate::core::commands::enqueue::image_inpaint::flux_pro_1_inpaint::handle_flux_pro_1_inpaint::handle_flux_pro_1_inpaint;
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, ResponseOrErrorType};
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
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_enqueue_failure_event::GenerationEnqueueFailureEvent;
use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent;
use crate::core::events::sendable_event_trait::SendableEvent;
use crate::core::model::contextual_image_edit_models::ContextualImageEditModel;
use crate::core::model::image_models::ImageModel;
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::provider_priority::ProviderPriorityStore;
use crate::core::state::task_database::TaskDatabase;
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::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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 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 errors::AnyhowError;
use log::{error, info, warn};
use openai_sora_client::recipes::image_remix_with_session_auto_renew::{image_remix_with_session_auto_renew, ImageRemixAutoRenewRequest};
use openai_sora_client::recipes::image_upload_from_file_with_session_auto_renew::{image_upload_from_file_with_session_auto_renew, ImageUploadFromFileAutoRenewRequest};
use openai_sora_client::recipes::maybe_upgrade_or_renew_session::maybe_upgrade_or_renew_session;
use openai_sora_client::requests::image_gen::common::{ImageSize, NumImages};
use openai_sora_client::sora_error::SoraError;
use serde_derive::{Deserialize, Serialize};
use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs};
use std::time::Duration;
use storyteller_client::error::storyteller_error::StorytellerError;
use storyteller_client::media_files::get_media_file::get_media_file;
use storyteller_client::utils::api_host::ApiHost;
use tauri::{AppHandle, Manager, State};
use tokens::tokens::media_files::MediaFileToken;
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ImageInpaintModel {
#[serde(rename = "flux_pro_1")]
FluxPro1,
}
#[derive(Deserialize, Debug)]
pub struct EnqueueInpaintImageCommand {
/// REQUIRED (Option<T> is just for error messages).
/// The model to use.
pub model: Option<ImageInpaintModel>,
/// REQUIRED (Option<T> is just for error messages).
/// The source image to edit.
pub image_media_token: Option<MediaFileToken>,
/// REQUIRED: Supply this *XOR* `mask_image_raw_bytes`.
/// The mask to focus the edit (already uploaded).
pub mask_image_media_token: Option<MediaFileToken>,
/// REQUIRED: Supply this *XOR* `mask_image_media_token`.
/// The mask to focus the edit (raw bytes).
pub mask_image_raw_bytes: Option<Vec<u8>>,
/// REQUIRED.
/// The user's image generation prompt.
pub prompt: String,
/// Number of images to generate.
pub image_count: Option<u32>,
/// If true, force the dimensions of the source image and mask image to match.
pub require_matching_dimensions: Option<bool>,
/// 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, Debug)]
#[serde(rename_all = "snake_case")]
pub enum EnqueueInpaintImageErrorType {
/// Caller didn't specify a model
ModelNotSpecified,
/// No model available for image generation
NoProviderAvailable,
/// No source image was supplied.
NoSourceImageSpecified,
/// No mask image was supplied.
NoMaskImageSpecified,
/// Too many mask images were supplied.
MultipleMaskImagesSpecified,
/// Bad mask image was supplied.
BadMaskImage,
/// Generic bad request error
BadRequest,
/// Generic server error
ServerError,
}
#[derive(Serialize)]
pub struct EnqueueImageInpaintSuccessResponse {
}
impl SerializeMarker for EnqueueImageInpaintSuccessResponse {}
#[tauri::command]
pub async fn enqueue_image_inpaint_command(
app: AppHandle,
request: EnqueueInpaintImageCommand,
app_data_root: State<'_, AppDataRoot>,
app_env_configs: State<'_, AppEnvConfigs>,
provider_priority_store: State<'_, ProviderPriorityStore>,
task_database: State<'_, TaskDatabase>,
storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
fal_creds_manager: State<'_, FalCredentialManager>,
fal_task_queue: State<'_, FalTaskQueue>,
sora_creds_manager: State<'_, SoraCredentialManager>,
sora_task_queue: State<'_, SoraTaskQueue>,
) -> ResponseOrErrorType<EnqueueImageInpaintSuccessResponse, EnqueueInpaintImageErrorType> {
info!("enqueue_image_inpaint_command called; model: {:?}, image_media_token: {:?}, prompt: {:?}",
&request.model, &request.image_media_token, &request.prompt);
let result = handle_request(
&request,
&app,
&app_data_root,
&app_env_configs,
&provider_priority_store,
&task_database,
&storyteller_creds_manager,
&fal_creds_manager,
&fal_task_queue,
&sora_creds_manager,
&sora_task_queue,
).await;
match result {
Err(err) => {
error!("Error enqueuing inpaint image: {:?}", err);
// TODO: Derive from err. Make service provider optional.
let event = GenerationEnqueueFailureEvent {
action: GenerationAction::GenerateImage,
service: GenerationServiceProvider::Artcraft, // FIXME: This is wrong.
model: None,
reason: None,
};
if let Err(err) = event.send(&app) {
error!("Failed to emit event: {:?}", err); // Fail open.
}
Err(err.to_tauri_response())
}
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.
}
Ok(EnqueueImageInpaintSuccessResponse {}.into())
}
}
}
pub async fn handle_request(
request: &EnqueueInpaintImageCommand,
app: &AppHandle,
app_data_root: &AppDataRoot,
app_env_configs: &AppEnvConfigs,
provider_priority_store: &ProviderPriorityStore,
task_database: &TaskDatabase,
storyteller_creds_manager: &StorytellerCredentialManager,
fal_creds_manager: &FalCredentialManager,
fal_task_queue: &FalTaskQueue,
sora_creds_manager: &SoraCredentialManager,
sora_task_queue: &SoraTaskQueue,
) -> Result<TaskEnqueueSuccess, InternalImageInpaintError> {
let success_event= match request.model {
None => {
return Err(InternalImageInpaintError::NoModelSpecified)
}
Some(ImageInpaintModel::FluxPro1) => {
handle_flux_pro_1_inpaint(
request,
app,
app_data_root,
app_env_configs,
provider_priority_store,
storyteller_creds_manager,
fal_creds_manager,
fal_task_queue,
sora_creds_manager,
sora_task_queue,
).await?
}
};
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)
}
@@ -1,128 +0,0 @@
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::EnqueueInpaintImageErrorType;
use crate::core::commands::response::failure_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus};
use errors::AnyhowError;
use serde::Serialize;
use storyteller_client::error::storyteller_error::StorytellerError;
#[derive(Debug)]
pub enum InternalImageInpaintError {
/// Request did not specify a model.
NoModelSpecified,
/// No provider is configured for the model type.
NoProviderAvailable,
/// A source image was not provided
NoSourceImageSpecified,
/// A mask image was not provided
NoMaskImageSpecified,
/// Both mask image media token and raw bytes were supplied.
MaskMediaTokenAndBytesSupplied,
/// Mask image has an unknown MIME type.
CouldNotDetermineMaskMimeType,
/// Could not encode the mask image to PNG.
CouldNotEncodeMask,
/// No storyteller credentials are set.
NeedsStorytellerCredentials,
/// An invalid number of images was requested.
InvalidNumberOfRequestedImages {
min: u32,
max: u32,
requested: u32,
},
/// An invalid number of input images was requested.
InvalidNumberOfInputImages {
min: u32,
max: u32,
requested: u32,
},
AnyhowError(AnyhowError),
StorytellerError(StorytellerError),
}
impl From<AnyhowError> for InternalImageInpaintError {
fn from(value: AnyhowError) -> Self {
Self::AnyhowError(value)
}
}
impl From<StorytellerError> for InternalImageInpaintError {
fn from(value: StorytellerError) -> Self {
Self::StorytellerError(value)
}
}
impl InternalImageInpaintError {
pub fn to_tauri_response<E: Serialize>(&self) -> CommandErrorResponseWrapper<EnqueueInpaintImageErrorType, E> {
let mut status = CommandErrorStatus::ServerError;
let mut error_type = EnqueueInpaintImageErrorType::ServerError;
let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem.".to_string();
match self {
InternalImageInpaintError::NoModelSpecified => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::ModelNotSpecified;
error_message = "No model specified for image generation".to_string();
}
InternalImageInpaintError::NoProviderAvailable => {
status = CommandErrorStatus::ServerError;
error_type = EnqueueInpaintImageErrorType::NoProviderAvailable;
error_message = "No configured provider available for image generation".to_string();
}
InternalImageInpaintError::NoSourceImageSpecified => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::NoSourceImageSpecified;
error_message = "No source image was provided".to_string();
}
InternalImageInpaintError::NoMaskImageSpecified => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::NoMaskImageSpecified;
error_message = "No mask image was provided".to_string();
}
InternalImageInpaintError::MaskMediaTokenAndBytesSupplied => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::MultipleMaskImagesSpecified;
error_message = "multiple mask images provided".to_string();
}
InternalImageInpaintError::CouldNotDetermineMaskMimeType => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::BadMaskImage;
error_message = "bad mask image (mime)".to_string();
}
InternalImageInpaintError::CouldNotEncodeMask => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::BadMaskImage;
error_message = "bad mask image (encode)".to_string();
}
InternalImageInpaintError::InvalidNumberOfRequestedImages { min, max, requested } => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::BadRequest;
error_message = format!("Invalid number of images requested ({}). Must be between {} and {}", requested, min, max);
}
InternalImageInpaintError::InvalidNumberOfInputImages { min, max, requested } => {
status = CommandErrorStatus::BadRequest;
error_type = EnqueueInpaintImageErrorType::BadRequest;
error_message = format!("Invalid number of input images ({}). Must be between {} and {}", requested, min, max);
}
InternalImageInpaintError::AnyhowError(_) => {}
InternalImageInpaintError::StorytellerError(_) => {}
InternalImageInpaintError::NeedsStorytellerCredentials => {}
}
CommandErrorResponseWrapper {
status,
error_message: Some(error_message.to_string()),
error_type: Some(error_type),
error_details: None,
}
}
}
@@ -1,63 +0,0 @@
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::EnqueueContextualEditImageCommand;
use crate::core::commands::enqueue::image_edit::errors::InternalContextualEditImageError;
use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_artcraft::handle_gpt_image_1_edit_artcraft;
use crate::core::commands::enqueue::image_edit::gpt_image_1::handle_gpt_image_1_edit_sora::handle_gpt_image_1_edit_sora;
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::EnqueueInpaintImageCommand;
use crate::core::commands::enqueue::image_inpaint::errors::InternalImageInpaintError;
use crate::core::commands::enqueue::image_inpaint::flux_pro_1_inpaint::handle_flux_pro_1_inpaint_artcraft::handle_flux_pro_1_inpaint_artcraft;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
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::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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 log::info;
use tauri::AppHandle;
pub(super) const MAX_IMAGES: usize = 10;
pub async fn handle_flux_pro_1_inpaint(
request: &EnqueueInpaintImageCommand,
app: &AppHandle,
app_data_root: &AppDataRoot,
app_env_configs: &AppEnvConfigs,
provider_priority_store: &ProviderPriorityStore,
storyteller_creds_manager: &StorytellerCredentialManager,
fal_creds_manager: &FalCredentialManager,
fal_task_queue: &FalTaskQueue,
sora_creds_manager: &SoraCredentialManager,
sora_task_queue: &SoraTaskQueue,
) -> Result<TaskEnqueueSuccess, InternalImageInpaintError> {
let priority = provider_priority_store.get_priority()?;
// TODO: Check if providers are available before proceeding.
info!("Providers by priority: {:?}", priority);
for provider in priority.iter() {
match provider {
Provider::Fal => {
// Fallthrough (for now)
}
Provider::Sora => {
// Fallthrough
}
Provider::Artcraft => {
info!("Dispatching flux pro 1 inpaint via Artcraft...");
return handle_flux_pro_1_inpaint_artcraft(
request,
app,
app_data_root,
app_env_configs,
storyteller_creds_manager
).await;
}
}
}
Err(InternalImageInpaintError::NoProviderAvailable)
}
@@ -1,181 +0,0 @@
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::EnqueueTextToImageRequest;
use crate::core::commands::enqueue::image::internal_image_error::InternalImageError;
use crate::core::commands::enqueue::image_bg_removal::enqueue_image_bg_removal_command::EnqueueImageBgRemovalCommand;
use crate::core::commands::enqueue::image_bg_removal::errors::InternalBgRemovalError;
use crate::core::commands::enqueue::image_edit::enqueue_contextual_edit_image_command::{EditImageQuality, EditImageSize, EnqueueContextualEditImageCommand};
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::EnqueueInpaintImageCommand;
use crate::core::commands::enqueue::image_inpaint::errors::InternalImageInpaintError;
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
use crate::core::model::contextual_image_edit_models::ContextualImageEditModel;
use crate::core::model::image_models::ImageModel;
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::core::utils::save_base64_image_to_temp_dir::save_base64_image_to_temp_dir;
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
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 anyhow::anyhow;
use artcraft_api_defs::generate::image::edit::gpt_image_1_edit_image::{GptImage1EditImageImageQuality, GptImage1EditImageImageSize, GptImage1EditImageNumImages, GptImage1EditImageRequest};
use artcraft_api_defs::generate::image::generate_flux_1_dev_text_to_image::GenerateFlux1DevTextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_1_schnell_text_to_image::GenerateFlux1SchnellTextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_pro_11_text_to_image::GenerateFluxPro11TextToImageRequest;
use artcraft_api_defs::generate::image::generate_flux_pro_11_ultra_text_to_image::GenerateFluxPro11UltraTextToImageRequest;
use artcraft_api_defs::generate::image::inpaint::flux_pro_1_inpaint_image::{FluxPro1InpaintImageNumImages, FluxPro1InpaintImageRequest};
use enums::common::generation_provider::GenerationProvider;
use enums::tauri::tasks::task_type::TaskType;
use fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs};
use fal_client::requests::webhook::image::enqueue_gpt_image_1_edit_image_webhook::GptEditImageQuality;
use idempotency::uuid::generate_random_uuid;
use images::encoding::image_bytes_to_png_bytes::image_bytes_to_png_bytes;
use images::encoding::image_bytes_to_png_bytes_with_dimensions::image_bytes_to_png_bytes_with_dimensions;
use images::mask_images::normalize_image_bytes_to_flux_mask::normalize_image_bytes_to_flux_mask;
use log::{error, info};
use mimetypes::mimetype_info::mimetype_info::MimetypeInfo;
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use storyteller_client::generate::image::edit::gpt_image_1_edit_image::gpt_image_1_edit_image;
use storyteller_client::generate::image::generate_flux_1_dev_text_to_image::generate_flux_1_dev_text_to_image;
use storyteller_client::generate::image::generate_flux_1_schnell_text_to_image::generate_flux_1_schnell_text_to_image;
use storyteller_client::generate::image::generate_flux_pro_11_text_to_image::generate_flux_pro_11_text_to_image;
use storyteller_client::generate::image::generate_flux_pro_11_ultra_text_to_image::generate_flux_pro_11_ultra_text_to_image;
use storyteller_client::generate::image::inpaint::flux_pro_1_inpaint_image::flux_pro_1_inpaint_image;
use storyteller_client::media_files::get_media_file::get_media_file;
use storyteller_client::media_files::upload_image_media_file_from_bytes::{upload_image_media_file_from_bytes, ImageType, UploadImageBytesArgs};
use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
use tauri::AppHandle;
use tokens::tokens::media_files::MediaFileToken;
pub async fn handle_flux_pro_1_inpaint_artcraft(
request: &EnqueueInpaintImageCommand,
app: &AppHandle,
app_data_root: &AppDataRoot,
app_env_configs: &AppEnvConfigs,
storyteller_creds_manager: &StorytellerCredentialManager,
) -> Result<TaskEnqueueSuccess, InternalImageInpaintError> {
let creds = match storyteller_creds_manager.get_credentials()? {
Some(creds) => creds,
None => {
return Err(InternalImageInpaintError::NeedsStorytellerCredentials);
},
};
let image_media_token = match &request.image_media_token {
Some(token) => token.clone(),
None => {
return Err(InternalImageInpaintError::NoSourceImageSpecified);
},
};
//let image_media_file = get_media_file(
// &app_env_configs.storyteller_host,
// &image_media_token,
//).await.map_err(|err| {
// error!("Failed to get media file for image: {:?}", err);
// InternalImageInpaintError::StorytellerError(err)
//})?;
let mask_media_token = get_mask(
request,
app_env_configs,
&creds,
).await?;
info!("Calling Artcraft flux pro 1 inpaint ...");
let uuid_idempotency_token = generate_random_uuid();
let num_images = match request.image_count {
None => None,
Some(1) => Some(FluxPro1InpaintImageNumImages::One),
Some(2) => Some(FluxPro1InpaintImageNumImages::Two),
Some(3) => Some(FluxPro1InpaintImageNumImages::Three),
Some(4) => Some(FluxPro1InpaintImageNumImages::Four),
Some(other) => {
return Err(InternalImageInpaintError::InvalidNumberOfRequestedImages {
min: 1,
max: 4,
requested: other,
});
},
};
let request = FluxPro1InpaintImageRequest {
uuid_idempotency_token,
prompt: Some(request.prompt.clone()),
image_media_token,
mask_media_token,
num_images,
};
let result = flux_pro_1_inpaint_image(
&app_env_configs.storyteller_host,
Some(&creds),
request,
).await;
let job_id = match result {
Ok(enqueued) => {
// TODO(bt,2025-07-05): Enqueue job token?
info!("Successfully enqueued Artcraft flux pro 1 inpaint. Job token: {}",
enqueued.inference_job_token);
enqueued.inference_job_token
}
Err(err) => {
error!("Failed to use Artcraft flux pro 1 inpaint: {:?}", err);
return Err(InternalImageInpaintError::StorytellerError(err));
}
};
Ok(TaskEnqueueSuccess {
provider: GenerationProvider::Artcraft,
model: Some(GenerationModel::FluxProKontextMax),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageGeneration,
})
}
async fn get_mask(
request: &EnqueueInpaintImageCommand,
app_env_configs: &AppEnvConfigs,
storyteller_creds: &StorytellerCredentialSet,
) -> Result<MediaFileToken, InternalImageInpaintError> {
if request.mask_image_media_token.is_some() && request.mask_image_raw_bytes.is_some() {
return Err(InternalImageInpaintError::MaskMediaTokenAndBytesSupplied);
}
if let Some(token) = request.mask_image_media_token.as_ref() {
return Ok(token.clone());
};
let image_bytes = request.mask_image_raw_bytes.as_ref()
.ok_or(InternalImageInpaintError::NoMaskImageSpecified)?;
let image_bytes = normalize_image_bytes_to_flux_mask(image_bytes)
.map_err(|err| {
error!("Failed to convert image bytes to png: {:?}", err);
InternalImageInpaintError::CouldNotEncodeMask
})?;
info!("Uploading image media file from bytes...");
let result = upload_image_media_file_from_bytes(UploadImageBytesArgs {
api_host: &app_env_configs.storyteller_host,
maybe_creds: Some(&storyteller_creds),
image_bytes: image_bytes.0,
image_type: ImageType::Png,
is_intermediate_system_file: true,
}).await
.map_err(|err| {
error!("Failed to upload image media file: {:?}", err);
InternalImageInpaintError::StorytellerError(err)
})?;
Ok(result.media_file_token)
}
@@ -1,2 +0,0 @@
pub (super) mod handle_flux_pro_1_inpaint;
mod handle_flux_pro_1_inpaint_artcraft;
@@ -1,3 +0,0 @@
pub mod enqueue_image_inpaint_command;
pub mod errors;
mod flux_pro_1_inpaint;
@@ -1,7 +1,6 @@
pub mod image;
pub mod image_bg_removal;
pub mod image_edit;
pub mod image_inpaint;
pub mod object;
pub mod task_enqueue_success;
pub mod object;
pub mod video;
@@ -4,7 +4,6 @@ use enums::common::generation_provider::GenerationProvider;
use enums::tauri::tasks::task_model_type::TaskModelType;
use enums::tauri::tasks::task_status::TaskStatus;
use enums::tauri::tasks::task_type::TaskType;
use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
use sqlite_tasks::error::SqliteTasksError;
use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs};
use tokens::tokens::sqlite::tasks::TaskId;
@@ -31,37 +30,18 @@ impl TaskEnqueueSuccess{
match self.provider {
GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft,
GenerationProvider::Fal => GenerationServiceProvider::Fal,
GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney,
GenerationProvider::Sora => GenerationServiceProvider::Sora,
}
}
pub async fn insert_into_task_database(&self, task_database: &TaskDatabase) -> Result<TaskId, SqliteTasksError> {
self.insert_into_task_database_with_frontend_payload(
task_database,
None,
None,
None,
).await
}
// TODO: This belongs somewhere else, not as a method of an event struct.
pub async fn insert_into_task_database_with_frontend_payload(
&self,
task_database: &TaskDatabase,
frontend_caller: Option<TauriCommandCaller>,
frontend_subscriber_id: Option<&str>,
frontend_subscriber_payload: Option<&str>,
) -> Result<TaskId, SqliteTasksError> {
// TODO: Move this mapping elsewhere, or remove the other models.
let model_type = match self.model {
None => None,
Some(GenerationModel::Flux1Dev) => Some(TaskModelType::Flux1Dev),
Some(GenerationModel::Flux1Schnell) => Some(TaskModelType::Flux1Schnell),
Some(GenerationModel::FluxPro1) => Some(TaskModelType::FluxPro1), // NB: This is for inpainting.
Some(GenerationModel::FluxPro11) => Some(TaskModelType::FluxPro11),
Some(GenerationModel::FluxPro11Ultra) => Some(TaskModelType::FluxPro11Ultra),
Some(GenerationModel::FluxProKontextMax) => Some(TaskModelType::FluxProKontextMax),
Some(GenerationModel::GptImage1) => Some(TaskModelType::GptImage1),
Some(GenerationModel::Recraft3) => Some(TaskModelType::Recraft3),
Some(GenerationModel::Kling21Pro) => Some(TaskModelType::Kling21Pro),
@@ -70,7 +50,6 @@ impl TaskEnqueueSuccess{
Some(GenerationModel::Veo2) => Some(TaskModelType::Veo2),
Some(GenerationModel::Hunyuan3d2_0) => Some(TaskModelType::Hunyuan3d2_0),
Some(GenerationModel::Hunyuan3d2_1) => Some(TaskModelType::Hunyuan3d2_1),
Some(GenerationModel::Midjourney) => Some(TaskModelType::Midjourney), // NB: This is a generic Midjourney model, version unknown.
// TODO: These seem wrong -
Some(GenerationModel::Kling1_6) => Some(TaskModelType::Kling16Pro), // NB: `VideoModel::Kling16Pro`.
@@ -82,12 +61,23 @@ impl TaskEnqueueSuccess{
db: task_database.get_connection(),
status: TaskStatus::Pending,
task_type: self.task_type,
model_type,
model_type: model_type,
provider: self.provider,
provider_job_id: self.provider_job_id.as_deref(),
frontend_caller,
frontend_subscriber_id,
frontend_subscriber_payload,
frontend_subscriber_id: None,
frontend_subscriber_payload: None,
}).await
}
// pub fn tauri_event_model(&self) -> GenerationModel {
// match self.model {
// ImageModel::Flux1Dev => GenerationModel::Flux1Dev,
// ImageModel::Flux1Schnell => GenerationModel::Flux1Schnell,
// ImageModel::FluxPro11 => GenerationModel::FluxPro11,
// ImageModel::FluxPro11Ultra => GenerationModel::FluxPro11Ultra,
// ImageModel::GptImage1 => GenerationModel::GptImage1,
// ImageModel::Recraft3 => GenerationModel::Recraft3,
// }
// }
}
@@ -1,5 +1,4 @@
use crate::core::events::sendable_event_error::SendableEventError;
use enums::tauri::ux::tauri_event_name::TauriEventName;
use log::info;
use serde::Serialize;
use std::fmt::Debug;
@@ -30,7 +29,7 @@ struct BasicEventWrapper<T: Serialize + Debug> {
pub trait BasicSendableEvent : Clone + Debug + Serialize {
/// This is the name of the event that the frontend subscribes to.
const FRONTEND_EVENT_NAME: TauriEventName;
const FRONTEND_EVENT_NAME: &'static str;
/// This is the status of the event: Success, Failure, etc.
const EVENT_STATUS: BasicEventStatus;
@@ -42,8 +41,8 @@ pub trait BasicSendableEvent : Clone + Debug + Serialize {
status: Self::EVENT_STATUS,
data: self.clone(),
};
info!("Sending event to frontend: {} ; payload = {:?}", Self::FRONTEND_EVENT_NAME.to_str(), wrapped);
let result = app.emit(Self::FRONTEND_EVENT_NAME.to_str(), wrapped);
info!("Sending event to frontend: {} ; payload = {:?}", Self::FRONTEND_EVENT_NAME, wrapped);
let result = app.emit(Self::FRONTEND_EVENT_NAME, wrapped);
if let Err(err) = result {
return Err(SendableEventError::from(err));
}
@@ -1,19 +0,0 @@
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 CanvasBackgroundRemovalCompleteEvent {
pub media_token: MediaFileToken,
pub image_cdn_url: Url,
pub maybe_frontend_subscriber_id: Option<String>,
pub maybe_frontend_subscriber_payload: Option<String>,
}
impl BasicSendableEvent for CanvasBackgroundRemovalCompleteEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::CanvasBgRemovedEvent;
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
}
@@ -1 +0,0 @@
pub mod canvas_background_removal_complete_event;
@@ -5,7 +5,6 @@ use serde::Serialize;
pub enum GenerationServiceProvider {
Artcraft,
Fal,
Midjourney,
Sora,
}
@@ -16,14 +15,10 @@ pub enum GenerationModel {
Flux1Dev,
#[serde(rename = "flux_1_schnell")]
Flux1Schnell,
#[serde(rename = "flux_pro_1")]
FluxPro1,
#[serde(rename = "flux_pro_1.1")]
FluxPro11,
#[serde(rename = "flux_pro_1.1_ultra")]
FluxPro11Ultra,
#[serde(rename = "flux_pro_kontext_max")]
FluxProKontextMax,
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "hunyuan_3d_2_0")]
@@ -31,10 +26,6 @@ pub enum GenerationModel {
#[serde(rename = "hunyuan_3d_2_1")]
Hunyuan3d2_1,
// Generic Midjourney model, version unknown.
#[serde(rename = "midjourney")]
Midjourney,
// TODO: Should be Kling16Pro
#[serde(rename = "kling_1.6")]
Kling1_6,
@@ -1,6 +1,5 @@
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use enums::tauri::ux::tauri_event_name::TauriEventName;
use serde_derive::Serialize;
#[derive(Clone, Debug, Serialize)]
@@ -12,6 +11,6 @@ pub struct GenerationCompleteEvent {
}
impl BasicSendableEvent for GenerationCompleteEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GenerationCompleteEvent;
const FRONTEND_EVENT_NAME: &'static str = "generation-complete-event";
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
}
@@ -1,6 +1,5 @@
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use enums::tauri::ux::tauri_event_name::TauriEventName;
use serde_derive::Serialize;
#[derive(Clone, Debug, Serialize)]
@@ -15,7 +14,7 @@ pub struct GenerationEnqueueFailureEvent {
}
impl BasicSendableEvent for GenerationEnqueueFailureEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GenerationEnqueueFailureEvent;
const FRONTEND_EVENT_NAME: &'static str = "generation-enqueue-failure-event";
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Failure;
}
@@ -1,6 +1,5 @@
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use enums::tauri::ux::tauri_event_name::TauriEventName;
use serde_derive::Serialize;
#[derive(Clone, Debug, Serialize)]
@@ -12,6 +11,6 @@ pub struct GenerationEnqueueSuccessEvent {
}
impl BasicSendableEvent for GenerationEnqueueSuccessEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GenerationEnqueueSuccessEvent;
const FRONTEND_EVENT_NAME: &'static str = "generation-enqueue-success-event";
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Success;
}
@@ -1,6 +1,5 @@
use crate::core::events::basic_sendable_event_trait::{BasicEventStatus, BasicSendableEvent};
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use enums::tauri::ux::tauri_event_name::TauriEventName;
use serde_derive::Serialize;
#[derive(Clone, Debug, Serialize)]
@@ -15,6 +14,6 @@ pub struct GenerationFailedEvent {
}
impl BasicSendableEvent for GenerationFailedEvent {
const FRONTEND_EVENT_NAME: TauriEventName = TauriEventName::GenerationFailedEvent;
const FRONTEND_EVENT_NAME: &'static str = "generation-failed-event";
const EVENT_STATUS: BasicEventStatus = BasicEventStatus::Failure;
}
@@ -1,5 +1,4 @@
pub mod basic_sendable_event_trait;
pub mod generation_events;
pub mod sendable_event_error;
pub mod sendable_event_trait;
pub mod functional_events;
pub mod sendable_event_trait;
@@ -1,7 +1,7 @@
use crate::core::lifecycle::startup::tasks::bootstrap_task_database::bootstrap_task_database;
use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
use crate::core::lifecycle::startup::tasks::load_provider_priority_state::load_provider_priority_state;
use crate::core::lifecycle::startup::tasks::set_app_log_level::set_app_log_level;
use crate::core::lifecycle::startup::tasks::initially_size_and_position_windows::initially_size_and_position_windows;
use crate::core::lifecycle::startup::tasks::spawn_discord_presence_thread::spawn_discord_presence_thread;
use crate::core::lifecycle::startup::tasks::spawn_fal_task_polling_thread::spawn_fal_task_polling_thread;
use crate::core::lifecycle::startup::tasks::spawn_main_window_thread::spawn_main_window_thread;
@@ -11,8 +11,6 @@ use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use crate::services::midjourney::threads::midjourney_long_polling_thread::midjourney_long_polling_thread;
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;
@@ -28,7 +26,6 @@ pub async fn handle_tauri_startup(
sora_task_queue: SoraTaskQueue,
fal_credential_manager: FalCredentialManager,
fal_task_queue: FalTaskQueue,
mj_creds_manager: MidjourneyCredentialManager,
) -> AnyhowResult<()> {
set_app_log_level(
@@ -75,15 +72,6 @@ pub async fn handle_tauri_startup(
&storyteller_creds_manager,
)?;
tauri::async_runtime::spawn(midjourney_long_polling_thread(
app.clone(),
app_env_configs.clone(),
root.clone(),
task_database.clone(),
mj_creds_manager.clone(),
storyteller_creds_manager.clone(),
));
spawn_discord_presence_thread()?;
initially_size_and_position_windows(&app, &root);
@@ -7,9 +7,6 @@ use serde_derive::Deserialize;
pub enum ContextualImageEditModel {
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "flux_pro_kontext_max")]
FluxProKontextMax,
// TODO: Flux Kontext models.
}
@@ -17,8 +17,4 @@ pub enum ImageModel {
GptImage1,
#[serde(rename = "recraft_3")]
Recraft3,
// Generic Midjourney model, version unknown.
#[serde(rename = "midjourney")]
Midjourney,
}
@@ -44,8 +44,4 @@ impl AppCredentialsDir {
pub fn get_fal_api_key_file_path(&self) -> PathBuf {
self.path.join("fal_api_key.txt")
}
pub fn get_midjourney_state_path(&self) -> PathBuf {
self.path.join("midjourney_state.json")
}
}
@@ -22,9 +22,7 @@ impl DataSubdir for AppStateDir{
impl AppStateDir {
pub fn get_tasks_sqlite_database_path(&self) -> PathBuf {
// NB: This must be incremented on any change to the database schema (including comments).
// The database houses ephemeral content and will be migrated automatically.
self.path.join("tasks_v3.sqlite")
self.path.join("tasks_v2.sqlite")
}
pub fn get_window_size_config_file(&self) -> PathBuf {
@@ -6,7 +6,6 @@ pub fn to_generation_service_provider(provider: GenerationProvider) -> Generatio
match provider {
GenerationProvider::Artcraft => GenerationServiceProvider::Artcraft,
GenerationProvider::Fal => GenerationServiceProvider::Fal,
GenerationProvider::Midjourney => GenerationServiceProvider::Midjourney,
GenerationProvider::Sora => GenerationServiceProvider::Sora,
}
}
+1 -15
View File
@@ -8,7 +8,6 @@ use crate::core::commands::app_preferences::update_app_preference_command::updat
use crate::core::commands::enqueue::image::enqueue_text_to_image_command::enqueue_text_to_image_command;
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_contextual_edit_image_command::enqueue_contextual_edit_image_command;
use crate::core::commands::enqueue::image_inpaint::enqueue_image_inpaint_command::enqueue_image_inpaint_command;
use crate::core::commands::enqueue::object::enqueue_image_to_3d_object_command::enqueue_image_to_3d_object_command;
use crate::core::commands::enqueue::video::enqueue_image_to_video_command::enqueue_image_to_video_command;
use crate::core::commands::flip_image::flip_image;
@@ -32,10 +31,6 @@ use crate::services::fal::commands::set_fal_api_key_command::set_fal_api_key_com
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
use crate::services::fal::threads::fal_task_polling_thread::fal_task_polling_thread;
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;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use crate::services::sora::commands::check_sora_session_command::check_sora_session_command;
use crate::services::sora::commands::open_sora_login_command::open_sora_login_command;
use crate::services::sora::commands::sora_image_generation_command::sora_image_generation_command;
@@ -89,10 +84,7 @@ pub fn run() {
.expect("AppEnvConfigs should be loaded from disk");
let app_env_configs_2 = app_env_configs.clone();
let midjourney_creds_manager = MidjourneyCredentialManager::initialize_from_disk_infallible(&app_data_root);
let midjourney_creds_manager_2 = midjourney_creds_manager.clone();
println!("Initializing backend runtime...");
tauri::Builder::default()
@@ -130,7 +122,6 @@ pub fn run() {
sora_tasks,
fal_creds,
fal_tasks,
midjourney_creds_manager_2,
).await;
if let Err(err) = result {
@@ -146,7 +137,6 @@ pub fn run() {
.manage(app_preferences)
.manage(fal_creds_manager)
.manage(fal_task_queue)
.manage(midjourney_creds_manager)
.manage(sora_creds_manager)
.manage(sora_task_queue)
.manage(storyteller_creds_manager_3)
@@ -154,7 +144,6 @@ pub fn run() {
check_sora_session_command,
enqueue_contextual_edit_image_command,
enqueue_image_bg_removal_command,
enqueue_image_inpaint_command,
enqueue_image_to_3d_object_command,
enqueue_image_to_video_command,
enqueue_text_to_image_command,
@@ -168,9 +157,6 @@ pub fn run() {
get_provider_order_command,
load_without_cors_command,
open_sora_login_command,
midjourney_clear_credentials_command,
midjourney_get_credential_info_command,
midjourney_open_login_command,
platform_info_command,
set_fal_api_key_command,
set_provider_order_command,
@@ -29,7 +29,7 @@ use reqwest::Url;
use serde_derive::{Deserialize, Serialize};
use storyteller_client::error::storyteller_error::StorytellerError;
use storyteller_client::media_files::get_media_file::{get_media_file, GetMediaFileSuccessResponse};
use storyteller_client::media_files::upload_image_media_file_from_file::{upload_image_media_file_from_file, UploadImageFromFileArgs};
use storyteller_client::media_files::upload_image_media_file_from_file::upload_image_media_file_from_file;
use storyteller_client::utils::api_host::ApiHost;
use tauri::{AppHandle, Manager, State};
use tokens::tokens::media_files::MediaFileToken;
@@ -256,12 +256,11 @@ pub async fn remove_background(
info!("Uploading image media file...");
let upload_result = upload_image_media_file_from_file(UploadImageFromFileArgs {
api_host: &app_env_configs.storyteller_host,
maybe_creds: Some(&creds),
path: result_filename,
is_intermediate_system_file: true, // NB: Probably not essential to keep this.
}).await?;
let upload_result = upload_image_media_file_from_file(
&app_env_configs.storyteller_host,
Some(&creds),
result_filename,
).await?;
// TODO: Don't re-request to simply build MediaLinks (or CDN URL). Get those from the upload API in one turn.
info!("Re-requesting media file...");
@@ -1,41 +0,0 @@
use crate::core::commands::response::shorthand::{Response, SimpleResponse};
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use errors::AnyhowResult;
use fal_client::creds::fal_api_key::FalApiKey;
use log::{error, info};
use serde_derive::Deserialize;
use tauri::State;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
#[tauri::command]
pub async fn midjourney_clear_credentials_command(
root: State<'_, AppDataRoot>,
creds_manager: State<'_, MidjourneyCredentialManager>,
) -> SimpleResponse {
info!("midjourney_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: &MidjourneyCredentialManager,
) -> AnyhowResult<()> {
creds.clear_credentials()?;
creds.persist_to_disk()?;
let creds_path = root.credentials_dir().get_midjourney_state_path();
if creds_path.exists() {
std::fs::remove_file(creds_path)?;
}
Ok(())
}
@@ -1,53 +0,0 @@
use crate::core::commands::response::shorthand::{Response, ResponseOrErrorMessage};
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
use crate::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use errors::AnyhowResult;
use log::{error, info};
use serde_derive::Serialize;
use tauri::State;
#[derive(Serialize)]
pub struct MidjourneyGetCredentialInfoResponse {
pub maybe_email: Option<String>,
pub can_clear_state: bool,
}
impl SerializeMarker for MidjourneyGetCredentialInfoResponse {}
#[tauri::command]
pub async fn midjourney_get_credential_info_command(
creds_manager: State<'_, MidjourneyCredentialManager>,
) -> ResponseOrErrorMessage<MidjourneyGetCredentialInfoResponse> {
info!("midjourney_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: &MidjourneyCredentialManager,
) -> AnyhowResult<MidjourneyGetCredentialInfoResponse> {
let mut can_clear_state = true;
let maybe_cookies = creds.maybe_copy_cookie_store()?;
let maybe_user_info = creds.maybe_copy_user_info()?;
if maybe_cookies.is_none() && maybe_user_info.is_none() {
can_clear_state = false;
}
let maybe_email = maybe_user_info
.map(|info| info.email.clone())
.flatten();
Ok(MidjourneyGetCredentialInfoResponse {
maybe_email,
can_clear_state,
})
}
@@ -1,38 +0,0 @@
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use crate::services::midjourney::windows::open_midjourney_login_window::open_midjourney_login_window;
use crate::services::sora::windows::sora_login_window::open_sora_login_window::open_sora_login_window;
use errors::AnyhowResult;
use log::{error, info};
use tauri::{AppHandle, State};
#[tauri::command]
pub async fn midjourney_open_login_command(
app: AppHandle,
app_data_root: State<'_, AppDataRoot>,
mj_creds_manager: State<'_, MidjourneyCredentialManager>,
) -> Result<String, String> {
info!("midjourney_open_login_command called");
do_open_login(&app, &app_data_root, &mj_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,
mj_creds_manager: &MidjourneyCredentialManager,
) -> AnyhowResult<()> {
info!("Building login window...");
open_midjourney_login_window(app, app_data_root, mj_creds_manager).await?;
info!("Done.");
Ok(())
}
@@ -1,3 +0,0 @@
pub mod midjourney_clear_credentials_command;
pub mod midjourney_get_credential_info_command;
pub mod midjourney_open_login_command;
@@ -1,5 +0,0 @@
pub mod commands;
pub mod state;
pub mod threads;
pub mod utils;
pub mod windows;
@@ -1,196 +0,0 @@
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
use crate::services::midjourney::state::serializable_midjourney_state::{SerializableMidjourneyState, SERIALIZABLE_MIDJOURNEY_STATE_VERSION};
use crate::services::storyteller::state::read_storyteller_credentials_from_disk::read_storyteller_credentials_from_disk;
use crate::services::storyteller::state::storyteller_credential_holder::StorytellerCredentialHolder;
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 midjourney_client::credentials::cookie_store_has_auth_cookies::cookie_store_has_auth_cookies;
use storyteller_client::credentials::storyteller_avt_cookie::StorytellerAvtCookie;
use storyteller_client::credentials::storyteller_credential_set::StorytellerCredentialSet;
use storyteller_client::credentials::storyteller_session_cookie::StorytellerSessionCookie;
#[derive(Clone)]
pub struct MidjourneyCredentialManager {
// TODO: Put last write/last read timestamps on these.
cookies: Arc<RwLock<Option<CookieStore>>>,
user_info: Arc<RwLock<Option<MidjourneyUserInfo>>>,
app_data_root: AppDataRoot,
}
impl MidjourneyCredentialManager {
pub fn initialize_empty(app_data_root: &AppDataRoot) -> Self {
Self {
cookies: Arc::new(RwLock::new(None)),
user_info: Arc::new(RwLock::new(None)),
app_data_root: app_data_root.clone(),
}
}
pub fn initialize_from_disk_infallible(app_data_root: &AppDataRoot) -> Self {
let mut cookies;
let mut user_info;
match read_midjourney_state_from_disk(app_data_root) {
Err(err) => {
warn!("Failed to read midjourney state from disk: {:?}", err);
cookies = Arc::new(RwLock::new(None));
user_info = Arc::new(RwLock::new(None));
},
Ok(None) => {
cookies = Arc::new(RwLock::new(None));
user_info = Arc::new(RwLock::new(None));
}
Ok(Some(state)) => {
let maybe_cookies = state.user_cookies
.map(|cookies| cookies.to_cookie_store());
let maybe_user_info = state.user_info
.map(|info| info.to_user_info());
cookies = Arc::new(RwLock::new(maybe_cookies));
user_info = Arc::new(RwLock::new(maybe_user_info));
}
};
Self {
cookies,
user_info,
app_data_root: app_data_root.clone(),
}
}
pub fn maybe_copy_cookie_store(&self) -> anyhow::Result<Option<CookieStore>> {
match self.cookies.read() {
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
Ok(store) => Ok(store.clone()),
}
}
pub fn maybe_copy_user_info(&self) -> anyhow::Result<Option<MidjourneyUserInfo>> {
match self.user_info.read() {
Err(err) => Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
Ok(info) => Ok(info.clone()),
}
}
pub fn replace_cookie_store(&self, store: CookieStore) -> anyhow::Result<()> {
match self.cookies.write() {
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
Ok(mut current_store) => {
*current_store = Some(store);
Ok(())
}
}
}
pub fn replace_user_info(&self, user_info: MidjourneyUserInfo) -> anyhow::Result<()> {
match self.user_info.write() {
Err(err) => Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
Ok(mut current_info) => {
*current_info = Some(user_info);
Ok(())
}
}
}
// NB: This is just a heuristic. We'll add better checks later.
pub fn session_appears_active(&self) -> anyhow::Result<bool> {
let maybe_cookies = match self.cookies.read() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
Ok(store) => store.clone(),
};
let maybe_user_info = match self.user_info.read() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
Ok(info) => info.clone(),
};
if maybe_cookies.is_none() || maybe_user_info.is_none() {
return Ok(false);
}
let cookies = maybe_cookies.unwrap();
let user_info = maybe_user_info.unwrap();
// NB: We consider the session active if we have auth cookies and a user id.
let has_user_id = user_info.user_id.is_some();
let maybe_has_auth_cookies = cookie_store_has_auth_cookies(&cookies);
// Misc cookies without login cookies are ~1055 length
// AUTH_I is ~1500 length
// AUTH_R is ~500 length
let maybe_has_big_enough_cookie = cookies.calculate_approx_cookie_character_length() > 2100;
// TODO: Heuristic should count.
// TODO: Consolidate with "login window thread" logic.
// TODO: Check timestamp of last valid request.
Ok(has_user_id && maybe_has_auth_cookies && maybe_has_big_enough_cookie && maybe_has_big_enough_cookie)
}
pub fn clear_credentials(&self) -> anyhow::Result<()> {
match self.cookies.write() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
Ok(mut store) => *store = None,
}
match self.user_info.write() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err)),
Ok(mut info) => *info = None,
}
Ok(())
}
pub fn persist_to_disk(&self) -> anyhow::Result<()> {
let maybe_cookies = match self.cookies.read() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
Ok(store) => store.clone(),
};
let maybe_user_info = match self.user_info.read() {
Err(err) => return Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err)),
Ok(info) => info.clone(),
};
let maybe_cookies= maybe_cookies
.map(|cookies| cookies.to_serializable());
let maybe_user_info= maybe_user_info
.map(|info| info.to_serializable());
if maybe_cookies.is_none() && maybe_user_info.is_none() {
info!("Nothing to write to disk, skipping.");
return Ok(());
}
let state = SerializableMidjourneyState {
version: SERIALIZABLE_MIDJOURNEY_STATE_VERSION,
user_cookies: maybe_cookies,
user_info: maybe_user_info,
};
let path = self.app_data_root.credentials_dir().get_midjourney_state_path();
let serialized = serde_json::to_string(&state)?;
std::fs::write(path, serialized)?;
Ok(())
}
}
fn read_midjourney_state_from_disk(root: &AppDataRoot) -> AnyhowResult<Option<SerializableMidjourneyState>> {
let midjourney_state_path= root.credentials_dir().get_midjourney_state_path();
if !midjourney_state_path.exists() {
return Ok(None);
}
let contents = read_to_string(&midjourney_state_path)?
.trim()
.to_string();
let state: SerializableMidjourneyState = serde_json::from_str(&contents)?;
Ok(Some(state))
}
@@ -1,30 +0,0 @@
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
use midjourney_client::recipes::get_user_info::GetUserInfoResponse;
use crate::services::midjourney::state::serializable_midjourney_state::SerializableMidjourneyUserInfo;
#[derive(Clone)]
pub struct MidjourneyUserInfo {
pub user_id: Option<MidjourneyUserId>,
pub email: Option<String>,
pub websocket_token: Option<String>,
}
impl MidjourneyUserInfo {
pub fn to_serializable(&self) -> SerializableMidjourneyUserInfo {
SerializableMidjourneyUserInfo {
user_id: self.user_id
.as_ref()
.map(|id| id.to_string()),
email: self.email.clone(),
websocket_token: self.websocket_token.clone(),
}
}
pub fn from_api_response(response: GetUserInfoResponse) -> Self {
Self {
user_id: response.user_id,
email: response.email,
websocket_token: response.websocket_token,
}
}
}
@@ -1,3 +0,0 @@
pub (super) mod serializable_midjourney_state;
pub mod midjourney_credential_manager;
pub mod midjourney_user_info;
@@ -1,41 +0,0 @@
use crate::services::midjourney::state::midjourney_user_info::MidjourneyUserInfo;
use cookie_store::serialized_cookie_store::SerializableCookieStore;
use midjourney_client::credentials::midjourney_user_id::MidjourneyUserId;
use serde_derive::{Deserialize, Serialize};
pub (super) const SERIALIZABLE_MIDJOURNEY_STATE_VERSION: u32 = 1;
#[derive(Serialize, Deserialize)]
pub (super) struct SerializableMidjourneyState {
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) user_info: Option<SerializableMidjourneyUserInfo>,
}
#[derive(Serialize, Deserialize)]
pub (super) struct SerializableMidjourneyUserInfo {
#[serde(skip_serializing_if = "Option::is_none")]
pub (super) user_id: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub (super) email: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub (super) websocket_token: Option<String>,
}
impl SerializableMidjourneyUserInfo {
pub fn to_user_info(&self) -> MidjourneyUserInfo {
MidjourneyUserInfo {
user_id: self.user_id
.as_deref()
.map(|id| MidjourneyUserId::from_str(id)),
email: self.email.clone(),
websocket_token: self.websocket_token.clone(),
}
}
}
@@ -1,261 +0,0 @@
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::services::midjourney::state::midjourney_credential_manager::MidjourneyCredentialManager;
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
use enums::common::generation_provider::GenerationProvider;
use enums::tauri::tasks::task_status::TaskStatus;
use errors::AnyhowResult;
use log::{error, info};
use midjourney_client::client::midjourney_hostname::MidjourneyHostname;
use midjourney_client::endpoints::imagine::{imagine, ImagineRequest};
use midjourney_client::utils::get_image_url::get_image_url;
use sqlite_tasks::queries::list_tasks_by_provider_and_tokens::{list_tasks_by_provider_and_tokens, ListTasksArgs, Task};
use std::collections::HashMap;
use std::fs::File;
use std::io::Write;
use std::path::PathBuf;
use tauri::AppHandle;
use url::Url;
use midjourney_client::utils::image_downloader_client::ImageDownloaderClient;
use sqlite_tasks::queries::update_task_status::{update_task_status, UpdateTaskArgs};
use storyteller_client::error::api_error::ApiError;
use storyteller_client::error::storyteller_error::StorytellerError;
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};
use crate::core::events::generation_events::generation_complete_event::GenerationCompleteEvent;
use crate::services::midjourney::utils::download_midjourney_image::download_midjourney_image;
/// This thread is responsible for picking up tasks that fell through the cracks of
/// the faster websocket thread.
pub async fn midjourney_long_polling_thread(
app_handle: AppHandle,
app_env_configs: AppEnvConfigs,
app_data_root: AppDataRoot,
task_database: TaskDatabase,
creds: MidjourneyCredentialManager,
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(std::time::Duration::from_millis(30_000)).await;
}
}
async fn polling_loop(
app_handle: &AppHandle,
app_env_configs: &AppEnvConfigs,
app_data_root: &AppDataRoot,
task_database: &TaskDatabase,
creds: &MidjourneyCredentialManager,
storyteller_creds_manager: &StorytellerCredentialManager,
) -> AnyhowResult<()> {
loop {
if !creds.session_appears_active()? {
tokio::time::sleep(std::time::Duration::from_millis(30_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 Midjourney polling.");
tokio::time::sleep(std::time::Duration::from_millis(10_000)).await;
continue;
}
};
let cookies = creds.maybe_copy_cookie_store()?;
let cookies = match cookies {
Some(cookies) => cookies,
None => {
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
continue;
}
};
let user_info = creds.maybe_copy_user_info()?;
let maybe_user_id = user_info
.map(|info| info.user_id)
.flatten();
let user_id = match maybe_user_id {
Some(user_id) => user_id,
None => {
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
continue;
}
};
let cookie_header = cookies.to_cookie_string();
let result = imagine(ImagineRequest {
hostname: MidjourneyHostname::Standard,
cookie_header,
user_id,
page_size: None,
}).await?;
let midjourney_items = result.items;
let midjourney_items_by_id = {
let mut hash = HashMap::new();
for item in midjourney_items.iter() {
if let Some(id) = &item.id {
hash.insert(id.to_string(), item.clone());
}
}
hash
};
//let midjourney_item_ids = midjourney_items.iter()
// .filter_map(|item| item.id.as_ref())
// .map(|id| id.to_string())
// .collect::<Vec<_>>();
let midjourney_item_ids = midjourney_items_by_id
.keys()
.map(|id| id.to_string())
.collect();
let local_tasks = list_tasks_by_provider_and_tokens(ListTasksArgs {
db: task_database.get_connection(),
provider: GenerationProvider::Midjourney,
provider_job_ids: Some(midjourney_item_ids),
}).await?;
let local_tasks = local_tasks.tasks;
// Map of Midjourney Job ID to Local Task.
let local_tasks_by_midjourney_job_id = local_tasks.iter()
.filter(|task| match task.status {
TaskStatus::Pending => true,
TaskStatus::Started => true,
TaskStatus::AttemptFailed => true,
_ => false,
})
.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>>();
// 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,
None => continue,
};
for index in 0..4 {
info!("Downloading generated Midjourney file...");
let download_path = download_midjourney_image(
&image_downloader,
midjourney_job_id,
index,
app_data_root
).await?;
let mut wait_delay = 0;
loop {
info!("Uploading to backend...");
let result = upload_image_media_file_from_file(UploadImageFromFileArgs {
api_host: &app_env_configs.storyteller_host,
maybe_creds: Some(&storyteller_creds),
path: &download_path,
is_intermediate_system_file: false,
}).await;
match result {
Ok(media_file) => {
info!("Successfully uploaded to backend: {:?}", media_file);
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(std::time::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
tokio::time::sleep(std::time::Duration::from_millis(100)).await;
}
let updated = update_task_status(UpdateTaskArgs {
db: task_database.get_connection(),
task_id: &local_task.id,
status: TaskStatus::CompleteSuccess,
}).await?;
if !updated {
continue; // If anything breaks with queries, don't spam events.
}
let event = GenerationCompleteEvent {
//media_file_token: result.media_file_token,
action: Some(GenerationAction::GenerateImage),
service: GenerationServiceProvider::Midjourney,
model: None,
};
if let Err(err) = event.send(&app_handle) {
error!("Failed to send GenerationCompleteEvent: {:?}", err); // Fail open
}
tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
}
//for task in tasks {
// let job_id = match task.provider_job_id {
// Some(job_id) => job_id,
// None => continue,
// };
// match task.status {
// TaskStatus::Pending => {}, // Fall-through
// TaskStatus::Started => {}, // Fall-through
// TaskStatus::AttemptFailed => {}, // Fall-through
// _ => {
// tasks_by_provider_job_id.remove(&job_id);
// continue
// },
// }
//}
tokio::time::sleep(std::time::Duration::from_millis(60_000)).await;
}
}

Some files were not shown because too many files have changed in this diff Show More