From 03336d8f7110482a00c610c79dd4dc0966a00889 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 16 May 2025 19:28:26 -0400 Subject: [PATCH] fal task polling and video download/upload --- Cargo.lock | 1 + crates/desktop/artcraft/Cargo.toml | 1 + .../src/threads/fal_task_polling_thread.rs | 116 +++++++++++++----- crates/lib/api_clients/fal_client/src/lib.rs | 5 +- .../src/utils/queue_status_checker.rs | 48 +++++--- .../storyteller_client/src/media_files/mod.rs | 1 + .../upload_video_media_file_from_file.rs | 91 ++++++++++++++ 7 files changed, 214 insertions(+), 49 deletions(-) create mode 100644 crates/lib/api_clients/storyteller_client/src/media_files/upload_video_media_file_from_file.rs diff --git a/Cargo.lock b/Cargo.lock index eb72aa232d..c71cf60273 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -850,6 +850,7 @@ dependencies = [ "tokens", "tokio", "tokio-util", + "url", "workspace-build-acceleration", ] diff --git a/crates/desktop/artcraft/Cargo.toml b/crates/desktop/artcraft/Cargo.toml index 25bb000aa1..15a3f9ff04 100644 --- a/crates/desktop/artcraft/Cargo.toml +++ b/crates/desktop/artcraft/Cargo.toml @@ -62,6 +62,7 @@ tempdir = "0.3.7" tempfile = "3.18.0" tokio = { version = "1.42", features = ["rt", "rt-multi-thread", "fs"] } tokio-util = { version = "0.7", features = ["codec"] } +url = { workspace = true } # NB: We need to vendor this in so that we can access the cookies inside via the private API tauri-plugin-http = { path = "../../vendor/tauri-plugin-http" } diff --git a/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs b/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs index 150e0fd290..23c2c8b4dd 100644 --- a/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs +++ b/crates/desktop/artcraft/src/threads/fal_task_polling_thread.rs @@ -8,9 +8,11 @@ use crate::state::fal::fal_task_queue::FalTaskQueue; use crate::state::sora::sora_credential_manager::SoraCredentialManager; use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager; use errors::AnyhowResult; +use fal_client::export::queue::Status; +use fal_client::fal_error_plus::FalErrorPlus; +use fal_client::utils::queue_status_checker::QueueStatusChecker; use log::{error, info, warn}; use openai_sora_client::requests::image_gen::image_gen_status::{Generation, TaskId, TaskStatus}; -use reqwest::Url; use serde_derive::Serialize; use std::fs::File; use std::io::Write; @@ -19,9 +21,12 @@ use storyteller_client::media_files::upload_image_media_file_from_file::upload_i use storyteller_client::utils::api_host::ApiHost; use tauri::{AppHandle, Emitter}; use tempdir::TempDir; -use fal_client::fal_error_plus::FalErrorPlus; -use fal_client::utils::queue_status_checker::QueueStatusChecker; +use tempfile::NamedTempFile; use tokens::tokens::media_files::MediaFileToken; +use url::Url; +use fal_client::model::fal_request_id::FalRequestId; +use storyteller_client::api_error::ApiError; +use storyteller_client::media_files::upload_video_media_file_from_file::{upload_video_media_file_from_file, UploadVideoMediaFileSuccessResponse}; pub async fn fal_task_polling_thread( app_handle: AppHandle, @@ -62,15 +67,15 @@ async fn polling_loop( info!("Task queue has {} pending tasks.", fal_task_queue.len()?); let key = fal_creds_manager.get_key_required()?; - + let storyteller_creds = storyteller_creds_manager.get_credentials_required()?; + let queue_status_checker = QueueStatusChecker::new(key); - // TODO: Clear completed tasks. - // let mut succeeded_tasks = Vec::new(); - // let mut failed_tasks = Vec::new(); - + let mut succeeded_tasks = Vec::new(); + //let mut failed_tasks = Vec::new(); // TODO + let len = fal_task_queue.len()?; - + for i in 0..len { let task = match fal_task_queue.get_index(i)? { Some(task) => task, @@ -80,24 +85,76 @@ async fn polling_loop( } }; - info!("Checking task status for task ID: {:?}", &task.enqueued_request.request_id); - + let task_id = task.enqueued_request.request_id.clone(); + + info!("Checking task status for task ID: {:?}", &task_id); + let result = queue_status_checker.check_status(&task.enqueued_request).await; - - match result { + + let url = match result { Err(err) => { warn!("Error checking job status: {:?}", err); tokio::time::sleep(std::time::Duration::from_millis(2_000)).await; + + // TODO: Some failures will be permanent job failures and the jobs will need to be removed. + continue; } Ok(status) => { info!("Status: {:?}", status); + match status.status { + Status::Completed => { + let result = queue_status_checker.get_download_url(&task.enqueued_request).await; + + let url = match result { + Ok(url) => url, + Err(err) => { + warn!("Failed to get download url: {:?}", err); + + // TODO: Some failures will be permanent job failures. + + continue; + } + }; + + url + } + _ => { + continue; + }, + } + } + }; + + + let download_file = download_generation(&url, &app_data_root).await?; + + let result = upload_video_media_file_from_file( + &ApiHost::Storyteller, + Some(&storyteller_creds), + download_file + ).await; + + match result { + Ok(success) => { + info!("Uploaded to API backend: {:?}", success.media_file_token); + + //let event = SoraImageGenerationCompleteEvent { + // media_file_token: success.media_file_token, + //}; + //event.send(&app_handle)?; + + succeeded_tasks.push(task_id); + } + Err(err) => { + warn!("Failed to upload video media file: {:?}", err); + continue; } } } - + // // TODO: The cursoring logic likely needs to improve. // for task in response.task_responses { // let status = TaskStatus::from_str(&task.status); @@ -157,31 +214,28 @@ async fn polling_loop( // } // } // -// let succeeded_task_ids : Vec<&TaskId> = succeeded_tasks -// .iter() -// .map(|task| &task.id) -// .collect(); -// -// fal_task_queue.remove_list(&succeeded_task_ids)?; + let succeeded_task_ids : Vec<&FalRequestId> = succeeded_tasks + .iter() + .map(|task| task) + .collect::>(); + + fal_task_queue.remove_list(&succeeded_task_ids)?; tokio::time::sleep(std::time::Duration::from_millis(6_000)).await; } } -async fn download_generation(generation: &Generation, app_data_root: &AppDataRoot) -> AnyhowResult { - let url = Url::parse(&generation.url)?; +async fn download_generation(url: &str, app_data_root: &AppDataRoot) -> AnyhowResult { + let url = Url::parse(&url)?; - let response = reqwest::get(&generation.url).await?; - let image_bytes = response.bytes().await?; + let response = reqwest::get(url.clone()).await?; + let response_bytes = response.bytes().await?; let ext = url.path().split(".").last().unwrap_or("png"); - let tempdir = app_data_root.temp_dir().path(); - let download_filename = format!("{}.{}", generation.id, ext); - let download_path = tempdir.join(download_filename); + let mut temp_file = app_data_root.temp_dir().new_named_temp_file_with_extension(ext)?; + + temp_file.write_all(&response_bytes)?; - let mut file = File::create(&download_path)?; - file.write_all(&image_bytes)?; - - Ok(download_path) + Ok(temp_file) } diff --git a/crates/lib/api_clients/fal_client/src/lib.rs b/crates/lib/api_clients/fal_client/src/lib.rs index 7eaa23d1c9..2524e89e74 100644 --- a/crates/lib/api_clients/fal_client/src/lib.rs +++ b/crates/lib/api_clients/fal_client/src/lib.rs @@ -1,8 +1,11 @@ //! A wrapper for the "fal" crate that bakes in a few extra recipes and utilities. -// Re-export extern crate fal; +// Re-export +pub mod export { + pub use fal::*; +} pub mod creds; pub mod fal_error_plus; diff --git a/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs index 256a16828a..8944cf16ad 100644 --- a/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs +++ b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs @@ -1,8 +1,11 @@ use fal::endpoints::fal_ai::kling_video::v1_6::pro::effects::I2VOutput; use fal::FalError; +use fal::prelude::fal_ai::imageutils::sam::RemoveBackgroundOutput; use fal::prelude::QueueResponse; use fal::queue::{Queue, QueueStatus}; use reqwest::Client; +use serde::de::DeserializeOwned; +use url::Url; use crate::creds::fal_api_key::FalApiKey; use crate::fal_error_plus::FalErrorPlus; use crate::model::enqueued_request::EnqueuedRequest; @@ -18,31 +21,42 @@ impl QueueStatusChecker { } pub async fn check_status(&self, request: &EnqueuedRequest) -> Result { + match &request.fal_endpoint { + FalEndpoint::Kling16ImageToVideo => { + let queue = self.make_queue::(&request); + let status = queue.status(false).await?; + Ok(status) + } + FalEndpoint::Other(endpoint) => { + Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint))) + } + } + } + + pub async fn get_download_url(&self, request: &EnqueuedRequest) -> Result { + match &request.fal_endpoint { + FalEndpoint::Kling16ImageToVideo => { + let queue = self.make_queue::(&request); + let status = queue.response().await?; + Ok(status.video.url) + } + FalEndpoint::Other(endpoint) => { + Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint))) + } + } + } + + pub fn make_queue(&self, request: &EnqueuedRequest) -> Queue { let client = Client::new(); let endpoint = request.fal_endpoint.url(); - + let payload = QueueResponse { request_id: request.request_id.0.clone(), response_url: request.response_url.clone(), status_url: request.status_url.clone(), cancel_url: request.cancel_url.clone(), }; - - let queue ; - - match &request.fal_endpoint { - FalEndpoint::Kling16ImageToVideo => { - // TODO: Cleanup / make concise - let f : Queue = Queue::new(client, endpoint, self.api_key.0.clone(), payload); - queue = f; - } - FalEndpoint::Other(endpoint) => { - return Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint))); - } - } - - let status = queue.status(false).await?; - Ok(status) + Queue::new(client, endpoint, self.api_key.0.clone(), payload) } } \ No newline at end of file diff --git a/crates/lib/api_clients/storyteller_client/src/media_files/mod.rs b/crates/lib/api_clients/storyteller_client/src/media_files/mod.rs index c9b77448b3..b10beebb3a 100644 --- a/crates/lib/api_clients/storyteller_client/src/media_files/mod.rs +++ b/crates/lib/api_clients/storyteller_client/src/media_files/mod.rs @@ -1,2 +1,3 @@ pub mod get_media_file; pub mod upload_image_media_file_from_file; +pub mod upload_video_media_file_from_file; diff --git a/crates/lib/api_clients/storyteller_client/src/media_files/upload_video_media_file_from_file.rs b/crates/lib/api_clients/storyteller_client/src/media_files/upload_video_media_file_from_file.rs new file mode 100644 index 0000000000..05cb9ae2ae --- /dev/null +++ b/crates/lib/api_clients/storyteller_client/src/media_files/upload_video_media_file_from_file.rs @@ -0,0 +1,91 @@ +use crate::api_error::ApiError; +use crate::credentials::storyteller_credential_set::StorytellerCredentialSet; +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; +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; + +// TODO(bt,2025-04-22): Share API definitions between client and server in common crate. + +#[derive(Deserialize, Debug)] +pub struct UploadVideoMediaFileSuccessResponse { + pub success: bool, + pub media_file_token: MediaFileToken, +} + +/// Upload a video media file from a file. +pub async fn upload_video_media_file_from_file>( + api_host: &ApiHost, + maybe_creds: Option<&StorytellerCredentialSet>, + path: P, +) -> Result { + + let url = get_route(api_host); + + debug!("Requesting {:?}", &url); + + //let mut file = File::open(path)?; + + let client = Client::builder() + .gzip(true) + .build()?; + + let form = Form::new() + .text("uuid_idempotency_token", generate_random_uuid()) + .file("file", path) + .await?; + + let mut request_builder = client.post(url) + .header("User-Agent", USER_AGENT) + .header("Accept", APPLICATION_JSON); + + if let Some(creds) = 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?; + + let response = filter_bad_response(response).await?; + let response_body = &response.text().await?; + + let media_file = serde_json::from_str(&response_body)?; + + Ok(media_file) +} + +fn get_route(api_host: &ApiHost) -> String { + let api_hostname = api_host.to_api_hostname(); + format!("https://{}/v1/media_files/upload/new_video", api_hostname) +} +