mirror of
https://github.com/storytold/artcraft.git
synced 2026-11-13 05:12:50 +00:00
fal video wip, polling more complete
This commit is contained in:
1 parent
6012421d1a
commit
a792ed9a09
15 files changed
+692
-16
No files matched your search
@@ -143,7 +143,7 @@ pub async fn fal_background_removal_command(
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to emit event: {:?}", err);
|
||||
//}
|
||||
|
||||
|
||||
Ok(FalBackgroundRemovalSuccessResponse {
|
||||
media_token: result.0.media_file.token,
|
||||
cdn_url: result.0.media_file.media_links.cdn_url,
|
||||
@@ -159,7 +159,6 @@ enum InnerError {
|
||||
AnyhowError(AnyhowError),
|
||||
StorytellerApiError(ApiError),
|
||||
DecodeError(DecodeError),
|
||||
//ImageError(image::ImageError),
|
||||
IoError(std::io::Error),
|
||||
}
|
||||
|
||||
@@ -187,12 +186,6 @@ impl From<DecodeError> for InnerError {
|
||||
}
|
||||
}
|
||||
|
||||
//impl From<image::ImageError> for InnerError {
|
||||
// fn from(value: image::ImageError) -> Self {
|
||||
// Self::ImageError(value)
|
||||
// }
|
||||
//}
|
||||
//
|
||||
impl From<std::io::Error> for InnerError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::IoError(value)
|
||||
|
||||
@@ -0,0 +1,246 @@
|
||||
use crate::commands::command_response_wrapper::{CommandErrorResponseWrapper, CommandErrorStatus, CommandResult, SerializeMarker};
|
||||
use crate::events::sendable_event_trait::SendableEvent;
|
||||
use crate::events::sora::sora_image_enqueue_failure_event::SoraImageEnqueueFailureEvent;
|
||||
use crate::events::sora::sora_image_enqueue_success_event::SoraImageEnqueueSuccessEvent;
|
||||
use crate::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use crate::state::fal::fal_credential_manager::FalCredentialManager;
|
||||
use crate::state::sora::read_sora_credentials_from_disk::read_sora_credentials_from_disk;
|
||||
use crate::state::sora::sora_credential_holder::SoraCredentialHolder;
|
||||
use crate::state::sora::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::state::sora::sora_task_queue::SoraTaskQueue;
|
||||
use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use crate::utils::get_url_file_extension::get_url_file_extension;
|
||||
use crate::utils::simple_http_download::simple_http_download;
|
||||
use crate::utils::simple_http_download_to_tempfile::simple_http_download_to_tempfile;
|
||||
use anyhow::anyhow;
|
||||
use base64::prelude::BASE64_STANDARD;
|
||||
use base64::{DecodeError, Engine};
|
||||
use errors::{AnyhowError, AnyhowResult};
|
||||
use fal_client::fal_error_plus::FalErrorPlus;
|
||||
use fal_client::requests::remove_background_rembg::remove_background_rembg;
|
||||
use image::codecs::png::{CompressionType, FilterType, PngEncoder};
|
||||
use image::{DynamicImage, EncodableLayout, ImageReader};
|
||||
use log::{debug, error, info, warn};
|
||||
use mimetypes::mimetype_info::mimetype_info::MimetypeInfo;
|
||||
use openai_sora_client::credentials::SoraCredentials;
|
||||
use openai_sora_client::creds::credential_migration::CredentialMigrationRef;
|
||||
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::requests::image_gen::sora_image_gen_remix::{sora_image_gen_remix, SoraImageGenRemixRequest};
|
||||
use openai_sora_client::requests::upload::upload_media_from_bytes::sora_media_upload_from_bytes;
|
||||
use openai_sora_client::requests::upload::upload_media_from_file::sora_media_upload_from_file;
|
||||
use openai_sora_client::sora_error::SoraError;
|
||||
use reqwest::Url;
|
||||
use serde_derive::{Deserialize, Serialize};
|
||||
use std::fs::read_to_string;
|
||||
use std::io::{Cursor, Read, Write};
|
||||
use std::ops::Add;
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
use storyteller_client::api_error::ApiError;
|
||||
use storyteller_client::api_error::ApiError::InternalServerError;
|
||||
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;
|
||||
use storyteller_client::utils::api_host::ApiHost;
|
||||
use tauri::{AppHandle, Emitter, Manager, State};
|
||||
use tempfile::NamedTempFile;
|
||||
use fal_client::requests::enqueue_kling_16_image_to_video::{enqueue_kling_16_image_to_video, Kling16Args, Kling16Duration};
|
||||
use filesys::file_read_bytes::file_read_bytes;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
#[derive(Deserialize)]
|
||||
pub struct FalKlingImageToVideoRequest {
|
||||
/// Image media file; the image to remove the background from.
|
||||
pub image_media_token: Option<MediaFileToken>,
|
||||
|
||||
/// Base64-encoded image
|
||||
pub base64_image: Option<String>,
|
||||
}
|
||||
|
||||
#[derive(Serialize, Debug)]
|
||||
#[serde(rename_all = "snake_case")]
|
||||
pub enum SoraImageRemixErrorType {
|
||||
/// Generic server error
|
||||
ServerError,
|
||||
/// No Fal API key available
|
||||
NeedsFalApiKey,
|
||||
}
|
||||
|
||||
|
||||
#[tauri::command]
|
||||
pub async fn fal_kling_image_to_video_command(
|
||||
request: FalKlingImageToVideoRequest,
|
||||
app_data_root: State<'_, AppDataRoot>,
|
||||
fal_creds_manager: State<'_, FalCredentialManager>,
|
||||
storyteller_creds_manager: State<'_, StorytellerCredentialManager>,
|
||||
) -> CommandResult<(), SoraImageRemixErrorType, ()> {
|
||||
|
||||
info!("fal_kling_image_to_video_command called; image media token: {:?}", request.image_media_token);
|
||||
|
||||
let has_credentials = fal_creds_manager
|
||||
.has_apparent_api_token()
|
||||
.unwrap_or(true); // NB: Failures would be lock issues
|
||||
|
||||
if !has_credentials {
|
||||
warn!("No API key found");
|
||||
return Err(CommandErrorResponseWrapper {
|
||||
status: CommandErrorStatus::Unauthorized,
|
||||
error_message: Some("You need to set a FAL api key".to_string()),
|
||||
error_type: Some(SoraImageRemixErrorType::NeedsFalApiKey),
|
||||
error_details: None,
|
||||
});
|
||||
}
|
||||
|
||||
let result = image_to_video(
|
||||
request,
|
||||
&app_data_root,
|
||||
&fal_creds_manager,
|
||||
&storyteller_creds_manager
|
||||
).await;
|
||||
|
||||
if let Err(err) = result {
|
||||
error!("error: {:?}", err);
|
||||
|
||||
//let event = SoraImageEnqueueFailureEvent {};
|
||||
//let result = event.send(&app);
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to emit event: {:?}", err);
|
||||
//}
|
||||
|
||||
let mut status = CommandErrorStatus::ServerError;
|
||||
let mut error_type = SoraImageRemixErrorType::ServerError;
|
||||
let mut error_message = "A server error occurred. Please try again. If it continues, please tell our staff about the problem.";
|
||||
|
||||
//match (err) {
|
||||
// _ => {},
|
||||
//}
|
||||
|
||||
return Err(CommandErrorResponseWrapper {
|
||||
status,
|
||||
error_message: Some(error_message.to_string()),
|
||||
error_type: Some(error_type),
|
||||
error_details: None,
|
||||
})
|
||||
}
|
||||
|
||||
//let event = SoraImageEnqueueSuccessEvent {};
|
||||
//let result = event.send(&app);
|
||||
//if let Err(err) = result {
|
||||
// error!("Failed to emit event: {:?}", err);
|
||||
//}
|
||||
|
||||
Ok(().into())
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
enum InnerError {
|
||||
FalError(FalErrorPlus),
|
||||
AnyhowError(AnyhowError),
|
||||
StorytellerApiError(ApiError),
|
||||
DecodeError(DecodeError),
|
||||
IoError(std::io::Error),
|
||||
}
|
||||
|
||||
impl From<AnyhowError> for InnerError {
|
||||
fn from(value: AnyhowError) -> Self {
|
||||
Self::AnyhowError(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<FalErrorPlus> for InnerError {
|
||||
fn from(value: FalErrorPlus) -> Self {
|
||||
Self::FalError(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<ApiError> for InnerError {
|
||||
fn from(value: ApiError) -> Self {
|
||||
Self::StorytellerApiError(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<DecodeError> for InnerError {
|
||||
fn from(value: DecodeError) -> Self {
|
||||
Self::DecodeError(value)
|
||||
}
|
||||
}
|
||||
|
||||
impl From<std::io::Error> for InnerError {
|
||||
fn from(value: std::io::Error) -> Self {
|
||||
Self::IoError(value)
|
||||
}
|
||||
}
|
||||
|
||||
pub async fn image_to_video(
|
||||
request: FalKlingImageToVideoRequest,
|
||||
app_data_root: &AppDataRoot,
|
||||
fal_creds_manager: &FalCredentialManager,
|
||||
storyteller_creds_manager: &StorytellerCredentialManager,
|
||||
) -> Result<(), InnerError> {
|
||||
|
||||
let api_key = fal_creds_manager.get_key_required()?;
|
||||
let creds = storyteller_creds_manager.get_credentials_required()?;
|
||||
|
||||
let mut temp_download;
|
||||
|
||||
if let Some(media_token) = request.image_media_token {
|
||||
temp_download = download_media_file(&app_data_root, &media_token).await?;
|
||||
|
||||
} else if let Some(base64_bytes) = request.base64_image {
|
||||
temp_download = persist_base64(&app_data_root, base64_bytes).await?;
|
||||
} else {
|
||||
return Err(InnerError::AnyhowError(anyhow!("No image media token or base64 image provided")));
|
||||
}
|
||||
|
||||
info!("Calling FAL image to video ...");
|
||||
|
||||
let filename = temp_download.path().to_path_buf();
|
||||
|
||||
let enqueued = enqueue_kling_16_image_to_video(Kling16Args {
|
||||
image_path: filename,
|
||||
api_key: &api_key,
|
||||
duration: Kling16Duration::Default,
|
||||
}).await?;
|
||||
|
||||
// TODO: Handle event polling
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn download_media_file(app_data_root: &AppDataRoot, token: &MediaFileToken) -> Result<NamedTempFile, InnerError> {
|
||||
let response = get_media_file(&ApiHost::Storyteller, token).await?;
|
||||
|
||||
let media_file_url = &response.media_file.media_links.cdn_url;
|
||||
|
||||
let extension_with_dot = get_url_file_extension(media_file_url)
|
||||
.map(|ext| format!(".{}", ext))
|
||||
.unwrap_or_else(|| ".png".to_string());
|
||||
|
||||
//let filename = format!("{}{}", response.media_file.token.as_str(), extension_with_dot);
|
||||
//let filename = app_data_root.downloads_dir().path().join(&filename);
|
||||
|
||||
let mut file = app_data_root.temp_dir().new_named_temp_file_with_extension(&extension_with_dot)?;
|
||||
|
||||
simple_http_download_to_tempfile(&media_file_url, &mut file).await?;
|
||||
|
||||
Ok(file) // NB: Must return TempFile to not drop / delete it
|
||||
}
|
||||
|
||||
async fn persist_base64(app_data_root: &AppDataRoot, base64_image: String) -> Result<NamedTempFile, InnerError> {
|
||||
let bytes = BASE64_STANDARD.decode(base64_image)?;
|
||||
|
||||
let extension = MimetypeInfo::get_for_bytes(&bytes)
|
||||
.map(|info| info.file_extension())
|
||||
.flatten()
|
||||
.map(|ext| ext.extension_with_period().to_string())
|
||||
.unwrap_or_else(|| ".png".to_string());
|
||||
|
||||
let mut file = app_data_root.temp_dir().new_named_temp_file_with_extension(&extension)?;
|
||||
|
||||
file.write_all(bytes.as_ref())?;
|
||||
|
||||
Ok(file) // NB: Must return TempFile to not drop / delete it
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod fal_background_removal_command;
|
||||
pub mod fal_kling_image_to_video_command;
|
||||
pub mod get_fal_api_key_command;
|
||||
pub mod set_fal_api_key_command;
|
||||
@@ -10,6 +10,7 @@ use tauri::Manager;
|
||||
use crate::commands::app_preferences::get_app_preferences_command::get_app_preferences_command;
|
||||
use crate::commands::app_preferences::update_app_preference_command::update_app_preferences_command;
|
||||
use crate::commands::fal::fal_background_removal_command::fal_background_removal_command;
|
||||
use crate::commands::fal::fal_kling_image_to_video_command::fal_kling_image_to_video_command;
|
||||
use crate::commands::fal::get_fal_api_key_command::get_fal_api_key_command;
|
||||
use crate::commands::fal::set_fal_api_key_command::set_fal_api_key_command;
|
||||
use crate::commands::flip_image::flip_image;
|
||||
@@ -22,12 +23,14 @@ use crate::commands::sora::sora_logout_command::sora_logout_command;
|
||||
use crate::state::app_preferences::app_preferences_manager::load_app_preferences_or_default;
|
||||
use crate::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::state::fal::fal_credential_manager::FalCredentialManager;
|
||||
use crate::state::fal::fal_task_queue::FalTaskQueue;
|
||||
use crate::state::main_window_position::MainWindowPosition;
|
||||
use crate::state::main_window_size::MainWindowSize;
|
||||
use crate::state::sora::sora_credential_manager::SoraCredentialManager;
|
||||
use crate::state::sora::sora_task_queue::SoraTaskQueue;
|
||||
use crate::state::storyteller::storyteller_credential_manager::StorytellerCredentialManager;
|
||||
use crate::threads::discord_presence_thread::discord_presence_thread;
|
||||
use crate::threads::fal_task_polling_thread::fal_task_polling_thread;
|
||||
use crate::threads::main_window_thread::main_window_thread::main_window_thread;
|
||||
use crate::threads::sora_task_polling_thread::sora_task_polling_thread;
|
||||
use crate::utils::webview_unsafe::webview_unsafe_for_app;
|
||||
@@ -59,11 +62,15 @@ pub fn run() {
|
||||
|
||||
println!("Attempting to read existing fal credentials...");
|
||||
let fal_creds_manager = FalCredentialManager::initialize_from_disk_infallible(&app_data_root);
|
||||
let fal_creds_manager_2 = fal_creds_manager.clone();
|
||||
|
||||
// Other state
|
||||
let sora_task_queue = SoraTaskQueue::new();
|
||||
let sora_task_queue_2 = sora_task_queue.clone();
|
||||
|
||||
let fal_task_queue = FalTaskQueue::new();
|
||||
let fal_task_queue_2 = fal_task_queue.clone();
|
||||
|
||||
println!("Initializing backend runtime...");
|
||||
|
||||
tauri::Builder::default()
|
||||
@@ -97,13 +104,19 @@ pub fn run() {
|
||||
eprintln!("Error setting webview unsafe: {:?}", err);
|
||||
}
|
||||
|
||||
// TODO(bt): Clean this up. We can just clone at the callsite. Also clean initialization
|
||||
let app_2 = app.clone();
|
||||
let app_3 = app.clone();
|
||||
let app_4 = app.clone();
|
||||
let app_5 = app.clone();
|
||||
let app_data_root_3 = app_data_root_2.clone();
|
||||
let app_data_root_4 = app_data_root_2.clone();
|
||||
let app_data_root_5 = app_data_root_2.clone();
|
||||
let sora_creds_manager_3 = sora_creds_manager_2.clone();
|
||||
let storyteller_creds_manager_3 = storyteller_creds_manager.clone();
|
||||
let storyteller_creds_manager_4 = storyteller_creds_manager.clone();
|
||||
let fal_creds_manager_3 = fal_creds_manager_2.clone();
|
||||
let fal_task_queue_3 = fal_task_queue_2.clone();
|
||||
|
||||
match MainWindowSize::from_filesystem_configs(&app_data_root_3) {
|
||||
Ok(None) => {}
|
||||
@@ -136,18 +149,21 @@ pub fn run() {
|
||||
tauri::async_runtime::spawn(main_window_thread(app_3, app_data_root_3, storyteller_creds_manager_2));
|
||||
tauri::async_runtime::spawn(sora_task_polling_thread(app_4, app_data_root_4, sora_creds_manager_3, storyteller_creds_manager_3, sora_task_queue_2));
|
||||
tauri::async_runtime::spawn(discord_presence_thread());
|
||||
tauri::async_runtime::spawn(fal_task_polling_thread(app_5, app_data_root_5, fal_creds_manager_3, storyteller_creds_manager_4, fal_task_queue_3));
|
||||
|
||||
Ok(())
|
||||
})
|
||||
.manage(app_data_root)
|
||||
.manage(app_preferences)
|
||||
.manage(sora_creds_manager)
|
||||
.manage(fal_creds_manager)
|
||||
.manage(fal_task_queue)
|
||||
.manage(sora_creds_manager)
|
||||
.manage(sora_task_queue)
|
||||
.manage(storyteller_creds_manager_3)
|
||||
.invoke_handler(tauri::generate_handler![
|
||||
check_sora_session_command,
|
||||
fal_background_removal_command,
|
||||
fal_kling_image_to_video_command,
|
||||
flip_image,
|
||||
get_app_preferences_command,
|
||||
get_fal_api_key_command,
|
||||
|
||||
@@ -0,0 +1,164 @@
|
||||
use chrono::{DateTime, Utc};
|
||||
use errors::AnyhowResult;
|
||||
use fal_client::model::enqueued_request::EnqueuedRequest;
|
||||
use fal_client::model::fal_request_id::FalRequestId;
|
||||
use indexmap::IndexMap;
|
||||
use std::collections::HashSet;
|
||||
use std::sync::{Arc, RwLock};
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FalTaskDetails {
|
||||
pub enqueue_time: DateTime<Utc>,
|
||||
pub enqueued_request: EnqueuedRequest,
|
||||
}
|
||||
|
||||
impl FalTaskDetails {
|
||||
pub fn for_request(request: EnqueuedRequest) -> Self {
|
||||
Self {
|
||||
enqueue_time: Utc::now(),
|
||||
enqueued_request: request,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
pub struct FalTaskQueue {
|
||||
// Insertion-order map of task IDs to task details
|
||||
queue: Arc<RwLock<IndexMap<FalRequestId, FalTaskDetails>>>
|
||||
}
|
||||
|
||||
impl FalTaskQueue {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
queue: Arc::new(RwLock::new(IndexMap::new()))
|
||||
}
|
||||
}
|
||||
|
||||
/// Determine if queue is empty
|
||||
/// This is an O(1) operation.
|
||||
pub fn is_empty(&self) -> AnyhowResult<bool> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => Ok(lock.is_empty())
|
||||
}
|
||||
}
|
||||
|
||||
/// Get length of set
|
||||
/// This is an O(1) operation.
|
||||
pub fn len(&self) -> AnyhowResult<usize> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => Ok(lock.len())
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the task at an index.
|
||||
/// This is an O(1) operation.
|
||||
pub fn get_index(&self, i: usize) -> AnyhowResult<Option<FalTaskDetails>> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => {
|
||||
if i >= lock.len() {
|
||||
Ok(None)
|
||||
} else {
|
||||
match lock.get_index(i) {
|
||||
None => Ok(None),
|
||||
Some((_key, value)) => {
|
||||
Ok(Some(value.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Determine if the key is in the queue
|
||||
/// This is an O(1) operation.
|
||||
pub fn contains_key(&self, request_id: &FalRequestId) -> AnyhowResult<bool> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => Ok(lock.contains_key(request_id))
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the first enqueued task
|
||||
/// This is an O(1) operation.
|
||||
pub fn first(&self) -> AnyhowResult<Option<FalTaskDetails>> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => match lock.first() {
|
||||
None => Ok(None),
|
||||
Some((request_id, task_details)) => Ok(Some(task_details.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Get the last enqueued task
|
||||
/// This is an O(1) operation.
|
||||
pub fn last(&self) -> AnyhowResult<Option<FalTaskDetails>> {
|
||||
match self.queue.read() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire read lock: {:?}", err))
|
||||
}
|
||||
Ok(lock) => match lock.last() {
|
||||
None => Ok(None),
|
||||
Some((request_id, task_details)) => Ok(Some(task_details.clone()))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Insert a new task. If a task already exists, we update its value.
|
||||
/// This is an O(1) operation.
|
||||
pub fn insert(&self, request: &EnqueuedRequest) -> AnyhowResult<Option<FalTaskDetails>> {
|
||||
match self.queue.write() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
|
||||
}
|
||||
Ok(mut lock) => {
|
||||
// NB: Probably fine to compute TaskDetails here for now. We may eventually expose this if we include additional info.
|
||||
let task_details = FalTaskDetails::for_request(request.clone());
|
||||
Ok(lock.insert(request.request_id.clone(), task_details))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a task by ID.
|
||||
/// This is an O(n) operation.
|
||||
pub fn remove(&self, request_id: &FalRequestId) -> AnyhowResult<Option<FalTaskDetails>> {
|
||||
match self.queue.write() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
|
||||
}
|
||||
Ok(mut lock) => {
|
||||
Ok(lock.shift_remove(request_id))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Remove a list of tasks by ID.
|
||||
/// This is an O(n) operation.
|
||||
pub fn remove_list(&self, request_ids: &[&FalRequestId]) -> AnyhowResult<()> {
|
||||
match self.queue.write() {
|
||||
Err(err) => {
|
||||
Err(anyhow::anyhow!("Failed to acquire write lock: {:?}", err))
|
||||
}
|
||||
Ok(mut lock) => {
|
||||
// TODO: Cleaner way to do this?
|
||||
let mut set : HashSet<&FalRequestId, _> = HashSet::with_capacity(request_ids.len());
|
||||
set.extend(request_ids);
|
||||
lock.retain(|k, v| !set.contains(k));
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod fal_credential_manager;
|
||||
pub mod fal_task_queue;
|
||||
pub mod read_fal_api_key_from_disk;
|
||||
@@ -0,0 +1,187 @@
|
||||
use crate::events::sendable_event_trait::SendableEvent;
|
||||
use crate::events::sora::sora_image_generation_complete_event::SoraImageGenerationCompleteEvent;
|
||||
use crate::events::sora::sora_image_generation_failed_event::SoraImageGenerationFailedEvent;
|
||||
use crate::state::data_dir::app_data_root::AppDataRoot;
|
||||
use crate::state::data_dir::trait_data_subdir::DataSubdir;
|
||||
use crate::state::fal::fal_credential_manager::FalCredentialManager;
|
||||
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 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;
|
||||
use std::path::PathBuf;
|
||||
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, Emitter};
|
||||
use tempdir::TempDir;
|
||||
use fal_client::fal_error_plus::FalErrorPlus;
|
||||
use fal_client::utils::queue_status_checker::QueueStatusChecker;
|
||||
use tokens::tokens::media_files::MediaFileToken;
|
||||
|
||||
pub async fn fal_task_polling_thread(
|
||||
app_handle: AppHandle,
|
||||
app_data_root: AppDataRoot,
|
||||
fal_creds_manager: FalCredentialManager,
|
||||
storyteller_creds_manager: StorytellerCredentialManager,
|
||||
fal_task_queue: FalTaskQueue,
|
||||
) -> ! {
|
||||
loop {
|
||||
let res = polling_loop(
|
||||
&app_handle,
|
||||
&fal_creds_manager,
|
||||
&storyteller_creds_manager,
|
||||
&fal_task_queue,
|
||||
&app_data_root,
|
||||
).await;
|
||||
if let Err(err) = res {
|
||||
error!("An error occurred: {:?}", err);
|
||||
}
|
||||
tokio::time::sleep(std::time::Duration::from_millis(30_000)).await;
|
||||
}
|
||||
}
|
||||
|
||||
async fn polling_loop(
|
||||
app_handle: &AppHandle,
|
||||
fal_creds_manager: &FalCredentialManager,
|
||||
storyteller_creds_manager: &StorytellerCredentialManager,
|
||||
fal_task_queue: &FalTaskQueue,
|
||||
app_data_root: &AppDataRoot,
|
||||
) -> AnyhowResult<()> {
|
||||
loop {
|
||||
if fal_task_queue.is_empty()? {
|
||||
// No need to poll if we don't have pending tasks.
|
||||
tokio::time::sleep(std::time::Duration::from_millis(1_000)).await;
|
||||
continue;
|
||||
}
|
||||
|
||||
info!("Task queue has {} pending tasks.", fal_task_queue.len()?);
|
||||
|
||||
let key = fal_creds_manager.get_key_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 len = fal_task_queue.len()?;
|
||||
|
||||
for i in 0..len {
|
||||
let task = match fal_task_queue.get_index(i)? {
|
||||
Some(task) => task,
|
||||
None => {
|
||||
warn!("Failed to get task ID at index {}", i);
|
||||
break;
|
||||
}
|
||||
};
|
||||
|
||||
info!("Checking task status for task ID: {:?}", &task.enqueued_request.request_id);
|
||||
|
||||
let result = queue_status_checker.check_status(&task.enqueued_request).await;
|
||||
|
||||
match result {
|
||||
Err(err) => {
|
||||
warn!("Error checking job status: {:?}", err);
|
||||
tokio::time::sleep(std::time::Duration::from_millis(2_000)).await;
|
||||
continue;
|
||||
}
|
||||
Ok(status) => {
|
||||
info!("Status: {:?}", status);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
|
||||
|
||||
// // TODO: The cursoring logic likely needs to improve.
|
||||
// for task in response.task_responses {
|
||||
// let status = TaskStatus::from_str(&task.status);
|
||||
//
|
||||
// match status {
|
||||
// TaskStatus::Succeeded => {
|
||||
// succeeded_tasks.push(task);
|
||||
// }
|
||||
// TaskStatus::Failed => {
|
||||
// failed_tasks.push(task);
|
||||
// }
|
||||
// TaskStatus::Queued => {}
|
||||
// TaskStatus::Running => {}
|
||||
// TaskStatus::Unknown(_) => {}
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// for task in failed_tasks.iter() {
|
||||
// if fal_task_queue.contains_key(&task.id)? {
|
||||
// let event = SoraImageGenerationFailedEvent {
|
||||
// prompt: task.prompt.clone(),
|
||||
// };
|
||||
// event.send(&app_handle)?;
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// let failed_task_ids: Vec<&TaskId> = failed_tasks
|
||||
// .iter()
|
||||
// .map(|task| &task.id)
|
||||
// .collect();
|
||||
//
|
||||
// fal_task_queue.remove_list(&failed_task_ids)?;
|
||||
//
|
||||
// let creds = storyteller_creds_manager.get_credentials_required()?;
|
||||
// let api_host = ApiHost::Storyteller;
|
||||
//
|
||||
// for task in succeeded_tasks.iter() {
|
||||
// if !fal_task_queue.contains_key(&task.id)? {
|
||||
// continue;
|
||||
// }
|
||||
// info!("Task succeeded: {:?}", task.id);
|
||||
//
|
||||
// for (i, generation) in task.generations.iter().enumerate() {
|
||||
// info!("Downloading generated file...");
|
||||
// let download_path = download_generation(generation, &app_data_root).await?;
|
||||
//
|
||||
// info!("Uploading to backend...");
|
||||
// let result = upload_image_media_file_from_file(&api_host, Some(&creds), download_path).await?;
|
||||
//
|
||||
// info!("Uploaded to API backend: {:?}", result.media_file_token);
|
||||
//
|
||||
// let event = SoraImageGenerationCompleteEvent {
|
||||
// media_file_token: result.media_file_token,
|
||||
// };
|
||||
//
|
||||
// event.send(&app_handle)?;
|
||||
// }
|
||||
// }
|
||||
//
|
||||
// let succeeded_task_ids : Vec<&TaskId> = succeeded_tasks
|
||||
// .iter()
|
||||
// .map(|task| &task.id)
|
||||
// .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<PathBuf> {
|
||||
let url = Url::parse(&generation.url)?;
|
||||
|
||||
let response = reqwest::get(&generation.url).await?;
|
||||
let image_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 file = File::create(&download_path)?;
|
||||
file.write_all(&image_bytes)?;
|
||||
|
||||
Ok(download_path)
|
||||
}
|
||||
@@ -1,3 +1,4 @@
|
||||
pub mod discord_presence_thread;
|
||||
pub mod fal_task_polling_thread;
|
||||
pub mod main_window_thread;
|
||||
pub mod sora_task_polling_thread;
|
||||
@@ -10,6 +10,8 @@ pub enum FalErrorPlus {
|
||||
AnyhowError(anyhow::Error),
|
||||
/// URL parse errors.
|
||||
UrlParseError(url::ParseError),
|
||||
/// An endpoint we don't support yet.
|
||||
UnhandledEndpoint(String),
|
||||
}
|
||||
|
||||
impl Display for FalErrorPlus {
|
||||
@@ -18,6 +20,7 @@ impl Display for FalErrorPlus {
|
||||
Self::FalError(err) => write!(f, "FalErrorPlus::FalError: {:?}", err),
|
||||
Self::AnyhowError(err) => write!(f, "FalErrorPlus::AnyhowError: {:?}", err),
|
||||
Self::UrlParseError(err) => write!(f, "FalErrorPlus::UrlParseError: {:?}", err),
|
||||
Self::UnhandledEndpoint(endpoint) => write!(f, "FalErrorPlus::UnhandledEndpoint: {:?}", endpoint),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::model::fal_endpoint::FalEndpoint;
|
||||
use crate::model::fal_request_id::FalRequestId;
|
||||
use errors::AnyhowResult;
|
||||
use fal::queue::{Queue, QueueResponse};
|
||||
use reqwest::Client;
|
||||
@@ -8,7 +9,7 @@ use serde::de::DeserializeOwned;
|
||||
#[derive(Clone, Debug)]
|
||||
pub struct EnqueuedRequest {
|
||||
pub fal_endpoint: FalEndpoint,
|
||||
pub request_id: String,
|
||||
pub request_id: FalRequestId,
|
||||
pub response_url: String,
|
||||
pub status_url: String,
|
||||
pub cancel_url: String,
|
||||
@@ -19,7 +20,7 @@ impl EnqueuedRequest {
|
||||
let endpoint = FalEndpoint::from_queue_response(queue);
|
||||
Ok(Self {
|
||||
fal_endpoint: endpoint.clone(),
|
||||
request_id: queue.payload.request_id.clone(),
|
||||
request_id: FalRequestId(queue.payload.request_id.clone()),
|
||||
response_url: queue.payload.response_url.clone(),
|
||||
status_url: queue.payload.status_url.clone(),
|
||||
cancel_url: queue.payload.cancel_url.clone(),
|
||||
@@ -30,7 +31,7 @@ impl EnqueuedRequest {
|
||||
let client = Client::new();
|
||||
let endpoint = self.fal_endpoint.url();
|
||||
let payload = QueueResponse {
|
||||
request_id: self.request_id.clone(),
|
||||
request_id: self.request_id.0.clone(),
|
||||
response_url: self.response_url.clone(),
|
||||
status_url: self.status_url.clone(),
|
||||
cancel_url: self.cancel_url.clone(),
|
||||
|
||||
@@ -0,0 +1,13 @@
|
||||
use std::fmt::Display;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use crate::model::fal_endpoint::FalEndpoint;
|
||||
|
||||
/// A strongly typed request ID for the Fal API.
|
||||
#[derive(Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize, Deserialize)]
|
||||
pub struct FalRequestId(pub String);
|
||||
|
||||
impl Display for FalRequestId {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
Display::fmt(&self.0, f)
|
||||
}
|
||||
}
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod enqueued_request;
|
||||
pub mod fal_endpoint;
|
||||
pub mod fal_endpoint;
|
||||
pub mod fal_request_id;
|
||||
@@ -1,2 +1,2 @@
|
||||
pub mod remove_background_rembg;
|
||||
mod enqueue_kling_16_image_to_video;
|
||||
pub mod enqueue_kling_16_image_to_video;
|
||||
pub mod remove_background_rembg;
|
||||
@@ -1,2 +1,3 @@
|
||||
pub mod clone_queue_payload;
|
||||
pub mod clone_queue_response;
|
||||
pub mod clone_queue_response;
|
||||
pub mod queue_status_checker;
|
||||
@@ -0,0 +1,48 @@
|
||||
use fal::endpoints::fal_ai::kling_video::v1_6::pro::effects::I2VOutput;
|
||||
use fal::FalError;
|
||||
use fal::prelude::QueueResponse;
|
||||
use fal::queue::{Queue, QueueStatus};
|
||||
use reqwest::Client;
|
||||
use crate::creds::fal_api_key::FalApiKey;
|
||||
use crate::fal_error_plus::FalErrorPlus;
|
||||
use crate::model::enqueued_request::EnqueuedRequest;
|
||||
use crate::model::fal_endpoint::FalEndpoint;
|
||||
|
||||
pub struct QueueStatusChecker {
|
||||
api_key: FalApiKey,
|
||||
}
|
||||
|
||||
impl QueueStatusChecker {
|
||||
pub fn new(api_key: FalApiKey) -> Self {
|
||||
Self { api_key }
|
||||
}
|
||||
|
||||
pub async fn check_status(&self, request: &EnqueuedRequest) -> Result<QueueStatus, FalErrorPlus> {
|
||||
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)
|
||||
}
|
||||
}
|
||||
Reference in new issue
Block a user