fal task polling and video download/upload

Cargo.lock

@@ -850,6 +850,7 @@ dependencies = [
  "tokens",
  "tokio",
  "tokio-util",
+ "url",
  "workspace-build-acceleration",
 ]
 

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" }

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::<Vec<_>>();
+
+    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<PathBuf> {
-  let url = Url::parse(&generation.url)?;
+async fn download_generation(url: &str, app_data_root: &AppDataRoot) -> AnyhowResult<NamedTempFile> {
+  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)
 }

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;

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<QueueStatus, FalErrorPlus> {
+    match &request.fal_endpoint {
+      FalEndpoint::Kling16ImageToVideo => {
+        let queue = self.make_queue::<I2VOutput>(&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<String, FalErrorPlus> {
+    match &request.fal_endpoint {
+      FalEndpoint::Kling16ImageToVideo => {
+        let queue = self.make_queue::<I2VOutput>(&request);
+        let status = queue.response().await?;
+        Ok(status.video.url)
+      }
+      FalEndpoint::Other(endpoint) => {
+        Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint)))
+      }
+    }
+  }
+  
+  pub fn make_queue<T: DeserializeOwned>(&self, request: &EnqueuedRequest) -> Queue<T> {
     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<I2VOutput> = 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)
   }
 }

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;

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<P: AsRef<Path>>(
+  api_host: &ApiHost,
+  maybe_creds: Option<&StorytellerCredentialSet>,
+  path: P,
+) -> Result<UploadVideoMediaFileSuccessResponse, ApiError> {
+
+  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)
+}
+