From a4e7a785d70699092ea91c923554180380d3b58f Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Mon, 21 Apr 2025 16:13:37 -0400 Subject: [PATCH] new sora tauri endpoints --- Cargo.lock | 1 + Cargo.toml | 4 +- crates/desktop/tauri-artcraft/Cargo.toml | 1 + .../tauri-artcraft/src/commands/sora/mod.rs | 3 +- .../sora/sora_image_generation_command.rs | 33 +++--- .../commands/sora/sora_image_remix_command.rs | 100 ++++++++++++++++++ crates/desktop/tauri-artcraft/src/lib.rs | 4 +- crates/service/web/storyteller_web/Cargo.toml | 2 +- 8 files changed, 129 insertions(+), 19 deletions(-) create mode 100644 crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs diff --git a/Cargo.lock b/Cargo.lock index 504cb464f4..8b93c0a164 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -12021,6 +12021,7 @@ dependencies = [ "tauri-plugin-log", "tempdir", "tempfile", + "tokens", "tokio", "tokio-util", "workspace-build-acceleration", diff --git a/Cargo.toml b/Cargo.toml index 1910fc149f..1f8864af57 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -168,12 +168,14 @@ sha3 = "0.10.8" thiserror = "2" toml = "0.7.4" url = "2.5.4" + # Internal libs crockford = { path = "crates/lib/data/crockford" } errors = { path = "crates/lib/errors" } filesys = { path = "crates/lib/files/filesys" } memory_caching = { path = "crates/lib/caching/memory_caching" } -shared_service_components = { path = "crates/lib/shared_service_components" } openai_sora_client = { path = "crates/lib/api_clients/openai_sora_client" } primitives = { path = "crates/lib/data/primitives" } +shared_service_components = { path = "crates/lib/shared_service_components" } subprocess_common = { path = "crates/lib/subprocess_common" } +tokens = { path = "crates/_types/tokens" } diff --git a/crates/desktop/tauri-artcraft/Cargo.toml b/crates/desktop/tauri-artcraft/Cargo.toml index 013a9de07c..920e0f5920 100644 --- a/crates/desktop/tauri-artcraft/Cargo.toml +++ b/crates/desktop/tauri-artcraft/Cargo.toml @@ -30,6 +30,7 @@ memory_caching = { workspace = true } openai_sora_client = { workspace = true } primitives = { workspace = true } subprocess_common = { workspace = true } +tokens = { workspace = true } # External #openssl = { version = "0.10", features = ["vendored"] } diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/mod.rs b/crates/desktop/tauri-artcraft/src/commands/sora/mod.rs index 9054cfe591..0c197b2d85 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/mod.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/mod.rs @@ -1,2 +1,3 @@ -pub mod sora_image_generation_command; pub mod open_sora_login_command; +pub mod sora_image_generation_command; +pub mod sora_image_remix_command; diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs index dfee9749bb..f578edbb98 100644 --- a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_generation_command.rs @@ -16,17 +16,14 @@ use crate::state::app_dir::AppDataRoot; #[tauri::command] pub async fn sora_image_generation_command( - image: &str, prompt: &str, - app: AppHandle, + image: Option<&str>, + _app: AppHandle, app_data_root: State<'_, AppDataRoot>, ) -> Result { info!("image_generation_command called; processing image..."); - let bytes = BASE64_STANDARD.decode(image) - .map_err(|err| format!("Base64 decode error: {}", err))?; - - generate_image(bytes, prompt, &app_data_root) + generate_image(image, prompt, &app_data_root) .await .map_err(|err| { error!("error: {:?}", err); @@ -36,19 +33,25 @@ pub async fn sora_image_generation_command( Ok("success".to_string()) } -pub async fn generate_image(file_bytes: Vec, prompt: &str, app_data_root: &AppDataRoot) -> AnyhowResult<()> { +pub async fn generate_image(maybe_image: Option<&str>, prompt: &str, app_data_root: &AppDataRoot) -> AnyhowResult<()> { let sora_credentials = get_credentials(app_data_root)?; - let filename = "image.png".to_string(); + let mut sora_media_tokens = vec![]; - let response= sora_media_upload_from_bytes(file_bytes, filename, &sora_credentials) - .await - .map_err(|err| { - error!("Failed to upload image to Sora: {:?}", err); - err - })?; + if let Some(image) = maybe_image { + let image_bytes = BASE64_STANDARD.decode(image)?; - let sora_media_tokens = vec![response.id]; + let filename = "image.png".to_string(); + + let response = sora_media_upload_from_bytes(image_bytes, filename, &sora_credentials) + .await + .map_err(|err| { + error!("Failed to upload image to Sora: {:?}", err); + err + })?; + + sora_media_tokens.push(response.id); + } let response = sora_image_gen_remix(SoraImageGenRemixRequest { prompt: prompt.to_string(), diff --git a/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs new file mode 100644 index 0000000000..00dc382776 --- /dev/null +++ b/crates/desktop/tauri-artcraft/src/commands/sora/sora_image_remix_command.rs @@ -0,0 +1,100 @@ +use base64::Engine; +use base64::prelude::BASE64_STANDARD; +use crate::state::app_dir::AppDataRoot; +use errors::AnyhowResult; +use image::codecs::png::{CompressionType, FilterType, PngEncoder}; +use image::{DynamicImage, ImageReader}; +use log::{error, info}; +use openai_sora_client::credentials::SoraCredentials; +use openai_sora_client::image_gen::common::{ImageSize, NumImages}; +use openai_sora_client::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest}; +use openai_sora_client::upload::upload_media_from_bytes::sora_media_upload_from_bytes; +use openai_sora_client::upload::upload_media_from_file::SoraMediaUploadRequest; +use serde_derive::Deserialize; +use std::fs::read_to_string; +use std::io::Cursor; +use tauri::{AppHandle, Manager, State}; +use tokens::tokens::media_files::MediaFileToken; + +#[derive(Deserialize)] +pub struct SoraImageRemixCommand { + /// Image media file; the engine or canvas snapshot (screenshot). + pub snapshot_media_token: MediaFileToken, + + /// The user's image generation prompt. + pub prompt: String, + + /// Turn off the system prompt. + pub disable_system_prompt: Option, + + /// Additional images to include (optional). Up to nine images. + pub maybe_additional_images: Option>, + + pub maybe_number_of_samples: Option, +} + +#[tauri::command] +pub async fn sora_image_remix_command( + _app: AppHandle, + request: SoraImageRemixCommand, + app_data_root: State<'_, AppDataRoot>, +) -> Result { + info!("image_generation_command called; processing image..."); + + generate_image(request, &app_data_root) + .await + .map_err(|err| { + error!("error: {:?}", err); + "there was an error".to_string() + })?; + + Ok("success".to_string()) +} + +pub async fn generate_image(request: SoraImageRemixCommand, app_data_root: &AppDataRoot) -> AnyhowResult<()> { + let sora_credentials = get_credentials(app_data_root)?; + + let sora_media_tokens = vec![]; + + // TODO(bt,2025-04-21): Download media tokens. + // Note: This is incredibly inefficient. We should keep a local cache. + // Also, if they've already been uploaded to OpenAI, we shouldn't continue to re-upload. + + let response = sora_image_gen_remix(SoraImageGenRemixRequest { + prompt: request.prompt.to_string(), + num_images: NumImages::One, + image_size: ImageSize::Square, + sora_media_tokens: sora_media_tokens.clone(), + credentials: &sora_credentials, + }).await + .map_err(|err| { + error!("Failed to call Sora image generation: {:?}", err); + err + })?; + + println!(">> TASK ID: {:?} ", response.task_id); + + Ok(()) +} + +fn get_credentials(app_data_root: &AppDataRoot) -> AnyhowResult { + let cookie_file = app_data_root.get_sora_cookie_file_path(); + + let bearer = read_to_string("/Users/bt/dev/storyteller/storyteller-rust/test_data/temp/bearer.txt")? + .trim() + .to_string(); + + let cookie= read_to_string(cookie_file)? + .trim() + .to_string(); + + let sentinel = read_to_string("/Users/bt/dev/storyteller/storyteller-rust/test_data/temp/sentinel.txt")? + .trim() + .to_string(); + + Ok(SoraCredentials { + bearer_token: bearer, + cookie: cookie, + sentinel: Some(sentinel), + }) +} \ No newline at end of file diff --git a/crates/desktop/tauri-artcraft/src/lib.rs b/crates/desktop/tauri-artcraft/src/lib.rs index e9cfceedc6..fdc96f79f6 100644 --- a/crates/desktop/tauri-artcraft/src/lib.rs +++ b/crates/desktop/tauri-artcraft/src/lib.rs @@ -8,6 +8,7 @@ pub mod utils; use crate::commands::flip_image::flip_image; use crate::commands::sora::open_sora_login_command::open_sora_login_command; use crate::commands::sora::sora_image_generation_command::sora_image_generation_command; +use crate::commands::sora::sora_image_remix_command::sora_image_remix_command; use crate::state::app_config::AppConfig; use crate::threads::sora_session_login_thread::sora_session_login_thread; @@ -57,8 +58,9 @@ pub fn run() { .manage(app_data_root) .invoke_handler(tauri::generate_handler![ flip_image, - sora_image_generation_command, open_sora_login_command, + sora_image_generation_command, + sora_image_remix_command, ]) .run(tauri::generate_context!("tauri.conf.json")) .expect("error while running tauri application"); diff --git a/crates/service/web/storyteller_web/Cargo.toml b/crates/service/web/storyteller_web/Cargo.toml index 821ae34557..d013aec500 100644 --- a/crates/service/web/storyteller_web/Cargo.toml +++ b/crates/service/web/storyteller_web/Cargo.toml @@ -68,7 +68,7 @@ reusable_types = { path = "../../../_types/reusable_types" } server_environment = { path = "../../../lib/server_environment" } shared_service_components = { workspace = true } thumbnail_generator = { path = "../../../lib/api_clients/thumbnail_generator" } -tokens = { path = "../../../_types/tokens" } +tokens = { workspace = true } tts_common = { path = "../../../lib/tts_common" } url_config = { path = "../../../lib/url_config" } user_input_common = { path = "../../../lib/user_input_common" }