Compare commits

...
112 changed files with 1617 additions and 759 deletions
@@ -5,8 +5,9 @@ use fal::endpoints::fal_ai::kling_video::v1_6::pro::image_to_video::{image_to_vi
use fal::webhook::WebhookResponse;
use reqwest::IntoUrl;
pub struct Kling16ProArgs<'a, U: IntoUrl, V: IntoUrl> {
pub struct Kling16ProArgs<'a, U: IntoUrl, T: IntoUrl, V: IntoUrl> {
pub image_url: U,
pub end_frame_image_url: Option<T>,
pub webhook_url: V,
pub prompt: &'a str,
pub api_key: &'a FalApiKey,
@@ -28,8 +29,8 @@ pub enum Kling16ProAspectRatio {
TallNineSixteen, // 9:16
}
pub async fn enqueue_kling_16_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>(
args: Kling16ProArgs<'_, U, V>
pub async fn enqueue_kling_16_pro_image_to_video_webhook<U: IntoUrl, T: IntoUrl, V: IntoUrl>(
args: Kling16ProArgs<'_, U, T, V>
) -> Result<WebhookResponse, FalErrorPlus> {
let duration = match args.duration {
Kling16Duration::Default => None,
@@ -45,15 +46,18 @@ pub async fn enqueue_kling_16_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>
let image_url = args.image_url.as_str().to_string();
let tail_image_url = args.end_frame_image_url
.map(|url| url.as_str().to_string());
let request = ProImageToVideoRequest {
image_url,
tail_image_url,
prompt: args.prompt.to_string(),
aspect_ratio,
duration,
// Maybe expose these later
cfg_scale: None,
negative_prompt: None,
tail_image_url: None,
};
let result = image_to_video(request)
@@ -64,35 +68,34 @@ pub async fn enqueue_kling_16_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>
result.map_err(|err| classify_fal_error(err))
}
/*
#[cfg(test)]
mod tests {
use crate::creds::fal_api_key::FalApiKey;
use crate::requests::webhook::video::enqueue_kling_16_image_to_video_webhook::{enqueue_kling_16_image_to_video_webhook, Kling16Args, Kling16Duration};
use crate::requests::webhook::video::enqueue_kling_16_pro_image_to_video_webhook::{enqueue_kling_16_pro_image_to_video_webhook, Kling16Duration, Kling16ProArgs, Kling16ProAspectRatio};
use errors::AnyhowResult;
use std::fs::read_to_string;
use testing::test_file_path::test_file_path;
use test_data::web::image_urls::{JUNO_AT_LAKE_IMAGE_URL, TALL_MOCHI_WITH_GLASSES_IMAGE_URL};
#[tokio::test]
#[ignore]
async fn test_kling16_video() -> AnyhowResult<()> {
let image = test_file_path("test_data/image/juno.jpg")?;
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 = Kling16Args {
image_path: image,
prompt: "a corgi looks out over the water",
let args = Kling16ProArgs {
image_url: TALL_MOCHI_WITH_GLASSES_IMAGE_URL,
end_frame_image_url: Some(JUNO_AT_LAKE_IMAGE_URL.to_string()),
prompt: "shiba in glasses runs to the lake and stands by the shore",
api_key: &api_key,
duration: Kling16Duration::Default,
aspect_ratio: Kling16ProAspectRatio::WideSixteenNine,
webhook_url: "https://example.com/webhook",
};
let result = enqueue_kling_16_image_to_video_webhook(args).await?;
let result = enqueue_kling_16_pro_image_to_video_webhook(args).await?;
Ok(())
}
}
*/
@@ -5,8 +5,9 @@ use fal::endpoints::fal_ai::kling_video::v2_1::pro::image_to_video::{image_to_vi
use fal::webhook::WebhookResponse;
use reqwest::IntoUrl;
pub struct Kling21ProArgs<'a, U: IntoUrl, V: IntoUrl> {
pub struct Kling21ProArgs<'a, U: IntoUrl, T: IntoUrl, V: IntoUrl> {
pub image_url: U,
pub end_frame_image_url: Option<T>,
pub webhook_url: V,
pub prompt: &'a str,
pub api_key: &'a FalApiKey,
@@ -28,8 +29,8 @@ pub enum Kling21ProAspectRatio {
TallNineSixteen, // 9:16
}
pub async fn enqueue_kling_21_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>(
args: Kling21ProArgs<'_, U, V>
pub async fn enqueue_kling_21_pro_image_to_video_webhook<U: IntoUrl, T: IntoUrl, V: IntoUrl>(
args: Kling21ProArgs<'_, U, T, V>
) -> Result<WebhookResponse, FalErrorPlus> {
let duration = match args.duration {
Kling21ProDuration::Default => None,
@@ -45,15 +46,18 @@ pub async fn enqueue_kling_21_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>
let image_url = args.image_url.as_str().to_string();
let tail_image_url = args.end_frame_image_url
.map(|url| url.as_str().to_string());
let request = ProImageToVideoRequest {
image_url,
prompt: args.prompt.to_string(),
tail_image_url,
aspect_ratio,
duration,
// Maybe expose these later
cfg_scale: None,
negative_prompt: None,
tail_image_url: None,
};
let result = image_to_video(request)
@@ -64,27 +68,26 @@ pub async fn enqueue_kling_21_pro_image_to_video_webhook<U: IntoUrl, V: IntoUrl>
result.map_err(|err| classify_fal_error(err))
}
#[cfg(test)]
mod tests {
use crate::creds::fal_api_key::FalApiKey;
use crate::requests::webhook::video::enqueue_kling_21_pro_image_to_video_webhook::{enqueue_kling_21_pro_image_to_video_webhook, Kling21ProArgs, Kling21ProAspectRatio, Kling21ProDuration};
use errors::AnyhowResult;
use std::fs::read_to_string;
use test_data::web::image_urls::{JUNO_AT_LAKE_IMAGE_URL, SUPER_WIDE_FALL_MOUNTAINS_IMAGE_URL};
#[tokio::test]
#[ignore]
async fn test_kling21_pro_video() -> AnyhowResult<()> {
let image_url = "https://cdn-2.fakeyou.com/media/3/4/h/f/s/34hfsmt8e38rvne6mwa4pwbxr6292sgy/image_34hfsmt8e38rvne6mwa4pwbxr6292sgy.png";
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 = Kling21ProArgs {
image_url: image_url,
prompt: "a shot of the mountains, the camera rotates around to show the mountain range",
image_url: SUPER_WIDE_FALL_MOUNTAINS_IMAGE_URL,
end_frame_image_url: Some(JUNO_AT_LAKE_IMAGE_URL.to_string()),
prompt: "a shot of the mountains, the camera pulls back to show a corgi waiting to jump into the lake",
api_key: &api_key,
duration: Kling21ProDuration::Default,
aspect_ratio: Kling21ProAspectRatio::WideSixteenNine,
@@ -6,8 +6,9 @@ use fal::endpoints::fal_ai::bytedance::seedance::v1::lite::image_to_video::Image
use fal::webhook::WebhookResponse;
use reqwest::IntoUrl;
pub struct Seedance1LiteArgs<'a, U: IntoUrl, V: IntoUrl> {
pub struct Seedance1LiteArgs<'a, U: IntoUrl, T: IntoUrl, V: IntoUrl> {
pub image_url: U,
pub end_frame_image_url: Option<T>,
pub webhook_url: V,
pub prompt: &'a str,
pub api_key: &'a FalApiKey,
@@ -29,8 +30,8 @@ pub enum Seedance1LiteResolution {
SevenTwentyP, // 720p
}
pub async fn enqueue_seedance_1_lite_image_to_video_webhook<U: IntoUrl, V: IntoUrl>(
args: Seedance1LiteArgs<'_, U, V>
pub async fn enqueue_seedance_1_lite_image_to_video_webhook<U: IntoUrl, T: IntoUrl, V: IntoUrl>(
args: Seedance1LiteArgs<'_, U, T, V>
) -> Result<WebhookResponse, FalErrorPlus> {
let duration = match args.duration {
Seedance1LiteDuration::FiveSeconds => Some("5".to_string()),
@@ -44,8 +45,12 @@ pub async fn enqueue_seedance_1_lite_image_to_video_webhook<U: IntoUrl, V: IntoU
let image_url = args.image_url.as_str().to_string();
let end_image_url = args.end_frame_image_url
.map(|url| url.as_str().to_string());
let request = ImageToVideoRequest {
image_url,
end_image_url,
prompt: args.prompt.to_string(),
duration,
resolution,
@@ -69,20 +74,20 @@ mod tests {
use crate::requests::webhook::video::enqueue_seedance_1_lite_image_to_video_webhook::{enqueue_seedance_1_lite_image_to_video_webhook, Seedance1LiteArgs, Seedance1LiteDuration, Seedance1LiteResolution};
use errors::AnyhowResult;
use std::fs::read_to_string;
use test_data::web::image_urls::{JUNO_AT_LAKE_IMAGE_URL, TALL_MOCHI_WITH_GLASSES_IMAGE_URL};
#[tokio::test]
#[ignore]
async fn test_kling21_pro_video() -> AnyhowResult<()> {
let image_url = "https://cdn-2.fakeyou.com/media/3/4/h/f/s/34hfsmt8e38rvne6mwa4pwbxr6292sgy/image_34hfsmt8e38rvne6mwa4pwbxr6292sgy.png";
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 = Seedance1LiteArgs {
image_url: image_url,
prompt: "a shot of the mountains as the sun sets and reveals the moon and stars",
image_url: TALL_MOCHI_WITH_GLASSES_IMAGE_URL,
end_frame_image_url: Some(JUNO_AT_LAKE_IMAGE_URL.to_string()),
prompt: "shiba in glasses runs to the lake and stands by the shore",
api_key: &api_key,
camera_fixed: false,
duration: Seedance1LiteDuration::FiveSeconds,
@@ -11,8 +11,6 @@ use crate::core::events::generation_events::common::{GenerationAction, Generatio
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;
@@ -6,8 +6,6 @@ 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, 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;
@@ -10,8 +10,6 @@ use crate::core::events::generation_events::common::{GenerationAction, Generatio
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;
@@ -27,6 +25,7 @@ use crate::services::storyteller::state::storyteller_credential_manager::Storyte
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};
@@ -41,9 +40,20 @@ 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;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ContextualImageEditModel {
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "flux_pro_kontext_max")]
FluxProKontextMax,
}
#[derive(Deserialize, Debug)]
pub struct EnqueueContextualEditImageCommand {
/// The model to use.
@@ -4,8 +4,6 @@ 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;
@@ -5,8 +5,6 @@ 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;
@@ -8,8 +8,6 @@ use crate::core::events::functional_events::show_provider_login_modal_event::Sho
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, 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::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;
@@ -1,7 +1,9 @@
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_dev_juggernaut_inpaint::handle_flux_dev_juggernaut_inpaint::handle_flux_dev_juggernaut_inpaint;
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::image_inpaint::flux_pro_kontext_inpaint::handle_flux_pro_kontext_inpaint::handle_flux_pro_kontext_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};
@@ -11,8 +13,6 @@ use crate::core::events::generation_events::common::{GenerationAction, Generatio
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;
@@ -44,15 +44,20 @@ 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_inpaint::flux_dev_juggernaut_inpaint::handle_flux_dev_juggernaut_inpaint::handle_flux_dev_juggernaut_inpaint;
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ImageInpaintModel {
// Mask-based inpainting models
#[serde(rename = "flux_dev_juggernaut")]
FluxDevJuggernaut,
#[serde(rename = "flux_pro_1")]
FluxPro1,
// Non-inpainting, instructive editing models
// NB: We're supporting these to keep the app simple and not over-complicate the javascript
#[serde(rename = "flux_pro_kontext_max")]
FluxProKontextMax,
}
#[derive(Deserialize, Debug)]
@@ -244,6 +249,21 @@ pub async fn handle_request(
sora_task_queue,
).await?
}
Some(ImageInpaintModel::FluxProKontextMax) => {
// Redirect to contextual image edit handler.
handle_flux_pro_kontext_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
@@ -7,8 +7,6 @@ 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;
@@ -7,8 +7,6 @@ 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;
@@ -132,7 +130,7 @@ pub async fn handle_flux_pro_1_inpaint_artcraft(
Ok(TaskEnqueueSuccess {
provider: GenerationProvider::Artcraft,
model: Some(GenerationModel::FluxProKontextMax),
model: Some(GenerationModel::FluxPro1),
provider_job_id: Some(job_id.to_string()),
task_type: TaskType::ImageInpaintEdit,
})
@@ -0,0 +1,59 @@
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_inpaint::enqueue_image_inpaint_command::EnqueueInpaintImageCommand;
use crate::core::commands::enqueue::image_inpaint::errors::InternalImageInpaintError;
use crate::core::commands::enqueue::image_inpaint::flux_pro_kontext_inpaint::handle_flux_pro_kontext_inpaint_artcraft::handle_flux_pro_kontext_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 async fn handle_flux_pro_kontext_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 kontext inpaint via Artcraft...");
return handle_flux_pro_kontext_inpaint_artcraft(
request,
app,
app_data_root,
app_env_configs,
storyteller_creds_manager
).await;
}
}
}
Err(InternalImageInpaintError::NoProviderAvailable)
}
@@ -0,0 +1,183 @@
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_edit::errors::InternalContextualEditImageError;
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::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::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 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::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 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_kontext_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)
//})?;
// TODO: Experiment in seeing if we can use the mask in any useful way.
// let mask_media_token = get_mask(
// request,
// app_env_configs,
// &creds,
// ).await?;
info!("Calling Artcraft flux pro kontext inpaint ...");
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(InternalImageInpaintError::InvalidNumberOfRequestedImages {
min: 1,
max: 4,
requested: other,
});
},
};
let request = FluxProKontextMaxEditImageRequest {
uuid_idempotency_token,
prompt: Some(request.prompt.clone()),
image_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 flux pro kontext inpaint. Job token: {}",
enqueued.inference_job_token);
enqueued.inference_job_token
}
Err(err) => {
error!("Failed to use Artcraft flux pro kontext 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::ImageInpaintEdit,
})
}
// 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)
// }
@@ -0,0 +1,2 @@
pub (super) mod handle_flux_pro_kontext_inpaint;
pub (super) mod handle_flux_pro_kontext_inpaint_artcraft;
@@ -1,4 +1,5 @@
mod flux_dev_juggernaut_inpaint;
mod flux_pro_1_inpaint;
mod flux_pro_kontext_inpaint;
pub mod enqueue_image_inpaint_command;
pub mod errors;
@@ -6,7 +6,6 @@ use crate::core::commands::response::shorthand::Response;
use crate::core::commands::response::success_response_wrapper::SerializeMarker;
use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent;
use crate::core::model::video_models::VideoModel;
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};
@@ -21,16 +20,45 @@ use serde_derive::{Deserialize, Serialize};
use tauri::{AppHandle, State};
use tokens::tokens::media_files::MediaFileToken;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum VideoModel {
#[serde(rename = "kling_1.6_pro")]
Kling16Pro,
#[serde(rename = "kling_2.1_pro")]
Kling21Pro,
#[serde(rename = "kling_2.1_master")]
Kling21Master,
#[serde(rename = "seedance_1.0_lite")]
Seedance10Lite,
#[serde(rename = "veo_2")]
Veo2,
}
#[derive(Deserialize)]
pub struct EnqueueImageToVideoRequest {
/// Image media file; the image to remove the background from.
/// TODO: In the future we may support base64 images, URLs, or file paths here.
pub image_media_token: Option<MediaFileToken>,
/// REQUIRED.
/// The model to use.
pub model: Option<VideoModel>,
/// Optional text prompt.
/// Currently REQUIRED.
/// Image media file; the starting frame of the video.
/// TODO: In the future we may support base64 images, URLs, or file paths here.
pub image_media_token: Option<MediaFileToken>,
/// Optional.
/// Image media file; the image to remove the background from.
/// TODO: In the future we may support base64 images, URLs, or file paths here.
pub end_frame_image_media_token: Option<MediaFileToken>,
/// Optional.
/// Text prompt used to direct the video.
pub prompt: Option<String>,
}
@@ -1,11 +1,10 @@
use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::EnqueueImageToVideoRequest;
use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::{EnqueueImageToVideoRequest, VideoModel};
use crate::core::commands::enqueue::image_to_video::internal_video_error::InternalVideoError;
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::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent;
use crate::core::model::video_models::VideoModel;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
@@ -76,6 +75,7 @@ pub async fn handle_video_artcraft(
let request = GenerateKling16ProImageToVideoRequest {
uuid_idempotency_token,
media_file_token: request.image_media_token,
end_frame_image_media_token: request.end_frame_image_media_token,
aspect_ratio: Some(GenerateKling16ProAspectRatio::WideSixteenNine),
prompt: request.prompt,
duration: None,
@@ -128,6 +128,7 @@ pub async fn handle_video_artcraft(
let request = GenerateKling21ProImageToVideoRequest {
uuid_idempotency_token,
media_file_token: request.image_media_token,
end_frame_image_media_token: request.end_frame_image_media_token,
aspect_ratio: Some(GenerateKling21ProAspectRatio::WideSixteenNine),
prompt: request.prompt,
duration: None,
@@ -154,6 +155,7 @@ pub async fn handle_video_artcraft(
let request = GenerateSeedance10LiteImageToVideoRequest {
uuid_idempotency_token,
media_file_token: request.image_media_token,
end_frame_image_media_token: request.end_frame_image_media_token,
prompt: request.prompt,
resolution: None,
duration: None,
@@ -1,11 +1,10 @@
use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::EnqueueImageToVideoRequest;
use crate::core::commands::enqueue::image_to_video::enqueue_image_to_video_command::{EnqueueImageToVideoRequest, VideoModel};
use crate::core::commands::enqueue::image_to_video::internal_video_error::InternalVideoError;
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::events::generation_events::generation_enqueue_success_event::GenerationEnqueueSuccessEvent;
use crate::core::model::video_models::VideoModel;
use crate::core::state::app_env_configs::app_env_configs::AppEnvConfigs;
use crate::core::state::data_dir::app_data_root::AppDataRoot;
use crate::core::utils::download_media_file_to_temp_dir::download_media_file_to_temp_dir;
@@ -12,7 +12,6 @@ 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_success_event::GenerationEnqueueSuccessEvent;
use crate::core::events::sendable_event_trait::SendableEvent;
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::{Provider, ProviderPriorityStore};
@@ -26,12 +25,35 @@ use crate::services::storyteller::state::storyteller_credential_manager::Storyte
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 fal_client::requests::queue::image_gen::enqueue_flux_pro_11_ultra_text_to_image::{enqueue_flux_pro_11_ultra_text_to_image, FluxPro11UltraTextToImageArgs};
use log::{error, info, warn};
use serde::{Deserialize, Serialize};
use sqlite_tasks::queries::create_task::{create_task, CreateTaskArgs};
use tauri::{AppHandle, State};
use enums::tauri::ux::tauri_command_caller::TauriCommandCaller;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ImageModel {
#[serde(rename = "flux_1_dev")]
Flux1Dev,
#[serde(rename = "flux_1_schnell")]
Flux1Schnell,
#[serde(rename = "flux_pro_11")]
FluxPro11,
#[serde(rename = "flux_pro_11_ultra")]
FluxPro11Ultra,
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "recraft_3")]
Recraft3,
// Generic Midjourney model, version unknown.
#[serde(rename = "midjourney")]
Midjourney,
}
#[derive(Deserialize)]
pub struct EnqueueTextToImageRequest {
@@ -1,10 +1,9 @@
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, TextToImageSize};
use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, ImageModel, TextToImageSize};
use crate::core::commands::enqueue::text_to_image::internal_image_error::InternalImageError;
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::services::storyteller::state::storyteller_credential_manager::StorytellerCredentialManager;
@@ -1,10 +1,9 @@
use crate::core::commands::enqueue::task_enqueue_success::TaskEnqueueSuccess;
use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::EnqueueTextToImageRequest;
use crate::core::commands::enqueue::text_to_image::enqueue_text_to_image_command::{EnqueueTextToImageRequest, ImageModel};
use crate::core::commands::enqueue::text_to_image::internal_image_error::InternalImageError;
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::services::fal::state::fal_credential_manager::FalCredentialManager;
use crate::services::fal::state::fal_task_queue::FalTaskQueue;
use anyhow::anyhow;
@@ -9,7 +9,6 @@ use crate::core::commands::enqueue::text_to_image::internal_image_error::Interna
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::{Provider, ProviderPriorityStore};
@@ -6,7 +6,6 @@ use crate::core::commands::enqueue::text_to_image::internal_image_error::Interna
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;
@@ -5,7 +5,6 @@ use crate::core::events::basic_sendable_event_trait::BasicSendableEvent;
use crate::core::events::functional_events::show_provider_login_modal_event::ShowProviderLoginModalEvent;
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::services::sora::state::sora_credential_manager::SoraCredentialManager;
use crate::services::sora::state::sora_task_queue::SoraTaskQueue;
use enums::common::generation_provider::GenerationProvider;
@@ -9,7 +9,6 @@ use crate::core::events::functional_events::show_provider_login_modal_event::Sho
use crate::core::events::generation_events::common::{GenerationAction, GenerationModel, GenerationServiceProvider};
use crate::core::events::generation_events::generation_enqueue_failure_event::GenerationEnqueueFailureEvent;
use crate::core::events::warning_events::flash_user_input_error_event::FlashUserInputErrorEvent;
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;
-1
View File
@@ -2,7 +2,6 @@ pub mod artcraft_error;
pub mod commands;
pub mod events;
pub mod lifecycle;
pub mod model;
pub mod state;
pub mod threads;
pub mod utils;
@@ -1,15 +0,0 @@
use serde_derive::Deserialize;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ContextualImageEditModel {
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "flux_pro_kontext_max")]
FluxProKontextMax,
// TODO: Flux Kontext models.
}
@@ -1,24 +0,0 @@
use serde_derive::Deserialize;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum ImageModel {
#[serde(rename = "flux_1_dev")]
Flux1Dev,
#[serde(rename = "flux_1_schnell")]
Flux1Schnell,
#[serde(rename = "flux_pro_11")]
FluxPro11,
#[serde(rename = "flux_pro_11_ultra")]
FluxPro11Ultra,
#[serde(rename = "gpt_image_1")]
GptImage1,
#[serde(rename = "recraft_3")]
Recraft3,
// Generic Midjourney model, version unknown.
#[serde(rename = "midjourney")]
Midjourney,
}
@@ -1,3 +0,0 @@
pub mod contextual_image_edit_models;
pub mod image_models;
pub mod video_models;
@@ -1,22 +0,0 @@
use serde_derive::Deserialize;
/// This is used in the Tauri command bridge.
/// Don't change the serializations without coordinating with the frontend.
#[derive(Deserialize, Debug, Copy, Clone)]
#[serde(rename_all = "snake_case")]
pub enum VideoModel {
#[serde(rename = "kling_1.6_pro")]
Kling16Pro,
#[serde(rename = "kling_2.1_pro")]
Kling21Pro,
#[serde(rename = "kling_2.1_master")]
Kling21Master,
#[serde(rename = "seedance_1.0_lite")]
Seedance10Lite,
#[serde(rename = "veo_2")]
Veo2,
}
@@ -40,7 +40,6 @@ pub async fn maybe_handle_text_to_image_complete_event(
batch_token,
).await?;
if let Err(err) = event.send(&app) {
error!("Failed to send TextToImageGenerationCompleteEvent: {:?}", err); // Fail open
}
+1 -1
View File
@@ -22,7 +22,7 @@ pub const MOUNTAIN_TREE_IMAGE_URL :&str = "https://cdn-2.fakeyou.com/media/3/4/h
pub const WHITE_HOUSE_SUNSET_IMAGE_URL : &str = "https://cdn-2.fakeyou.com/media/q/t/z/w/k/qtzwka8dd0g2nfyk94k1qh3zzfwqtfhn/asset_qtzwka8dd0g2nfyk94k1qh3zzfwqtfhn.jpg";
// media token: m_z61yg5w7dnnvhb1f9sqt8jf02y61bv
pub const WIDE_FALL_MOUNTAINS_IMAGE_URL : &str = "https://cdn-2.fakeyou.com/media/w/6/m/j/p/w6mjpcqkky7f2f3xkyvkcv3qeedkmejs/asset_w6mjpcqkky7f2f3xkyvkcv3qeedkmejs.jpg";
pub const SUPER_WIDE_FALL_MOUNTAINS_IMAGE_URL: &str = "https://cdn-2.fakeyou.com/media/w/6/m/j/p/w6mjpcqkky7f2f3xkyvkcv3qeedkmejs/asset_w6mjpcqkky7f2f3xkyvkcv3qeedkmejs.jpg";
// media token: m_m1bz02z1kkzanxy6rb4vk1kvq9de9g
pub const JUNO_AT_LAKE_IMAGE_URL : &str = "https://cdn-2.fakeyou.com/media/s/v/c/p/p/svcpp441gyb6mjtt8413kf5zxw1ehy4t/upload_svcpp441gyb6mjtt8413kf5zxw1ehy4t.jpg";
@@ -13,7 +13,11 @@ pub struct GenerateKling16ProImageToVideoRequest {
/// Source of the image file to convert to video.
/// It must be an image.
pub media_file_token: Option<MediaFileToken>,
/// Optional.
/// Image media file; the image to remove the background from.
pub end_frame_image_media_token: Option<MediaFileToken>,
/// Optional text prompt.
pub prompt: Option<String>,
@@ -13,6 +13,10 @@ pub struct GenerateKling21ProImageToVideoRequest {
/// Source of the image file to convert to video.
/// It must be an image.
pub media_file_token: Option<MediaFileToken>,
/// Optional.
/// Image media file; the image to remove the background from.
pub end_frame_image_media_token: Option<MediaFileToken>,
/// Optional text prompt.
pub prompt: Option<String>,
@@ -13,7 +13,11 @@ pub struct GenerateSeedance10LiteImageToVideoRequest {
/// Source of the image file to convert to video.
/// It must be an image.
pub media_file_token: Option<MediaFileToken>,
/// Optional.
/// Image media file; the image to remove the background from.
pub end_frame_image_media_token: Option<MediaFileToken>,
/// Optional text prompt.
pub prompt: Option<String>,
@@ -27,6 +27,7 @@ use errors::AnyhowResult;
pub mod billing;
pub mod configs;
pub mod error;
pub mod http_server;
pub mod state;
pub mod threads;
@@ -30,6 +30,7 @@ use crate::configs::static_api_tokens::StaticApiTokenSet;
pub mod billing;
pub mod configs;
pub mod error;
pub mod http_server;
pub mod state;
pub mod threads;
@@ -0,0 +1,37 @@
use std::error::Error;
use std::fmt::{Debug, Display, Formatter};
/// Any generic errors arising internally.
/// An error class that we do not expose via the API.
/// This should be massaged before we send it to the user.
#[derive(Debug)]
pub enum InternalError {
/// Issue with HMAC decoding
/// This is typically when using development cookies against production,
/// or vice versa.
VisitorCookieError(String),
/// Decode error with AVT cookie missing a required field.
VisitorCookieMissingField(&'static str),
/// Catch instances of anyhow::Error.
AnyhowError(anyhow::Error),
}
impl Error for InternalError {}
impl Display for InternalError {
fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
match self {
InternalError::VisitorCookieError(reason) => write!(f, "Visitor cookie error: {:?}", reason),
InternalError::VisitorCookieMissingField(field) => write!(f, "Visitor cookie missing field: {}", field),
InternalError::AnyhowError(err) => write!(f, "Anyhow error: {}", err),
}
}
}
impl From<anyhow::Error> for InternalError {
fn from(err: anyhow::Error) -> Self {
InternalError::AnyhowError(err)
}
}
@@ -0,0 +1 @@
pub mod internal_error;
@@ -16,21 +16,6 @@ const QUALITY : u8 = 95;
pub struct MediaLinksBuilder {}
impl MediaLinksBuilder {
#[deprecated(note = "we need to start passing the domain and the server environment to generate the CDN")]
pub fn from_media_path(
domain: MediaDomain,
bucket_path: &MediaFileBucketPath,
) -> MediaLinks {
Self::from_media_path_and_env(domain, ServerEnvironment::Production, bucket_path)
}
#[deprecated(note = "we need to start passing the domain and the server environment to generate the CDN")]
pub fn from_rooted_path(
domain: MediaDomain,
rooted_path: &str,
) -> MediaLinks {
Self::from_rooted_path_and_env(domain, ServerEnvironment::Production, rooted_path)
}
pub fn from_media_path_and_env(
domain: MediaDomain,
@@ -148,7 +133,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn wav_file() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.wav");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.wav");
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/foo/bar.wav");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -156,7 +141,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn glb_file() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.glb");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.glb");
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/foo/bar.glb");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -164,7 +149,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn jpg_image() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.jpg");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.jpg");
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/foo/bar.jpg");
assert_eq!(links.maybe_thumbnail_template, Some("https://cdn-2.fakeyou.com/cdn-cgi/image/width={WIDTH},quality=95/foo/bar.jpg".to_string()));
assert_eq!(links.maybe_video_previews, None);
@@ -172,7 +157,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn mp4_video() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.mp4");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.mp4");
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/foo/bar.mp4");
assert_eq!(links.maybe_thumbnail_template, None);
let video_previews = links.maybe_video_previews.expect("should have previews");
@@ -190,7 +175,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
fn wav_file() {
// https://storage.googleapis.com/vocodes-public/media/9/4/a/2/7/94a27nmbd0bqmd10tg0pp3hz45zytf67/fakeyou_94a27nmbd0bqmd10tg0pp3hz45zytf67.wav
let media_path = MediaFileBucketPath::from_object_hash("94a27nmbd0bqmd10tg0pp3hz45zytf67", Some("fakeyou_"), Some(".wav"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/media/9/4/a/2/7/94a27nmbd0bqmd10tg0pp3hz45zytf67/fakeyou_94a27nmbd0bqmd10tg0pp3hz45zytf67.wav");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -200,7 +185,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
fn png_image() {
/// https://storage.googleapis.com/vocodes-public/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png
let media_path = MediaFileBucketPath::from_object_hash("37mb3gh8fmj85y21thvbv08bzv24atjt", Some("upload_"), Some(".png"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png");
assert_eq!(links.maybe_thumbnail_template, Some("https://cdn-2.fakeyou.com/cdn-cgi/image/width={WIDTH},quality=95/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png".to_string()));
assert_eq!(links.maybe_video_previews, None);
@@ -210,7 +195,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
fn mp4_video() {
// https://storage.googleapis.com/vocodes-public/media/t/6/c/n/y/t6cnyw4g3e8k7carkk2bvrt6nd3fycjv/storyteller_t6cnyw4g3e8k7carkk2bvrt6nd3fycjv.mp4
let media_path = MediaFileBucketPath::from_object_hash("t6cnyw4g3e8k7carkk2bvrt6nd3fycjv", Some("storyteller_"), Some(".mp4"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn-2.fakeyou.com/media/t/6/c/n/y/t6cnyw4g3e8k7carkk2bvrt6nd3fycjv/storyteller_t6cnyw4g3e8k7carkk2bvrt6nd3fycjv.mp4");
assert_eq!(links.maybe_thumbnail_template, None);
let video_previews = links.maybe_video_previews.expect("should have previews");
@@ -229,10 +214,11 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
mod rooted_path {
use super::*;
use server_environment::ServerEnvironment;
#[test]
fn wav_file() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.wav");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.wav");
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/foo/bar.wav");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -240,7 +226,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn glb_file() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.glb");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.glb");
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/foo/bar.glb");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -248,7 +234,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn jpg_image() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.jpg");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.jpg");
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/foo/bar.jpg");
assert_eq!(links.maybe_thumbnail_template, Some("https://cdn.storyteller.ai/cdn-cgi/image/width={WIDTH},quality=95/foo/bar.jpg".to_string()));
assert_eq!(links.maybe_video_previews, None);
@@ -256,7 +242,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
#[test]
fn mp4_video() {
let links = MediaLinksBuilder::from_rooted_path(DOMAIN, "/foo/bar.mp4");
let links = MediaLinksBuilder::from_rooted_path_and_env(DOMAIN, ServerEnvironment::Production, "/foo/bar.mp4");
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/foo/bar.mp4");
assert_eq!(links.maybe_thumbnail_template, None);
let video_previews = links.maybe_video_previews.expect("should have previews");
@@ -269,12 +255,13 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
mod media_path {
use super::*;
use server_environment::ServerEnvironment;
#[test]
fn wav_file() {
// https://storage.googleapis.com/vocodes-public/media/9/4/a/2/7/94a27nmbd0bqmd10tg0pp3hz45zytf67/fakeyou_94a27nmbd0bqmd10tg0pp3hz45zytf67.wav
let media_path = MediaFileBucketPath::from_object_hash("94a27nmbd0bqmd10tg0pp3hz45zytf67", Some("fakeyou_"), Some(".wav"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/media/9/4/a/2/7/94a27nmbd0bqmd10tg0pp3hz45zytf67/fakeyou_94a27nmbd0bqmd10tg0pp3hz45zytf67.wav");
assert_eq!(links.maybe_thumbnail_template, None);
assert_eq!(links.maybe_video_previews, None);
@@ -284,7 +271,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
fn png_image() {
/// https://storage.googleapis.com/vocodes-public/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png
let media_path = MediaFileBucketPath::from_object_hash("37mb3gh8fmj85y21thvbv08bzv24atjt", Some("upload_"), Some(".png"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png");
assert_eq!(links.maybe_thumbnail_template, Some("https://cdn.storyteller.ai/cdn-cgi/image/width={WIDTH},quality=95/media/3/7/m/b/3/37mb3gh8fmj85y21thvbv08bzv24atjt/upload_37mb3gh8fmj85y21thvbv08bzv24atjt.png".to_string()));
assert_eq!(links.maybe_video_previews, None);
@@ -294,7 +281,7 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
fn mp4_video() {
// https://storage.googleapis.com/vocodes-public/media/t/6/c/n/y/t6cnyw4g3e8k7carkk2bvrt6nd3fycjv/storyteller_t6cnyw4g3e8k7carkk2bvrt6nd3fycjv.mp4
let media_path = MediaFileBucketPath::from_object_hash("t6cnyw4g3e8k7carkk2bvrt6nd3fycjv", Some("storyteller_"), Some(".mp4"));
let links = MediaLinksBuilder::from_media_path(DOMAIN, &media_path);
let links = MediaLinksBuilder::from_media_path_and_env(DOMAIN, ServerEnvironment::Production, &media_path);
assert_eq!(links.cdn_url.as_str(), "https://cdn.storyteller.ai/media/t/6/c/n/y/t6cnyw4g3e8k7carkk2bvrt6nd3fycjv/storyteller_t6cnyw4g3e8k7carkk2bvrt6nd3fycjv.mp4");
assert_eq!(links.maybe_thumbnail_template, None);
let video_previews = links.maybe_video_previews.expect("should have previews");
@@ -4,12 +4,12 @@ use actix_web::HttpRequest;
use anyhow::anyhow;
use log::warn;
use crate::error::internal_error::InternalError;
use crate::http_server::cookies::anonymous_visitor_tracking::avt_cookie_payload::AvtCookiePayload;
use cookies::jwt_signer::JwtSigner;
use errors::AnyhowResult;
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
use crate::http_server::cookies::anonymous_visitor_tracking::avt_cookie_payload::AvtCookiePayload;
const VISITOR_COOKIE_NAME : &str = "visitor";
/// Handle "anonymous visitor tracking" cookies.
@@ -70,7 +70,7 @@ impl AvtCookieManager {
Ok(payload)
}
pub fn decode_cookie_payload_from_request(&self, request: &HttpRequest) -> AnyhowResult<Option<AvtCookiePayload>> {
pub fn decode_cookie_payload_from_request(&self, request: &HttpRequest) -> Result<Option<AvtCookiePayload>, InternalError> {
let cookie = match request.cookie(VISITOR_COOKIE_NAME) {
None => return Ok(None),
Some(cookie) => cookie,
@@ -79,7 +79,7 @@ impl AvtCookieManager {
match self.decode_cookie_payload(&cookie) {
Err(e) => {
warn!("Visitor cookie decode error: {:?}", e);
Err(anyhow!("Could not decode visitor cookie: {:?}", e))
Err(InternalError::VisitorCookieError(e.to_string()))
},
Ok(payload) => Ok(Some(payload)),
}
@@ -1,9 +1,7 @@
use std::collections::BTreeMap;
use std::str::FromStr;
use anyhow::anyhow;
use errors::AnyhowResult;
use crate::error::internal_error::InternalError;
use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
/**
@@ -38,18 +36,19 @@ impl AvtCookiePayload {
}
}
pub fn from_map(map: BTreeMap<String, String>) -> AnyhowResult<Self> {
pub fn from_map(map: BTreeMap<String, String>) -> Result<Self, InternalError> {
let avt_token = map.get("avt_token")
.ok_or(anyhow!("avt_token not present"))?
.to_string();
.ok_or_else(|| InternalError::VisitorCookieMissingField("avt_token"))?;
let cookie_version = map
.get("cookie_version")
.ok_or(anyhow!("avt_token not present"))?;
.ok_or_else(|| InternalError::VisitorCookieMissingField("cookie_version"))?;
let avt_token = AnonymousVisitorTrackingToken::new_from_str(&avt_token);
let avt_token = AnonymousVisitorTrackingToken::new_from_str(avt_token);
let cookie_version = u32::from_str(cookie_version)?;
let cookie_version = u32::from_str(cookie_version)
.map_err(|e| InternalError::VisitorCookieError(
format!("invalid integer for cookie_version: {:?}, version: {}", e, cookie_version)))?;
Ok(AvtCookiePayload {
avt_token,
@@ -5,6 +5,8 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use crate::util::lookup::fetch_all_required_media_files::fetch_all_required_media_files;
use crate::util::traits::into_media_links_trait::IntoMediaLinks;
use actix_web::web::Json;
use actix_web::{web, HttpRequest};
use artcraft_api_defs::generate::video::generate_kling_1_6_pro_image_to_video::{GenerateKling16ProAspectRatio, GenerateKling16ProImageToVideoRequest};
@@ -24,6 +26,7 @@ use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::FalCategory;
use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job_for_fal_queue::InsertGenericInferenceForFalArgs;
use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::batch_get_media_files_by_tokens_with_connection;
use mysql_queries::queries::media_files::get::get_media_file::get_media_file_with_connection;
use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs};
use sqlx::Acquire;
@@ -73,14 +76,35 @@ pub async fn generate_kling_1_6_pro_video_handler(
// }
//};
let media_file_token = match &request.media_file_token {
let start_frame_media_file_token = match &request.media_file_token {
Some(token) => token,
None => {
warn!("No media file token provided");
return Err(CommonWebError::BadInputWithSimpleMessage("No media file token provided".to_string()));
}
};
let mut tokens = vec![
start_frame_media_file_token.clone(),
];
let maybe_end_frame_image_media_token = request.end_frame_image_media_token.as_ref();
if let Some(end_frame_token) = maybe_end_frame_image_media_token {
tokens.push(end_frame_token.clone());
}
let media_files = fetch_all_required_media_files(
&mut mysql_connection,
&tokens,
).await?;
for media_file in media_files.iter() {
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
}
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
@@ -91,42 +115,30 @@ pub async fn generate_kling_1_6_pro_video_handler(
error!("Error inserting idempotency token: {:?}", err);
CommonWebError::BadInputWithSimpleMessage("invalid idempotency token".to_string())
})?;
const IS_MOD : bool = false;
let media_file_lookup_result = get_media_file_with_connection(
media_file_token,
IS_MOD,
&mut mysql_connection,
).await;
let media_file = match media_file_lookup_result {
Ok(Some(media_file)) => media_file,
Ok(None) => {
warn!("MediaFile not found: {:?}", media_file_token);
return Err(CommonWebError::NotFound);
},
Err(err) => {
warn!("Error looking up media_file: {:?}", err);
return Err(CommonWebError::ServerError);
}
let media_domain = get_media_domain(&http_request);
let start_frame_url = media_files.iter()
.find(|file| &file.token == start_frame_media_file_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("Start frame media file not found after fetch");
CommonWebError::NotFound
})?;
let maybe_end_frame_url = match maybe_end_frame_image_media_token {
None => None,
Some(end_frame_token) => Some(media_files.iter()
.find(|file| &file.token == end_frame_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("End frame media file not found after fetch");
CommonWebError::NotFound
})?),
};
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
let media_domain = get_media_domain(&http_request);
let bucket_path = MediaFileBucketPath::from_object_hash(
&media_file.public_bucket_directory_hash,
media_file.maybe_public_bucket_prefix.as_deref(),
media_file.maybe_public_bucket_extension.as_deref());
let media_links = MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&bucket_path);
info!("Fal webhook URL: {}", server_state.fal.webhook_url);
let prompt = request.prompt
@@ -146,9 +158,10 @@ pub async fn generate_kling_1_6_pro_video_handler(
Some(GenerateKling16ProDuration::TenSeconds) => Kling16Duration::TenSeconds,
None => Kling16Duration::FiveSeconds,
};
let args = Kling16ProArgs {
image_url: media_links.cdn_url,
image_url: start_frame_url,
end_frame_image_url: maybe_end_frame_url,
webhook_url: &server_state.fal.webhook_url,
duration,
prompt,
@@ -159,7 +172,7 @@ pub async fn generate_kling_1_6_pro_video_handler(
let fal_result = enqueue_kling_16_pro_image_to_video_webhook(args)
.await
.map_err(|err| {
warn!("Error calling remove_background_rembg_webhook: {:?}", err);
warn!("Error calling enqueue_kling_16_pro_image_to_video_webhook: {:?}", err);
CommonWebError::ServerError
})?;
@@ -206,6 +219,8 @@ pub async fn generate_kling_1_6_pro_video_handler(
}
};
// TODO(bt,2025-08-26): Save reference images.
let db_result = insert_generic_inference_job_for_fal_queue(InsertGenericInferenceForFalArgs {
uuid_idempotency_token: &request.uuid_idempotency_token,
maybe_external_third_party_id: &external_job_id,
@@ -28,6 +28,8 @@ use mysql_queries::queries::media_files::get::get_media_file::{get_media_file, g
use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs};
use sqlx::Acquire;
use utoipa::ToSchema;
use crate::util::lookup::fetch_all_required_media_files::fetch_all_required_media_files;
use crate::util::traits::into_media_links_trait::IntoMediaLinks;
/// Kling 2.1 Pro Image to Video
#[utoipa::path(
@@ -73,14 +75,35 @@ pub async fn generate_kling_2_1_pro_video_handler(
// }
//};
let media_file_token = match &request.media_file_token {
let start_frame_media_file_token = match &request.media_file_token {
Some(token) => token,
None => {
warn!("No media file token provided");
return Err(CommonWebError::BadInputWithSimpleMessage("No media file token provided".to_string()));
}
};
let mut tokens = vec![
start_frame_media_file_token.clone(),
];
let maybe_end_frame_image_media_token = request.end_frame_image_media_token.as_ref();
if let Some(end_frame_token) = maybe_end_frame_image_media_token {
tokens.push(end_frame_token.clone());
}
let media_files = fetch_all_required_media_files(
&mut mysql_connection,
&tokens,
).await?;
for media_file in media_files.iter() {
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
}
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
@@ -91,42 +114,30 @@ pub async fn generate_kling_2_1_pro_video_handler(
error!("Error inserting idempotency token: {:?}", err);
CommonWebError::BadInputWithSimpleMessage("invalid idempotency token".to_string())
})?;
const IS_MOD : bool = false;
let media_file_lookup_result = get_media_file_with_connection(
media_file_token,
IS_MOD,
&mut mysql_connection,
).await;
let media_file = match media_file_lookup_result {
Ok(Some(media_file)) => media_file,
Ok(None) => {
warn!("MediaFile not found: {:?}", media_file_token);
return Err(CommonWebError::NotFound);
},
Err(err) => {
warn!("Error looking up media_file: {:?}", err);
return Err(CommonWebError::ServerError);
}
let media_domain = get_media_domain(&http_request);
let start_frame_url = media_files.iter()
.find(|file| &file.token == start_frame_media_file_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("Start frame media file not found after fetch");
CommonWebError::NotFound
})?;
let maybe_end_frame_url = match maybe_end_frame_image_media_token {
None => None,
Some(end_frame_token) => Some(media_files.iter()
.find(|file| &file.token == end_frame_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("End frame media file not found after fetch");
CommonWebError::NotFound
})?),
};
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
let media_domain = get_media_domain(&http_request);
let bucket_path = MediaFileBucketPath::from_object_hash(
&media_file.public_bucket_directory_hash,
media_file.maybe_public_bucket_prefix.as_deref(),
media_file.maybe_public_bucket_extension.as_deref());
let media_links = MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&bucket_path);
info!("Fal webhook URL: {}", server_state.fal.webhook_url);
let prompt = request.prompt
@@ -148,7 +159,8 @@ pub async fn generate_kling_2_1_pro_video_handler(
};
let args = Kling21ProArgs {
image_url: media_links.cdn_url,
image_url: start_frame_url,
end_frame_image_url: maybe_end_frame_url,
webhook_url: &server_state.fal.webhook_url,
duration,
prompt,
@@ -5,6 +5,8 @@ use crate::http_server::common_responses::media::media_links_builder::MediaLinks
use crate::http_server::endpoints::media_files::helpers::get_media_domain::get_media_domain;
use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format;
use crate::state::server_state::ServerState;
use crate::util::lookup::fetch_all_required_media_files::fetch_all_required_media_files;
use crate::util::traits::into_media_links_trait::IntoMediaLinks;
use actix_web::web::Json;
use actix_web::{web, HttpRequest};
use artcraft_api_defs::generate::video::generate_seedance_1_0_lite_image_to_video::GenerateSeedance10LiteDuration;
@@ -74,14 +76,35 @@ pub async fn generate_seedance_1_0_lite_image_to_video_handler(
// }
//};
let media_file_token = match &request.media_file_token {
let start_frame_media_file_token = match &request.media_file_token {
Some(token) => token,
None => {
warn!("No media file token provided");
return Err(CommonWebError::BadInputWithSimpleMessage("No media file token provided".to_string()));
}
};
let mut tokens = vec![
start_frame_media_file_token.clone(),
];
let maybe_end_frame_image_media_token = request.end_frame_image_media_token.as_ref();
if let Some(end_frame_token) = maybe_end_frame_image_media_token {
tokens.push(end_frame_token.clone());
}
let media_files = fetch_all_required_media_files(
&mut mysql_connection,
&tokens,
).await?;
for media_file in media_files.iter() {
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
}
if let Err(reason) = validate_idempotency_token_format(&request.uuid_idempotency_token) {
return Err(CommonWebError::BadInputWithSimpleMessage(reason));
}
@@ -92,42 +115,30 @@ pub async fn generate_seedance_1_0_lite_image_to_video_handler(
error!("Error inserting idempotency token: {:?}", err);
CommonWebError::BadInputWithSimpleMessage("invalid idempotency token".to_string())
})?;
const IS_MOD : bool = false;
let media_file_lookup_result = get_media_file_with_connection(
media_file_token,
IS_MOD,
&mut mysql_connection,
).await;
let media_file = match media_file_lookup_result {
Ok(Some(media_file)) => media_file,
Ok(None) => {
warn!("MediaFile not found: {:?}", media_file_token);
return Err(CommonWebError::NotFound);
},
Err(err) => {
warn!("Error looking up media_file: {:?}", err);
return Err(CommonWebError::ServerError);
}
let media_domain = get_media_domain(&http_request);
let start_frame_url = media_files.iter()
.find(|file| &file.token == start_frame_media_file_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("Start frame media file not found after fetch");
CommonWebError::NotFound
})?;
let maybe_end_frame_url = match maybe_end_frame_image_media_token {
None => None,
Some(end_frame_token) => Some(media_files.iter()
.find(|file| &file.token == end_frame_token)
.map(|file| file.to_media_links(media_domain, server_state.server_environment))
.map(|file| file.cdn_url)
.ok_or_else(|| {
warn!("End frame media file not found after fetch");
CommonWebError::NotFound
})?),
};
if !media_file.media_type.is_jpg_or_png_or_legacy_image() {
return Err(CommonWebError::BadInputWithSimpleMessage("Media file must be a JPG or PNG image".to_string()));
}
let media_domain = get_media_domain(&http_request);
let bucket_path = MediaFileBucketPath::from_object_hash(
&media_file.public_bucket_directory_hash,
media_file.maybe_public_bucket_prefix.as_deref(),
media_file.maybe_public_bucket_extension.as_deref());
let media_links = MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&bucket_path);
info!("Fal webhook URL: {}", server_state.fal.webhook_url);
let prompt = request.prompt
@@ -148,7 +159,8 @@ pub async fn generate_seedance_1_0_lite_image_to_video_handler(
};
let args = Seedance1LiteArgs {
image_url: media_links.cdn_url,
image_url: start_frame_url,
end_frame_image_url: maybe_end_frame_url,
webhook_url: &server_state.fal.webhook_url,
api_key: &server_state.fal.api_key,
duration,
@@ -31,6 +31,7 @@ use mysql_queries::queries::generic_inference::web::batch_get_inference_job_stat
use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus;
use r2d2_redis::redis::Commands;
use redis_common::redis_keys::RedisKeys;
use server_environment::ServerEnvironment;
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use utoipa::{IntoParams, ToSchema};
@@ -203,7 +204,11 @@ pub async fn batch_get_inference_job_status_handler(
if tokens.is_empty() {
// NB: MediaDomain doesn't matter since it's an empty list.
return Ok(records_to_response(Vec::new(), MediaDomain::Storyteller));
return Ok(records_to_response(
Vec::new(),
server_state.server_environment,
MediaDomain::Storyteller
));
}
// NB: Since this is publicly exposed, we don't query sensitive data.
@@ -249,18 +254,19 @@ pub async fn batch_get_inference_job_status_handler(
let media_domain = get_media_domain(&http_request);
Ok(records_to_response(records, media_domain))
Ok(records_to_response(records, server_state.server_environment, media_domain))
}
fn records_to_response(
records: Vec<GenericInferenceJobStatus>,
server_environment: ServerEnvironment,
media_domain: MediaDomain,
) -> Json<BatchGetInferenceJobStatusSuccessResponse> {
Json(BatchGetInferenceJobStatusSuccessResponse {
success: true,
job_states: records.into_iter()
.map(|record| {
db_record_to_response_payload(record, None, media_domain)
db_record_to_response_payload(record, None, server_environment, media_domain)
})
.collect::<Vec<_>>(),
})
@@ -269,6 +275,7 @@ fn records_to_response(
fn db_record_to_response_payload(
record: GenericInferenceJobStatus,
maybe_extra_status_description: Option<String>,
server_environment: ServerEnvironment,
media_domain: MediaDomain,
) -> BatchInferenceJobStatusResponsePayload {
let inference_category = record.request_details.inference_category;
@@ -374,7 +381,11 @@ fn db_record_to_response_payload(
BatchResultDetailsResponse {
entity_type: result_details.entity_type,
entity_token: result_details.entity_token,
media_links: MediaLinksBuilder::from_rooted_path(media_domain, &public_bucket_media_path),
media_links: MediaLinksBuilder::from_rooted_path_and_env(
media_domain,
server_environment,
&public_bucket_media_path
),
maybe_public_bucket_media_path: Some(public_bucket_media_path),
maybe_successfully_completed_at: result_details.maybe_successfully_completed_at,
}
@@ -29,6 +29,7 @@ use mysql_queries::queries::generic_inference::web::get_inference_job_status::ge
use mysql_queries::queries::generic_inference::web::job_status::GenericInferenceJobStatus;
use r2d2_redis::redis::{Commands, RedisResult};
use redis_common::redis_keys::RedisKeys;
use server_environment::ServerEnvironment;
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use utoipa::ToSchema;
@@ -251,6 +252,7 @@ pub async fn get_inference_job_status_handler(
let record_for_response = record_to_payload(
record,
maybe_extra_status_description,
server_state.server_environment,
media_domain,
);
@@ -263,6 +265,7 @@ pub async fn get_inference_job_status_handler(
fn record_to_payload(
record: GenericInferenceJobStatus,
maybe_extra_status_description: Option<String>,
server_environment: ServerEnvironment,
media_domain: MediaDomain,
) -> InferenceJobStatusResponsePayload {
let inference_category = record.request_details.inference_category;
@@ -368,7 +371,11 @@ fn record_to_payload(
ResultDetailsResponse {
entity_type: result_details.entity_type,
entity_token: result_details.entity_token,
media_links: MediaLinksBuilder::from_rooted_path(media_domain, &public_bucket_media_path),
media_links: MediaLinksBuilder::from_rooted_path_and_env(
media_domain,
server_environment,
&public_bucket_media_path
),
maybe_public_bucket_media_path: Some(public_bucket_media_path),
maybe_successfully_completed_at: result_details.maybe_successfully_completed_at,
}
@@ -384,6 +391,7 @@ mod tests {
use crate::http_server::endpoints::inference_job::get::get_inference_job_status_handler::record_to_payload;
use enums::by_table::generic_inference_jobs::inference_category::InferenceCategory;
use mysql_queries::queries::generic_inference::web::job_status::{GenericInferenceJobStatus, RequestDetails, ResultDetails};
use server_environment::ServerEnvironment;
use url::Url;
#[test]
@@ -407,7 +415,7 @@ mod tests {
};
let payload =
record_to_payload(status, None, MediaDomain::Storyteller);
record_to_payload(status, None, ServerEnvironment::Production, MediaDomain::Storyteller);
assert!(payload.maybe_result.is_some());
@@ -446,7 +454,7 @@ mod tests {
};
let payload =
record_to_payload(status, None, MediaDomain::FakeYou);
record_to_payload(status, None, ServerEnvironment::Production, MediaDomain::FakeYou);
assert!(payload.maybe_result.is_some());
@@ -31,6 +31,7 @@ use mysql_queries::queries::generic_inference::web::list_session_jobs::{list_ses
use primitives::numerics::i64_to_u64_zero_clamped::i64_to_u64_zero_clamped;
use r2d2_redis::redis::Commands;
use redis_common::redis_keys::RedisKeys;
use server_environment::ServerEnvironment;
use tokens::tokens::generic_inference_jobs::InferenceJobToken;
use tokens::tokens::media_files::MediaFileToken;
use utoipa::{IntoParams, ToSchema};
@@ -197,16 +198,17 @@ pub async fn list_session_jobs_handler(
let media_domain = get_media_domain(&http_request);
records_to_response(records, media_domain)
records_to_response(records, server_state.server_environment, media_domain)
}
fn records_to_response(
records: Vec<GenericInferenceJobStatus>,
server_environment: ServerEnvironment,
media_domain: MediaDomain,
) -> Result<Json<ListSessionJobsSuccessResponse>, ListSessionJobsError> {
let mut records = records.into_iter()
.map(|record| {
db_record_to_response_payload(record, None, media_domain)
db_record_to_response_payload(record, None, server_environment, media_domain)
})
.collect::<Vec<_>>();
@@ -240,6 +242,7 @@ fn records_to_response(
fn db_record_to_response_payload(
record: GenericInferenceJobStatus,
maybe_extra_status_description: Option<String>,
server_environment: ServerEnvironment,
media_domain: MediaDomain,
) -> ListSessionJobsItem {
let inference_category = record.request_details.inference_category;
@@ -359,7 +362,11 @@ fn db_record_to_response_payload(
entity_type: result_details.entity_type,
entity_token: result_details.entity_token,
maybe_batch_token: result_details.maybe_batch_token,
media_links: MediaLinksBuilder::from_rooted_path(media_domain, &public_bucket_media_path),
media_links: MediaLinksBuilder::from_rooted_path_and_env(
media_domain,
server_environment,
&public_bucket_media_path
),
maybe_public_bucket_media_path: Some(public_bucket_media_path),
maybe_successfully_completed_at: result_details.maybe_successfully_completed_at,
}
@@ -333,7 +333,11 @@ pub async fn batch_get_media_files_handler(
maybe_media_subtype: result.maybe_media_subtype,
maybe_engine_extension,
maybe_batch_token: result.maybe_batch_token,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -112,7 +112,11 @@ pub async fn list_batch_generated_redux_media_files_handler(
media_class: result.media_class,
media_type: result.media_type,
maybe_batch_token: result.maybe_batch_token,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
cover_image: MediaFileCoverImageDetailsBuilder::from_optional_db_fields(
&result.token,
media_domain,
@@ -288,7 +288,11 @@ pub async fn list_featured_media_files_handler(
media_type: m.media_type,
maybe_engine_category: m.maybe_engine_category,
maybe_animation_type: m.maybe_animation_type,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -277,7 +277,11 @@ pub async fn list_media_files_by_batch_token_handler(
maybe_origin_model_type: record.maybe_origin_model_type
.map(|t| PublicMediaFileModelType::from_enum(t)),
maybe_origin_model_token: record.maybe_origin_model_token,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -345,7 +345,11 @@ pub async fn list_media_files_handler(
maybe_origin_model_type: record.maybe_origin_model_type
.map(|m| PublicMediaFileModelType::from_enum(m)),
maybe_origin_model_token: record.maybe_origin_model_token,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -215,7 +215,11 @@ pub async fn list_pinned_media_files_handler(
media_type: m.media_type,
maybe_engine_category: m.maybe_engine_category,
maybe_animation_type: m.maybe_animation_type,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -227,7 +227,11 @@ pub async fn search_featured_media_files_handler(
media_type: result.media_type,
maybe_engine_category: result.maybe_engine_category,
maybe_animation_type: result.maybe_animation_type,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -254,7 +254,11 @@ pub async fn search_session_media_files_handler(
media_type: result.media_type,
maybe_engine_category: result.maybe_engine_category,
maybe_animation_type: result.maybe_animation_type,
media_links: MediaLinksBuilder::from_media_path(media_domain, &public_bucket_path),
media_links: MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
&public_bucket_path
),
public_bucket_path: public_bucket_path
.get_full_object_path_str()
.to_string(),
@@ -230,7 +230,10 @@ pub async fn list_user_bookmarks_for_user_handler(
});
let maybe_media_file_media_links = maybe_media_file_bucket_path.as_ref()
.map(|bucket_path| MediaLinksBuilder::from_media_path(media_domain, bucket_path));
.map(|bucket_path| MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_state.server_environment,
bucket_path));
let mut maybe_media_file_cover = None;
let mut maybe_model_weight_cover = None;
@@ -92,6 +92,7 @@ use crate::util::troll_user_bans::troll_user_ban_list::TrollUserBanList;
pub mod billing;
pub mod configs;
pub mod error;
pub mod http_server;
pub mod state;
pub mod threads;
@@ -0,0 +1,37 @@
use crate::http_server::common_responses::common_web_error::CommonWebError;
use log::{error, warn};
use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::{batch_get_media_files_by_tokens_with_connection, MediaFilesByTokensRecord};
use sqlx::pool::PoolConnection;
use sqlx::MySql;
use tokens::tokens::media_files::MediaFileToken;
const CAN_SEE_DELETED: bool = false;
/// Look up all the media tokens. If any are missing, return an error.
pub async fn fetch_all_required_media_files(
mysql_connection: &mut PoolConnection<MySql>,
media_file_tokens: &[MediaFileToken],
) -> Result<Vec<MediaFilesByTokensRecord>, CommonWebError> {
let result = batch_get_media_files_by_tokens_with_connection(
mysql_connection,
&media_file_tokens,
CAN_SEE_DELETED,
).await;
let media_files = match result {
Ok(files) => files,
Err(err) => {
error!("Error getting media files by tokens: {:?}", err);
return Err(CommonWebError::ServerError);
}
};
if media_files.len() != media_file_tokens.len() {
warn!("Wrong number of media files returned for tokens: {} found for {} tokens", media_files.len(), media_file_tokens.len());
return Err(CommonWebError::BadInputWithSimpleMessage(
format!("Not all media files could be found. Media files found: {}, tokens provided: {}",
media_files.len(), media_file_tokens.len())));
}
Ok(media_files)
}
@@ -0,0 +1 @@
pub mod fetch_all_required_media_files;
@@ -11,7 +11,9 @@ pub mod enroll_in_studio;
pub mod generate_random_username;
pub mod http_download_url_to_bytes;
pub mod http_download_url_to_tempfile;
pub mod lookup;
pub mod placeholder_images;
pub mod read_toml_file_to_struct;
pub mod title_to_url_slug;
pub mod traits;
pub mod troll_user_bans;
@@ -0,0 +1,28 @@
use crate::http_server::common_responses::media::media_domain::MediaDomain;
use crate::http_server::common_responses::media::media_links_builder::MediaLinksBuilder;
use artcraft_api_defs::common::responses::media_links::MediaLinks;
use bucket_paths::legacy::typified_paths::public::media_files::bucket_file_path::MediaFileBucketPath;
use mysql_queries::queries::media_files::get::batch_get_media_files_by_tokens::MediaFilesByTokensRecord;
use server_environment::ServerEnvironment;
pub trait IntoMediaLinks {
fn to_media_links(
&self,
media_domain: MediaDomain,
server_environment: ServerEnvironment,
) -> MediaLinks;
}
impl IntoMediaLinks for MediaFilesByTokensRecord {
fn to_media_links(&self, media_domain: MediaDomain, server_environment: ServerEnvironment) -> MediaLinks {
let bucket_path = MediaFileBucketPath::from_object_hash(
&self.public_bucket_directory_hash,
self.maybe_public_bucket_prefix.as_deref(),
self.maybe_public_bucket_extension.as_deref());
MediaLinksBuilder::from_media_path_and_env(
media_domain,
server_environment,
&bucket_path
)
}
}
@@ -0,0 +1 @@
pub mod into_media_links_trait;
@@ -5,6 +5,29 @@ use serde::{Deserialize, Serialize};
#[allow(unused_imports)]
use std::collections::HashMap;
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct ImageToVideoRequest {
pub image_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub end_image_url: Option<String>,
/// The resolution of the generated video frame
#[serde(skip_serializing_if = "Option::is_none")]
pub resolution: Option<String>,
/// The duration of the generated video in seconds
#[serde(skip_serializing_if = "Option::is_none")]
pub duration: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub camera_fixed: Option<bool>,
pub prompt: String,
pub seed: i32,
}
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct File {
/// The mime type of the file.
@@ -39,21 +62,6 @@ pub struct I2VOutput {
pub video: File,
}
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct ImageToVideoRequest {
/// The resolution of the generated video frame
#[serde(skip_serializing_if = "Option::is_none")]
pub resolution: Option<String>,
/// The duration of the generated video in seconds
#[serde(skip_serializing_if = "Option::is_none")]
pub duration: Option<String>,
pub image_url: String,
#[serde(skip_serializing_if = "Option::is_none")]
pub camera_fixed: Option<bool>,
pub prompt: String,
pub seed: i32,
}
#[derive(Debug, Serialize, Deserialize, Default)]
pub struct ValidationError {
pub loc: Vec<serde_json::Value>,
@@ -13,3 +13,17 @@ export const normalizeCanvas = (canvas: HTMLCanvasElement, width: number, height
ctx.drawImage(canvas, 0, 0, width, height);
return newCanvas;
}
export const drawAlphaMask = (canvas: OffscreenCanvas, width: number, height: number) => {
const ctx = canvas.getContext("2d");
if (!ctx) {
throw new Error("Failed to get canvas context");
}
// Fill the entire screen with black and composite to keep non-transparent pixels
ctx.save();
ctx.globalCompositeOperation = "source-in";
ctx.fillStyle = "black";
ctx.fillRect(0, 0, width, height);
ctx.restore();
}
@@ -36,7 +36,7 @@ import { useSignals } from "@preact/signals-react/runtime";
import { setLogoutStates } from "~/signals/authentication/utilities";
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import { useTabStore } from "~/pages/Stores/TabState";
import { TabId, useTabStore } from "~/pages/Stores/TabState";
import {
is3DEditorInitialized,
is3DSceneLoaded,
@@ -59,7 +59,23 @@ interface Props {
const SWITCHER_THROTTLE_TIME = 500; // milliseconds
const appMenuTabs: MenuIconItem[] = [
// NB: See `TabState` for the default tab.
const appMenuTabs: MenuIconItem<TabId>[] = [
{
id: "IMAGE",
label: "Text to Image",
icon: <FontAwesomeIcon icon={faImage} />,
},
{
id: "VIDEO",
label: "Image to Video",
icon: <FontAwesomeIcon icon={faFilm} />,
},
{
id: "EDIT",
label: "Edit Image",
icon: <FontAwesomeIcon icon={faPenNib} />,
},
{
id: "2D",
label: "2D Canvas",
@@ -76,21 +92,6 @@ const appMenuTabs: MenuIconItem[] = [
description: "Precision control. Great for AI film.",
large: true,
},
{
id: "IMAGE",
label: "Prompt to Image",
icon: <FontAwesomeIcon icon={faImage} />,
},
{
id: "VIDEO",
label: "Image to Video",
icon: <FontAwesomeIcon icon={faFilm} />,
},
{
id: "EDIT",
label: "Edit Image",
icon: <FontAwesomeIcon icon={faPenNib} />,
},
];
export const topNavMediaId = signal<string>("");
@@ -153,6 +154,9 @@ export const TopBar = ({ pageName, loginSignUpPressed }: Props) => {
url,
mediaToken: mediaId || "",
};
// Add it to state history
useEditStore.getState().addHistoryImageBundle({ images: [baseImage] });
useEditStore.getState().setBaseImageInfo(baseImage);
// Close gallery modal and lightbox if open
@@ -11,14 +11,14 @@ import Konva from "konva";
import { setCanvasRenderBitmap } from "../../signals/canvasRenderBitmap";
import { captureStageImageBitmap } from "./hooks/useUpdateSnapshot";
import { ContextMenuContainer } from "./components/ui/ContextMenu";
import { ModelInfo } from "@storyteller/model-list";
import { ImageModel } from "@storyteller/model-list";
import {
CANVAS_2D_PAGE_MODEL_LIST,
ClassyModelSelector,
ModelPage,
ModelSelector,
useModelSelectorStore,
} from "@storyteller/ui-model-selector";
import { useCanvasBgRemovedEvent } from "@storyteller/tauri-api";
import { getSelectedImageModel } from "@storyteller/ui-model-selector";
const PAGE_ID : ModelPage = ModelPage.Canvas2D;
@@ -50,20 +50,12 @@ export const DecodeBase64ToImage = async (
const PageDraw = () => {
const canvasWidth = useRef<number>(1024);
const canvasHeight = useRef<number>(1024);
const { selectedModels } = useModelSelectorStore();
const [isSelecting, setIsSelecting] = useState<boolean>(false);
const stageRef = useRef<Konva.Stage>({} as Konva.Stage);
const transformerRefs = useRef<{ [key: string]: Konva.Transformer }>({});
const store = useSceneStore();
const selectedModel =
selectedModels[PAGE_ID] ||
CANVAS_2D_PAGE_MODEL_LIST[0]?.label;
const selectedModelInfo: ModelInfo | undefined =
CANVAS_2D_PAGE_MODEL_LIST.find(
(m) => m.label === selectedModel,
)?.modelInfo;
const selectedImageModel : ImageModel | undefined = getSelectedImageModel(PAGE_ID);
useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
useUndoRedoHotkeys({ undo: store.undo, redo: store.redo });
@@ -305,7 +297,8 @@ const PageDraw = () => {
}}
onEnqueuePressed={onEnqueuedPressed}
onFitPressed={onFitPressed}
selectedModelInfo={selectedModelInfo}
//selectedModelInfo={selectedModelInfo}
selectedImageModel={selectedImageModel}
/>
</div>
<SideToolbar
@@ -446,7 +439,7 @@ const PageDraw = () => {
</ContextMenuContainer>
</div>
<div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
<ModelSelector
<ClassyModelSelector
items={CANVAS_2D_PAGE_MODEL_LIST}
page={PAGE_ID}
panelTitle="Select Model"
@@ -9,12 +9,14 @@ import { JobProvider, useJobContext } from "../JobContext";
// Set this value on when enqueue is pressed nasty global variable.
import { getCanvasRenderBitmap } from "../../../signals/canvasRenderBitmap";
const PromptEditor: React.FC<PromptEditorProps> = ({
onImageStyleChange,
onEnqueuePressed,
onAspectRatioChange,
onFitPressed,
selectedModelInfo,
//selectedModelInfo,
selectedImageModel,
}) => {
const [images, setImages] = useState<ImageStyle[]>([]);
@@ -44,7 +46,8 @@ const PromptEditor: React.FC<PromptEditorProps> = ({
onEnqueuePressed={onEnqueuePressed}
onAspectRatioChange={onAspectRatioChange}
onFitPressed={onFitPressed}
selectedModelInfo={selectedModelInfo}
//selectedModelInfo={selectedModelInfo}
selectedImageModel={selectedImageModel}
/>
</JobProvider>
@@ -1,4 +1,5 @@
import { ModelInfo } from "libs/model-list/src/lib/ModelInfo";
import { ImageModel } from "@storyteller/model-list";
import { ModelInfo } from "@storyteller/model-list";
export interface PromptEditorProps {
initialPrompt?: string;
@@ -11,6 +12,7 @@ export interface PromptEditorProps {
onEnqueuePressed?: () => void;
onFitPressed?: () => void;
selectedModelInfo?: ModelInfo;
selectedImageModel?: ImageModel;
}
export type AspectRatio = "1:1" | "3:2" | "2:3";
@@ -1,11 +1,19 @@
import { Button } from "@storyteller/ui-button";
import { useImageEditCompleteEvent } from "@storyteller/tauri-events";
import { faClockRotateLeft, faTrash, faTrashXmark, faXmark } from "@fortawesome/pro-solid-svg-icons";
import { useRef, useState } from "react";
import {
faClockRotateLeft,
faTrash,
faTrashXmark,
faXmark,
} from "@fortawesome/pro-solid-svg-icons";
import { Fragment, useRef, useState } from "react";
import { twMerge } from "tailwind-merge";
import { BaseSelectorImage } from "./BaseImageSelector";
import { Tooltip } from "@storyteller/ui-tooltip";
import { isActionReminderOpen, showActionReminder } from "@storyteller/ui-action-reminder-modal";
import {
isActionReminderOpen,
showActionReminder,
} from "@storyteller/ui-action-reminder-modal";
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
export interface ImageBundle {
@@ -15,21 +23,18 @@ export interface ImageBundle {
interface HistoryStackProps {
onClear: () => void;
onImageSelect?: (image: BaseSelectorImage) => void;
startingBundles: ImageBundle[];
imageBundles: ImageBundle[];
onNewImageBundle?: (newBundle: ImageBundle) => void;
selectedImageToken?: string;
}
export const HistoryStack = ({
onClear,
onImageSelect = () => { },
startingBundles = [],
imageBundles,
onNewImageBundle = () => { },
selectedImageToken
}: HistoryStackProps) => {
const [imageBundles, setImageBundles] =
useState<ImageBundle[]>(startingBundles);
const [selectedImageToken, setSelectedImageToken] = useState<string | null>(
startingBundles.length > 0 && startingBundles[0].images.length > 0
? startingBundles[0].images[0].mediaToken
: null,
);
useImageEditCompleteEvent(async (event) => {
const newBundle: ImageBundle = {
images: event.edited_images.map(
@@ -41,12 +46,7 @@ export const HistoryStack = ({
),
};
setImageBundles([...imageBundles, newBundle]);
const newlySelected = newBundle.images[0];
if (newlySelected) {
setSelectedImageToken(newlySelected.mediaToken);
onImageSelect(newlySelected);
}
onNewImageBundle(newBundle);
});
// This is used to force image reloads in different sessions
@@ -54,27 +54,23 @@ export const HistoryStack = ({
const sessionRandBuster = useRef(Math.random());
const handleClear = () => {
setImageBundles(startingBundles);
setSelectedImageToken(null);
onClear();
};
return (
<div className="max-h-1/2 h-auto w-20 overflow-y-auto rounded-lg">
<div className="flex glass rounded-lg flex-col-reverse items-center justify-center gap-2 p-1">
<div className="h-auto w-20 rounded-lg">
<div className={"glass max-h-[50vh] flex flex-col-reverse items-center justify-center gap-2 overflow-y-auto rounded-lg p-1.5"}>
{imageBundles.map((bundle, index) => (
<>
{index < imageBundles.length - 1 && <hr className="h-0.5 w-3/4 rounded-md border-none bg-white/10" />}
{bundle.images.map((image, imgIndex) => (
<Fragment key={index}>
{bundle.images.map((image) => (
<Button
key={imgIndex}
key={image.mediaToken}
className={twMerge(
"aspect-square relative h-full w-full border-2 bg-transparent p-0 m-1 hover:bg-transparent hover:opacity-80",
"relative aspect-square h-full w-full border-2 bg-transparent p-0 hover:bg-transparent hover:opacity-80",
selectedImageToken === image.mediaToken &&
"border-primary hover:opacity-100",
)}
onClick={() => {
setSelectedImageToken(image.mediaToken);
onImageSelect(image);
}}
>
@@ -87,35 +83,54 @@ export const HistoryStack = ({
/>
</Button>
))}
</>
{index < imageBundles.length - 1 && (
<hr className="min-h-0.5 h-0.5 w-3/4 rounded-md border-none bg-white/10" key={"hr" + index} />
)}
</Fragment>
))}
<FontAwesomeIcon icon={faClockRotateLeft} className="p-1 text-gray-400" />
<FontAwesomeIcon
icon={faClockRotateLeft}
className="p-1 text-gray-400"
/>
</div>
<div className="glass mt-4 rounded-full border-red/50 border-2 size-12 flex justify-center items-center mx-auto">
<button
className="flex h-full w-full aspect-square rounded-full items-center justify-center text-white transition-colors hover:bg-red/50"
onClick={() =>
showActionReminder({
reminderType: "default",
title: "Reset All",
primaryActionIcon: faTrashXmark,
primaryActionBtnClassName: "bg-red hover:bg-red/80",
message: (
<p className="text-sm text-white/70">
Are you sure you want to reset all? This will clear all your
work and cannot be undone.
</p>
),
primaryActionText: "Reset all",
onPrimaryAction: () => {
handleClear();
isActionReminderOpen.value = false;
},
})
}
>
<FontAwesomeIcon icon={faXmark} className="h-5 w-5 text-xl" />
</button>
<div className="mt-3 flex justify-center">
<div className="glass w-fit rounded-xl border-2 border-red/50 shadow-lg hover:border-red/80">
<div className="relative h-full">
<Tooltip
content="Reset All"
position="left"
closeOnClick={true}
className="ms-1 rounded-md bg-red px-3 py-1"
delay={100}
>
<button
className="flex h-10 w-10 items-center justify-center rounded-lg border-2 border-transparent text-white transition-colors hover:bg-red/50"
onClick={() =>
showActionReminder({
reminderType: "default",
title: "Reset All",
primaryActionIcon: faTrashXmark,
primaryActionBtnClassName: "bg-red hover:bg-red/80",
message: (
<p className="text-sm text-white/70">
Are you sure you want to reset all? This will clear all
your work and cannot be undone.
</p>
),
primaryActionText: "Reset all",
onPrimaryAction: () => {
handleClear();
isActionReminderOpen.value = false;
},
})
}
>
<FontAwesomeIcon icon={faXmark} className="h-5 w-5 text-xl" />
</button>
</Tooltip>
</div>
</div>
</div>
</div>
);
@@ -2,7 +2,6 @@ import { useState, useRef, useEffect, useCallback } from "react";
import Konva from "konva"; // just for types
import {
EnqueueImageInpaint,
EnqueueImageInpaintModel,
} from "@storyteller/tauri-api";
import { ContextMenuContainer } from "../PageDraw/components/ui/ContextMenu";
import { useCopyPasteHotkeys } from "../PageDraw/hooks/useCopyPasteHotkeys";
@@ -11,34 +10,24 @@ import { useUndoRedoHotkeys } from "../PageDraw/hooks/useUndoRedoHotkeys";
import PromptEditor from "./PromptEditor/PromptEditor";
import { ActiveEditTool, useEditStore } from "./stores/EditState";
import { EditPaintSurface } from "./EditPaintSurface";
import { normalizeCanvas } from "../../Helpers/CanvasHelpers";
import { drawAlphaMask, normalizeCanvas } from "../../Helpers/CanvasHelpers";
import { BaseImageSelector, BaseSelectorImage } from "./BaseImageSelector";
import DrawToolControlBar from "./DrawToolControlBar";
import {
ClassyModelSelector,
getSelectedImageModel,
IMAGE_EDITOR_PAGE_MODEL_LIST,
ModelPage,
ModelSelector,
useModelSelectorStore,
} from "@storyteller/ui-model-selector";
import { lookupModelByTauriId, ModelInfo } from "@storyteller/model-list";
import { ImageModel } from "@storyteller/model-list";
import { HistoryStack, ImageBundle } from "./HistoryStack";
import { ModelTag } from "libs/model-list/src/lib/ModelTag";
const PAGE_ID: ModelPage = ModelPage.ImageEditor;
const PageEdit = () => {
//useStateSceneLoader();
const { selectedModels } = useModelSelectorStore();
const selectedModelLabel =
selectedModels[PAGE_ID] || IMAGE_EDITOR_PAGE_MODEL_LIST[0]?.label;
const selectedModel = IMAGE_EDITOR_PAGE_MODEL_LIST.find(
(m) => m.label === selectedModelLabel,
);
const selectedModelInfo: ModelInfo | undefined =
selectedModel?.modelInfo;
const selectedImageModel: ImageModel | undefined = getSelectedImageModel(PAGE_ID);
// State for canvas dimensions
const canvasWidth = useRef<number>(1024);
@@ -57,6 +46,8 @@ const PageEdit = () => {
// Use the Zustand store
const store = useEditStore();
const addHistoryImageBundle = useEditStore((state) => state.addHistoryImageBundle);
const historyImageBundles = useEditStore((state) => state.historyImageBundles)
// Pass store actions directly as callbacks
useDeleteHotkeys({ onDelete: store.deleteSelectedItems });
@@ -165,18 +156,6 @@ const PageEdit = () => {
// eslint-disable-next-line react-hooks/exhaustive-deps
}, []);
// Auto-fit the canvas when base image is loaded
useEffect(() => {
if (!store.baseImageBitmap) {
console.log("No base image bitmap available, skipping auto-fit");
return;
}
// When baseImageBitmap changes from null to an actual image, call onFitPressed
// Use a slight delay to ensure the image dimensions are properly loaded
onFitPressed();
}, [onFitPressed, store.baseImageBitmap]);
// Create a function to use the left layer ref and download the bitmap from it
const getMaskArrayBuffer = async (): Promise<Uint8Array> => {
if (
@@ -208,6 +187,10 @@ const PageEdit = () => {
rect.height(),
);
// Convert colored canvas to alpha mask
// NOTE: This isn't needed because the tauri backend uses the alpha channel anyway
// drawAlphaMask(fittedCanvas, rect.width(), rect.height());
const blob = await fittedCanvas.convertToBlob({ type: "image/png" });
const arrayBuffer = await blob.arrayBuffer();
@@ -235,7 +218,7 @@ const PageEdit = () => {
try {
await EnqueueImageInpaint({
model: selectedModelInfo,
model: selectedImageModel,
image_media_token: editedImageToken,
mask_image_raw_bytes: arrayBuffer,
prompt: prompt,
@@ -249,19 +232,23 @@ const PageEdit = () => {
[
generationCount,
isEnqueuing,
selectedModelInfo,
selectedImageModel,
store.baseImageInfo?.mediaToken,
],
);
const modelConfig = lookupModelByTauriId(selectedModelInfo!.tauri_id);
const supportsMaskedInpainting = modelConfig?.tags?.includes(ModelTag.MaskedInpainting) ?? false;
//const modelConfig = lookupModelByTauriId(selectedImageModel!.tauriId);
//const supportsMaskedInpainting = modelConfig?.tags?.includes(ModelTag.MaskedInpainting) ?? false;
const supportsMaskedInpainting = selectedImageModel?.usesInpaintingMask ?? false;
useEffect(() => {
if (!supportsMaskedInpainting && (store.activeTool !== "select" || store.lineNodes.length > 0)) {
// TODO: Implement a new mode for unsupported masking and hide the nodes layer instead of clearing
store.setActiveTool("select");
store.clearLineNodes();
}
}, [store, supportsMaskedInpainting]);
if (!supportsMaskedInpainting && (store.activeTool !== "select" || store.lineNodes.length > 0)) {
// TODO: Implement a new mode for unsupported masking and hide the nodes layer instead of clearing
store.setActiveTool("select");
store.clearLineNodes();
}
/*
*
@@ -283,6 +270,7 @@ const PageEdit = () => {
>
<BaseImageSelector
onImageSelect={(image: BaseSelectorImage) => {
addHistoryImageBundle({ images: [image] });
store.setBaseImageInfo(image);
}}
showLoading={
@@ -315,11 +303,13 @@ const PageEdit = () => {
onClear={() => {
store.RESET();
}}
startingBundles={[{ images: [store.baseImageInfo] } as ImageBundle]}
imageBundles={historyImageBundles}
onImageSelect={(baseImage) => {
store.clearLineNodes();
store.setBaseImageInfo(baseImage);
}}
onNewImageBundle={addHistoryImageBundle}
selectedImageToken={store.baseImageInfo?.mediaToken}
/>
</div>
<div
@@ -327,7 +317,7 @@ const PageEdit = () => {
}`}
>
<PromptEditor
modelInfo={selectedModelInfo}
selectedImageModel={selectedImageModel}
onModeChange={(mode: string) => {
store.setActiveTool(mode as ActiveEditTool);
}}
@@ -382,9 +372,9 @@ const PageEdit = () => {
canvasWidth.current = width;
canvasHeight.current = height;
}}
fillColor={store.fillColor}
//fillColor={store.fillColor}
activeTool={store.activeTool}
brushColor={store.brushColor}
brushColor={"rgba(39, 187, 245, 0.54)"}
brushSize={store.brushSize}
onSelectionChange={setIsSelecting}
stageRef={stageRef}
@@ -395,7 +385,7 @@ const PageEdit = () => {
</ContextMenuContainer>
</div>
<div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
<ModelSelector
<ClassyModelSelector
items={IMAGE_EDITOR_PAGE_MODEL_LIST}
page={PAGE_ID}
panelTitle="Select Model"
@@ -7,7 +7,7 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
onGenerateClick,
isDisabled,
onFitPressed,
modelInfo,
selectedImageModel,
generationCount,
onGenerationCountChange,
supportsMaskedInpainting
@@ -20,7 +20,7 @@ const PromptEditor: React.FC<PromptBoxEditProps> = ({
onGenerateClick={onGenerateClick}
isDisabled={isDisabled}
onFitPressed={onFitPressed}
modelInfo={modelInfo}
selectedImageModel={selectedImageModel}
generationCount={generationCount}
onGenerationCountChange={onGenerationCountChange}
supportsMaskedInpainting={supportsMaskedInpainting}
@@ -3,6 +3,7 @@ import { LineNode, generateId } from "~/pages/PageDraw/stores/SceneState";
import { Node } from "~/pages/PageDraw/Node";
import { BaseSelectorImage } from "../BaseImageSelector";
import { FetchProxy } from "libs/tauri-utils/src/lib/FetchProxy";
import { ImageBundle } from "../HistoryStack";
export type ActiveEditTool = "select" | "edit" | "expand";
export type EditOperation = "add" | "minus";
@@ -126,6 +127,11 @@ interface EditState {
// pushHistory: () => void;
// clearHistory: () => void;
RESET: () => void;
// For the history stack
historyImageBundles: ImageBundle[];
clearHistoryImages: () => void;
addHistoryImageBundle: (bundle: ImageBundle) => void;
}
// Add a flag to prevent saving state during restoration
@@ -172,6 +178,9 @@ export const useEditStore = create<EditState>((set, get, store) => ({
baseImageInfo: null,
baseImageBitmap: null,
// History stack state
historyImageBundles: [],
// Actions
addNode: (node: Node) => {
set((state) => {
@@ -1009,5 +1018,13 @@ export const useEditStore = create<EditState>((set, get, store) => ({
RESET: () => {
set(store.getInitialState());
}
},
clearHistoryImages: () => {
set({ historyImageBundles: [] });
},
addHistoryImageBundle(bundle) {
set((state) => ({ historyImageBundles: [...state.historyImageBundles, bundle] }));
},
}));
@@ -68,21 +68,20 @@ import {
import {
STAGE_3D_PAGE_MODEL_LIST,
ModelPage,
ModelSelector,
useModelSelectorStore,
getSelectedImageModel,
ClassyModelSelector,
} from "@storyteller/ui-model-selector";
import { LoginModal, useLoginModalStore } from "@storyteller/ui-login-modal";
import PageDraw from "../PageDraw/PageDraw";
import { useTabStore } from "../Stores/TabState";
import PageEdit from "../PageEdit/PageEdit";
import { ModelInfo } from "@storyteller/model-list";
import { ImageModel } from "@storyteller/model-list";
const PAGE_ID : ModelPage = ModelPage.Stage3D;
export const PageEditor = () => {
useSignals();
const { triggerRecheck } = useLoginModalStore();
const { selectedModels } = useModelSelectorStore();
const { status } = authentication;
useEffect(() => {
@@ -96,12 +95,6 @@ export const PageEditor = () => {
event.stopPropagation();
};
// TODO: For 3d Promptbox to accept this later on
// const { selectedModels } = useModelSelectorStore();
// const selectedModel =
// selectedModels[ModelCategory.Editor3D] ||
// videoGenerationModels[0]?.label;
useEffect(() => {
timelineHeight.value = 0; //timelineHeight.value = 208;
sidePanelWidth.value = 340;
@@ -110,14 +103,7 @@ export const PageEditor = () => {
};
}, []);
const selectedModel =
selectedModels[PAGE_ID] ||
STAGE_3D_PAGE_MODEL_LIST[0]?.label;
const selectedModelInfo: ModelInfo | undefined =
STAGE_3D_PAGE_MODEL_LIST.find(
(m) => m.label === selectedModel,
)?.modelInfo;
const selectedImageModel : ImageModel | undefined = getSelectedImageModel(PAGE_ID);
const height =
dndTimelineHeight.value > -1
@@ -526,7 +512,7 @@ export const PageEditor = () => {
handleCameraNameChange={handleCameraNameChange}
handleCameraFocalLengthChange={handleCameraFocalLengthChange}
onAspectRatioSelect={onAspectRatioSelect}
selectedModelInfo={selectedModelInfo}
selectedImageModel={selectedImageModel}
setEnginePrompt={(prompt) => {
console.log("setEnginePrompt", prompt);
if (!editorEngine) {
@@ -548,7 +534,7 @@ export const PageEditor = () => {
/>
<div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
<ModelSelector
<ClassyModelSelector
items={STAGE_3D_PAGE_MODEL_LIST}
page={PAGE_ID}
panelTitle="Select Model"
@@ -5,15 +5,14 @@ import BackgroundGallery from "./BackgroundGallery";
import {
TEXT_TO_IMAGE_PAGE_MODEL_LIST,
ModelPage,
ModelSelector,
useModelSelectorStore,
ClassyModelSelector,
getSelectedImageModel,
} from "@storyteller/ui-model-selector";
import { ImageModel } from "@storyteller/model-list";
interface TextToImageProps {
imageMediaId?: string;
imageUrl?: string;
}
import { ModelInfo } from "@storyteller/model-list";
import { useTextToImageGenerationCompleteEvent } from "@storyteller/tauri-events";
import { useTextToImageStore } from "./TextToImageStore";
import { animated, useSpring } from "@react-spring/web";
@@ -27,19 +26,12 @@ const PAGE_ID: ModelPage = ModelPage.TextToImage;
const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
const containerRef = useRef<HTMLDivElement>(null);
const { selectedModels } = useModelSelectorStore();
const batches = useTextToImageStore((s) => s.batches);
const startBatch = useTextToImageStore((s) => s.startBatch);
const completeBatch = useTextToImageStore((s) => s.completeBatch);
const resetBatches = useTextToImageStore((s) => s.reset);
const selectedModel =
selectedModels[PAGE_ID] || TEXT_TO_IMAGE_PAGE_MODEL_LIST[0]?.label;
const selectedModelInfo: ModelInfo | undefined =
TEXT_TO_IMAGE_PAGE_MODEL_LIST.find(
(m) => m.label === selectedModel,
)?.modelInfo;
const selectedImageModel : ImageModel | undefined = getSelectedImageModel(PAGE_ID);
const jobContext: JobContextType = {
jobTokens: [],
@@ -76,6 +68,10 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
config: { tension: 200, friction: 28, mass: 1.1 },
});
// Show the batches in reverse order, with the newest items at top.
// Like Midjourney instead of a "chat history" style.
const inverseBatch = batches.toReversed();
return (
<div
ref={containerRef}
@@ -95,19 +91,8 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
{hasAnyBatches && (
<div className="h-full w-full overflow-y-auto pb-40">
<div className="mx-auto flex max-w-screen-2xl flex-col gap-8 pr-2">
{batches.map((batch) => (
{inverseBatch.map((batch) => (
<div key={batch.id} className="flex items-start gap-4">
<div>
<div className="glass inline-block w-[320px] shrink-0 rounded-xl px-4 py-3 text-left text-sm text-white/90">
<div>{batch.prompt}</div>
</div>
<div className="mt-2 flex justify-end">
<Badge
label={batch.modelLabel}
className="px-2 py-1 text-xs opacity-70"
/>
</div>
</div>
<div className="grid flex-1 grid-cols-4 gap-4">
{batch.status === "pending" && batch.images.length === 0
? Array.from({
@@ -147,6 +132,17 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
</button>
))}
</div>
<div>
<div className="glass inline-block w-[320px] shrink-0 rounded-xl px-4 py-3 text-left text-sm text-white/90">
<div>{batch.prompt}</div>
</div>
<div className="mt-2 flex justify-end">
<Badge
label={batch.modelLabel}
className="px-2 py-1 text-xs opacity-70"
/>
</div>
</div>
</div>
))}
</div>
@@ -171,12 +167,12 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
useJobContext={() => {
return jobContext;
}}
model={selectedModel}
modelInfo={selectedModelInfo}
selectedModel={selectedImageModel}
imageMediaId={imageMediaId}
url={imageUrl ?? undefined}
onEnqueuePressed={async (prompt, count, subscriberId) => {
startBatch(prompt, count, selectedModel, subscriberId);
const modelLabel = selectedImageModel?.fullName ?? "";
startBatch(prompt, count, modelLabel, subscriberId);
}}
/>
</animated.div>
@@ -184,7 +180,7 @@ const TextToImage = ({ imageMediaId, imageUrl }: TextToImageProps) => {
{!showPromptAtBottom && <BackgroundGallery />}
<div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
<ModelSelector
<ClassyModelSelector
items={TEXT_TO_IMAGE_PAGE_MODEL_LIST}
page={PAGE_ID}
mode="hoverSelect"
@@ -3,12 +3,12 @@ import { JobContextType } from "@storyteller/common";
import { PromptBoxVideo } from "@storyteller/ui-promptbox";
import BackgroundGallery from "./BackgroundGallery";
import {
ModelPage,
ModelSelector,
useModelSelectorStore,
ClassyModelSelector,
IMAGE_TO_VIDEO_PAGE_MODEL_LIST,
ModelPage,
} from "@storyteller/ui-model-selector";
import { ModelInfo } from "@storyteller/model-list";
import { VideoModel } from "@storyteller/model-list";
import { getSelectedVideoModel } from "@storyteller/ui-model-selector";
const PAGE_ID : ModelPage = ModelPage.ImageToVideo;
@@ -19,15 +19,8 @@ interface ImageToVideoProps {
const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
const containerRef = useRef<HTMLDivElement>(null);
const { selectedModels } = useModelSelectorStore();
const selectedModel =
selectedModels[PAGE_ID] ||
IMAGE_TO_VIDEO_PAGE_MODEL_LIST[0]?.label;
const selectedModelInfo: ModelInfo | undefined = IMAGE_TO_VIDEO_PAGE_MODEL_LIST.find(
(m) => m.label === selectedModel,
)?.modelInfo;
const selectedVideoModel : VideoModel | undefined = getSelectedVideoModel(PAGE_ID);
const jobContext: JobContextType = {
jobTokens: [],
@@ -56,8 +49,7 @@ const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
console.log("Using job context");
return jobContext;
}}
model={selectedModel}
modelInfo={selectedModelInfo}
selectedModel={selectedVideoModel}
imageMediaId={imageMediaId}
url={imageUrl ?? undefined}
onEnqueuePressed={async () => {
@@ -68,7 +60,7 @@ const ImageToVideo = ({ imageMediaId, imageUrl }: ImageToVideoProps) => {
<BackgroundGallery />
<div className="absolute bottom-6 left-6 z-20 flex items-center gap-2">
<ModelSelector
<ClassyModelSelector
items={IMAGE_TO_VIDEO_PAGE_MODEL_LIST}
page={PAGE_ID}
panelTitle="Select Model"
@@ -3,7 +3,7 @@ import { useSceneStore } from '../PageDraw/stores/SceneState';
export type TabId = '2D' | '3D' | 'VIDEO' | 'IMAGE' | 'EDIT';
const DEFAULT_TAB : TabId = '2D';
const DEFAULT_TAB : TabId = 'IMAGE';
interface TabState {
// Current active tab
@@ -1,4 +1,5 @@
export * from "./lib/model-selector";
export * from "./lib/model-selector-store";
export * from "./lib/model-lists";
export * from "./lib/model-selectors-for-pages";
export * from "./lib/model-pages";
export * from "./lib/classy-model-selector";
export * from "./lib/classy-model-selector-store";
@@ -0,0 +1,37 @@
import { create } from "zustand";
import { ModelPage } from "./model-pages";
import { ImageModel, Model, VideoModel } from "@storyteller/model-list";
interface ClassyModelSelectorState {
selectedModels: { [page in ModelPage]?: Model };
setSelectedModel: (page: ModelPage, model: Model) => void;
}
export const useClassyModelSelectorStore = create<ClassyModelSelectorState>((set) => ({
selectedModels: {},
setSelectedModel: (page, model) =>
set((state) => ({
selectedModels: {
...state.selectedModels,
[page]: model,
},
})),
}));
export const getSelectedImageModel = (page: ModelPage) : ImageModel | undefined => {
const { selectedModels } = useClassyModelSelectorStore();
const maybeModel = selectedModels[page];
if (!maybeModel || !(maybeModel instanceof ImageModel)) {
return undefined;
}
return maybeModel;
};
export const getSelectedVideoModel = (page: ModelPage) : VideoModel | undefined => {
const { selectedModels } = useClassyModelSelectorStore();
const maybeModel = selectedModels[page];
if (!maybeModel || !(maybeModel instanceof VideoModel)) {
return undefined;
}
return maybeModel;
};
@@ -1,9 +1,9 @@
import { PopoverMenu, type PopoverItem } from "@storyteller/ui-popover";
import { useModelSelectorStore } from "./model-selector-store";
import { useClassyModelSelectorStore } from "./classy-model-selector-store";
import { useEffect, useMemo } from "react";
import { ModelPage } from "./model-pages";
interface ModelSelectorProps {
interface ClassyModelSelectorProps {
items: Omit<PopoverItem, "selected">[];
page: ModelPage;
mode?: "hoverSelect" | "default" | "toggle" | "button";
@@ -14,31 +14,31 @@ interface ModelSelectorProps {
triggerLabel?: string;
}
export function ModelSelector({
export function ClassyModelSelector({
items,
page,
...popoverProps
}: ModelSelectorProps) {
const { selectedModels, setSelectedModel } = useModelSelectorStore();
const selectedModel = selectedModels[page] || items[0]?.label;
}: ClassyModelSelectorProps) {
const { selectedModels, setSelectedModel } = useClassyModelSelectorStore();
const selectedModel = selectedModels[page] || items[0]?.model;
// For the first mount, make sure the selected model is set for other components to listen
useEffect(() => {
// Initialize selected model if not set
if (!selectedModels[page] && items[0]) {
setSelectedModel(page, items[0].label);
setSelectedModel(page, items[0].model!);
}
}, []);
const handleModelSelect = (item: PopoverItem) => {
setSelectedModel(page, item.label);
setSelectedModel(page, item.model!);
};
const modelList = useMemo(
() =>
items.map((item) => ({
...item,
selected: item.label === selectedModel,
selected: item.model === selectedModel,
})),
[items, selectedModel]
);
@@ -53,4 +53,4 @@ export function ModelSelector({
);
}
export default ModelSelector;
export default ClassyModelSelector;
@@ -1,57 +0,0 @@
import type { PopoverItem } from "@storyteller/ui-popover";
import { faClock, faFilm, faImage } from "@fortawesome/pro-solid-svg-icons";
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import {
getCreatorIcon,
ModelConfig,
TEXT_TO_IMAGE_PAGE_MODELS,
IMAGE_TO_VIDEO_PAGE_MODELS,
CANVAS_2D_PAGE_MODELS,
STAGE_3D_PAGE_MODELS,
IMAGE_EDITOR_PAGE_MODELS,
} from "@storyteller/model-list";
export type ModelList = Omit<PopoverItem, "selected">[];
const withIcon = (creatorIcon: any, fallback: any) => creatorIcon || fallback;
const buildItems = (
models: ModelConfig[],
fallbackIcon: any
) =>
models.map((m) => ({
label: m.label,
icon: withIcon(getCreatorIcon(m.info.creator), fallbackIcon),
description: m.description,
badges: m.badges?.map((b) => ({
label: b.label,
icon: <FontAwesomeIcon icon={faClock} />,
})),
modelConfig: m, // Access to full object.
modelInfo: m.info,
}));
export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems(
TEXT_TO_IMAGE_PAGE_MODELS,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems(
IMAGE_TO_VIDEO_PAGE_MODELS,
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
);
export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems(
CANVAS_2D_PAGE_MODELS,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const STAGE_3D_PAGE_MODEL_LIST : ModelList = buildItems(
STAGE_3D_PAGE_MODELS,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const IMAGE_EDITOR_PAGE_MODEL_LIST : ModelList = buildItems(
IMAGE_EDITOR_PAGE_MODELS,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
@@ -1,18 +0,0 @@
import { create } from "zustand";
import { ModelPage } from "./model-pages";
interface ModelSelectorState {
selectedModels: { [page in ModelPage]?: string };
setSelectedModel: (page: ModelPage, modelLabel: string) => void;
}
export const useModelSelectorStore = create<ModelSelectorState>((set) => ({
selectedModels: {},
setSelectedModel: (page, modelLabel) =>
set((state) => ({
selectedModels: {
...state.selectedModels,
[page]: modelLabel,
},
})),
}));
@@ -1,12 +0,0 @@
import { render } from '@testing-library/react';
import ModelSelector from './model-selector';
describe('ModelSelector', () => {
it('should render successfully', () => {
const { baseElement } = render(<ModelSelector />);
expect(baseElement).toBeTruthy();
});
});
@@ -0,0 +1,71 @@
import type { PopoverItem } from "@storyteller/ui-popover";
import { faClock, faFilm, faImage } from "@fortawesome/pro-solid-svg-icons";
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import {
getCreatorIcon,
Model,
VIDEO_MODELS,
IMAGE_MODELS,
ALL_MODELS_BY_ID,
} from "@storyteller/model-list";
import { ModelTag } from "@storyteller/model-list";
export type ModelList = Omit<PopoverItem, "selected">[];
const withIcon = (creatorIcon: any, fallback: any) => creatorIcon || fallback;
const buildItems2 = (
models: Model[],
fallbackIcon: any
) =>
models.map((model : Model) => ({
label: model.selectorName,
icon: withIcon(getCreatorIcon(model.creator), fallbackIcon),
description: model.selectorDescription,
badges: model.toLegacyBadges()?.map((b) => ({
label: b.label,
icon: <FontAwesomeIcon icon={faClock} />,
})),
modelConfig: model.toLegacyModelConfig(), // Access to full object.
model: model,
}));
/**
* IN-PROGRESS MIGRATION (messy for now)
* We're gradually going to phase out ModelList, ModelConfig, etc.
* We won't index by name, but rather id, or simply will always have full
* access to the object directly.
*/
export const CANVAS_2D_PAGE_MODEL_LIST : ModelList = buildItems2(
IMAGE_MODELS.filter((m) => m.tags?.includes(ModelTag.InstructiveEdit)),
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const STAGE_3D_PAGE_MODEL_LIST : ModelList = buildItems2(
IMAGE_MODELS.filter((m) => m.tags?.includes(ModelTag.InstructiveEdit)),
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const IMAGE_EDITOR_PAGE_MODEL_LIST : ModelList = buildItems2(
//[
// ALL_MODELS_BY_ID.get("flux_pro_inpaint")!,
// ALL_MODELS_BY_ID.get("flux_dev_juggernaut_inpaint")!,
// ALL_MODELS_BY_ID.get("flux_pro_kontext_max")!,
//],
IMAGE_MODELS.filter((m) => m.canEditImages),
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
export const IMAGE_TO_VIDEO_PAGE_MODEL_LIST : ModelList = buildItems2(
VIDEO_MODELS,
<FontAwesomeIcon icon={faFilm} className="h-4 w-4" />
);
export const TEXT_TO_IMAGE_PAGE_MODEL_LIST : ModelList = buildItems2(
IMAGE_MODELS
.filter((model) => !model.usesInpaintingMask) // We can't use masked inpainting models here
.filter((model) => model.id !== "flux_pro_kontext_max") // NB: Flux Pro Kontext Max requires an input image
,
<FontAwesomeIcon icon={faImage} className="h-4 w-4" />
);
@@ -13,7 +13,7 @@ import {
faChevronUp,
faCircleCheck,
} from "@fortawesome/pro-solid-svg-icons";
import { ModelInfo } from "@storyteller/model-list";
import { Model, ModelInfo } from "@storyteller/model-list";
export interface PopoverItem {
label: string;
@@ -28,6 +28,7 @@ export interface PopoverItem {
icon?: ReactNode;
}>;
modelInfo?: ModelInfo;
model?: Model; // NB: Let's migrate to using this.
}
interface PopoverMenuProps {
@@ -35,19 +35,13 @@ import { IsDesktopApp } from "@storyteller/tauri-utils";
import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
import { PromptsApi } from "@storyteller/api";
import {
// SoraImageRemix,
// SoraImageRemixAspectRatio,
// CheckSoraSession,
// SoraSessionState,
// waitForSoraLogin,
EnqueueContextualEditImage,
EnqueueContextualEditImageSize,
} from "@storyteller/tauri-api";
// import { showActionReminder } from "@storyteller/ui-action-reminder-modal";
// import { invoke } from "@tauri-apps/api/core";
import { usePrompt2DStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
import { getCapabilitiesForModel } from "@storyteller/model-list";
import { ImageModel } from "@storyteller/model-list";
export type AspectRatio = "1:1" | "3:2" | "2:3";
@@ -61,7 +55,7 @@ interface PromptBox2DProps {
assetFile: File;
progressCallback: (newState: UploaderState) => void;
}) => Promise<void>;
selectedModelInfo?: ModelInfo;
selectedImageModel?: ImageModel;
getCanvasRenderBitmap: () => MaybeCanvasRenderBitmapType;
EncodeImageBitmapToBase64: (imageBitmap: ImageBitmap) => Promise<string>;
useJobContext: () => JobContextType;
@@ -77,7 +71,7 @@ export const PromptBox2D = ({
useJobContext,
onEnqueuePressed,
onAspectRatioChange,
selectedModelInfo,
selectedImageModel,
onFitPressed,
}: PromptBox2DProps) => {
useSignals();
@@ -129,7 +123,7 @@ export const PromptBox2D = ({
// Update generation count options from selected model capabilities
useEffect(() => {
const caps = getCapabilitiesForModel(selectedModelInfo);
const caps = getCapabilitiesForModel(selectedImageModel);
const items: PopoverItem[] = Array.from(
{ length: caps.maxGenerationCount },
(_, i) => i + 1
@@ -146,7 +140,7 @@ export const PromptBox2D = ({
);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [selectedModelInfo]);
}, [selectedImageModel]);
// Keep UI selection in sync when store value changes
useEffect(() => {
@@ -318,31 +312,7 @@ export const PromptBox2D = ({
setPrompt(e.target.value);
};
// Helper to show Sora login reminder and wait for login
// const handleSoraLoginReminder = async () => {
// return new Promise<void>((resolve) => {
// showActionReminder({
// reminderType: "soraLogin",
// onPrimaryAction: async () => {
// await invoke("open_sora_login_command");
// await waitForSoraLogin();
// toast.success("Logged in to Sora!");
// resolve();
// },
// });
// });
// };
const handleTauriEnqueue = async () => {
// NB(bt): This needs to move to an error handler.
// // Check if the Sora session is valid
// const soraSession = await CheckSoraSession();
// if (soraSession.state !== SoraSessionState.Valid) {
// setIsEnqueueing(false);
// await handleSoraLoginReminder();
// return;
// }
const api = new PromptsApi();
let image = getCanvasRenderBitmap();
if (image === undefined) {
@@ -379,7 +349,7 @@ export const PromptBox2D = ({
const aspectRatio = getCurrentAspectRatio();
const generateResponse = await EnqueueContextualEditImage({
model: selectedModelInfo,
model: selectedImageModel,
scene_image_media_token: snapshotMediaToken.data!,
image_media_tokens: referenceImages.map((image) => image.mediaToken),
disable_system_prompt: !useSystemPrompt,
@@ -53,7 +53,7 @@ import {
// import { showActionReminder } from "@storyteller/ui-action-reminder-modal";
import { usePrompt3DStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
import { ModelInfo } from "@storyteller/model-list";
import { ImageModel } from "@storyteller/model-list";
interface PromptBox3DProps {
cameras: Signal<Camera[]>;
@@ -73,7 +73,7 @@ interface PromptBox3DProps {
handleCameraFocalLengthChange: (id: string, value: number) => void;
onAspectRatioSelect: (newRatio: CameraAspectRatio) => void;
setEnginePrompt: (prompt: string) => void;
selectedModelInfo?: ModelInfo;
selectedImageModel?: ImageModel;
snapshotCurrentFrame:
| ((shouldDownload?: boolean) => {
base64Snapshot: string;
@@ -100,7 +100,7 @@ export const PromptBox3D = ({
handleCameraFocalLengthChange,
onAspectRatioSelect,
setEnginePrompt,
selectedModelInfo,
selectedImageModel,
snapshotCurrentFrame,
}: PromptBox3DProps) => {
useSignals();
@@ -450,7 +450,7 @@ export const PromptBox3D = ({
const aspectRatio = getCurrentAspectRatio();
const generateResponse = await EnqueueContextualEditImage({
model: selectedModelInfo,
model: selectedImageModel,
scene_image_media_token: snapshotResult.data!,
image_media_tokens: referenceImages.map((image) => image.mediaToken),
disable_system_prompt: !useSystemPrompt,
@@ -15,7 +15,7 @@ import { ButtonIconSelect } from "@storyteller/ui-button-icon-select";
import { PopoverMenu, PopoverItem } from "@storyteller/ui-popover";
import { Tooltip } from "@storyteller/ui-tooltip";
import { useEffect, useRef, useState } from "react";
import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
import { ImageModel, getCapabilitiesForModel } from "@storyteller/model-list";
export interface PromptBoxEditProps {
onModeChange?: (mode: string) => void;
@@ -23,7 +23,7 @@ export interface PromptBoxEditProps {
onGenerateClick: (prompt: string) => void;
isDisabled?: boolean;
onFitPressed?: () => void | Promise<void>;
modelInfo?: ModelInfo;
selectedImageModel?: ImageModel;
generationCount?: number;
onGenerationCountChange?: (count: number) => void;
supportsMaskedInpainting?: boolean;
@@ -35,7 +35,7 @@ export const PromptBoxEdit = ({
onGenerateClick,
isDisabled,
onFitPressed,
modelInfo,
selectedImageModel,
generationCount: generationCountProp,
onGenerationCountChange,
supportsMaskedInpainting = false,
@@ -69,7 +69,7 @@ export const PromptBoxEdit = ({
// Build generation count options based on selected model
useEffect(() => {
const caps = getCapabilitiesForModel(modelInfo);
const caps = getCapabilitiesForModel(selectedImageModel);
console.log(">>> caps", caps);
const items: PopoverItem[] = Array.from(
{ length: caps.maxGenerationCount },
@@ -91,7 +91,7 @@ export const PromptBoxEdit = ({
onGenerationCountChange?.(clamped);
}
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [modelInfo]);
}, [selectedImageModel]);
// Keep UI selection in sync when store/value changes
useEffect(() => {
@@ -28,7 +28,7 @@ import {
import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import { IconProp } from "@fortawesome/fontawesome-svg-core";
import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
import { ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
import { ImageModel, ModelInfo, getCapabilitiesForModel } from "@storyteller/model-list";
import { usePromptImageStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
@@ -39,8 +39,7 @@ interface PromptBoxImageProps {
count: number,
subscriberId: string
) => void | Promise<void>;
model: string;
modelInfo?: ModelInfo;
selectedModel?: ImageModel;
imageMediaId?: string;
url?: string;
}
@@ -48,8 +47,7 @@ interface PromptBoxImageProps {
export const PromptBoxImage = ({
useJobContext,
onEnqueuePressed,
model,
modelInfo,
selectedModel,
imageMediaId,
url,
}: PromptBoxImageProps) => {
@@ -133,7 +131,7 @@ export const PromptBoxImage = ({
useEffect(() => {
// Build generation count options based on selected model
const caps = getCapabilitiesForModel(modelInfo);
const caps = getCapabilitiesForModel(selectedModel);
const defaultCount = Math.min(
Math.max(1, caps.defaultGenerationCount ?? 1),
caps.maxGenerationCount
@@ -151,7 +149,7 @@ export const PromptBoxImage = ({
}));
setGenerationCountList(items);
// eslint-disable-next-line react-hooks/exhaustive-deps
}, [modelInfo]);
}, [selectedModel]);
useEffect(() => {
setGenerationCountList((prev) =>
@@ -285,7 +283,7 @@ export const PromptBoxImage = ({
const generateResponse = await EnqueueTextToImage({
prompt: prompt,
model: modelInfo,
model: selectedModel,
aspect_ratio: aspectRatio,
number_images: generationCount,
frontend_caller: "text_to_image",
@@ -21,15 +21,14 @@ import { FontAwesomeIcon } from "@fortawesome/react-fontawesome";
import { IconProp } from "@fortawesome/fontawesome-svg-core";
import { IsDesktopApp } from "@storyteller/tauri-utils";
import { GalleryItem, GalleryModal } from "@storyteller/ui-gallery-modal";
import { ModelInfo } from "@storyteller/model-list";
import { VideoModel } from "@storyteller/model-list";
import { usePromptVideoStore, RefImage } from "./promptStore";
import { gtagEvent } from "@storyteller/google-analytics";
interface PromptBoxVideoProps {
useJobContext: () => JobContextType;
onEnqueuePressed?: () => void | Promise<void>;
model: string;
modelInfo?: ModelInfo;
selectedModel?: VideoModel;
imageMediaId?: string;
url?: string;
}
@@ -37,8 +36,7 @@ interface PromptBoxVideoProps {
export const PromptBoxVideo = ({
useJobContext,
onEnqueuePressed,
model,
modelInfo,
selectedModel,
imageMediaId,
url,
}: PromptBoxVideoProps) => {
@@ -194,7 +192,7 @@ export const PromptBoxVideo = ({
const handleEnqueue = async () => {
setIsEnqueueing(true);
console.log("PromptBoxVideo - Prompting with model", modelInfo);
console.log("PromptBoxVideo - Prompting with model", selectedModel);
gtagEvent("enqueue_video");
@@ -206,7 +204,7 @@ export const PromptBoxVideo = ({
}, 10000);
const generateResponse = await EnqueueImageToVideo({
model: modelInfo,
model: selectedModel,
image_media_token: referenceImages[0].mediaToken,
prompt: prompt,
});
+16 -7
View File
@@ -1,7 +1,16 @@
export * from "./lib/ModelConfig.js";
export * from "./lib/ModelCreator.js";
export * from "./lib/ModelCreatorIcons.js";
export * from "./lib/ModelInfo.js";
export * from "./lib/ModelMapping.js";
export * from "./lib/Models.js";
export * from "./lib/ModelsByPage.js";
export * from "./lib/classes/ImageModel.js";
export * from "./lib/classes/Model.js";
export * from "./lib/classes/VideoModel.js";
export * from "./lib/classes/metadata/ModelCreator.js";
export * from "./lib/classes/metadata/ModelCreatorIcons.js";
export * from "./lib/classes/metadata/ModelMapping.js";
export * from "./lib/classes/metadata/ModelTag.js";
export * from "./lib/legacy/ModelConfig.js";
export * from "./lib/legacy/ModelInfo.js";
export * from "./lib/legacy/Models.js";
export * from "./lib/legacy/ModelsByPage.js";
export * from "./lib/lists/AllModels.js";
export * from "./lib/lists/AllModelsById.js";
export * from "./lib/lists/ImageModels.js";
export * from "./lib/lists/VideoModels.js";
@@ -1,36 +0,0 @@
import { ALL_MODELS_BY_ID } from "./Models.js";
import { ModelConfig } from "./ModelConfig.js";
import { ModelTag } from "./ModelTag.js";
import { getModelsByCategory } from "./Models.js";
/**
* We might not want to always programatically build the list of models, because
* we likely want to control the order of models in the UI to prioritize
* more popular models. We may also want to selectively prune non-useful models.
*/
export const TEXT_TO_IMAGE_PAGE_MODELS : ModelConfig[] =
getModelsByCategory("image")
// Remove inpainting and editing models.
.filter((m) => m.id !== "flux_pro_inpaint")
.filter((m) => m.id !== "flux_pro_kontext_max");
export const IMAGE_TO_VIDEO_PAGE_MODELS : ModelConfig[] =
getModelsByCategory("video");
export const CANVAS_2D_PAGE_MODELS : ModelConfig[] =
getModelsByCategory("image")
// Remove non-instructive models.
.filter((m) => m.tags?.includes(ModelTag.InstructiveEdit))
export const STAGE_3D_PAGE_MODELS : ModelConfig[] =
getModelsByCategory("image")
// Remove non-instructive models.
.filter((m) => m.tags?.includes(ModelTag.InstructiveEdit))
export const IMAGE_EDITOR_PAGE_MODELS : ModelConfig[] = [
// Explicit list for now.
ALL_MODELS_BY_ID.get("flux_pro_inpaint")!,
ALL_MODELS_BY_ID.get("flux_dev_juggernaut_inpaint")!,
ALL_MODELS_BY_ID.get("flux_pro_kontext_max")!,
];
@@ -0,0 +1,49 @@
import { ModelCreator } from "src/index.js";
import { Model } from "./Model.js";
import { ModelCategory } from "../legacy/ModelConfig.js";
import { ModelTag } from "./metadata/ModelTag.js";
export class ImageModel extends Model {
// Maximum number of images that can be generated at once
readonly maxGenerationCount: number;
// Default number of images that can be generated at once
readonly defaultGenerationCount: number;
// Signals image editing models that focus on editing a single image.
readonly canEditImages: boolean;
// For inpainting models, does it require sending a mask?
readonly usesInpaintingMask: boolean;
constructor(args: {
id: string;
tauriId: string;
fullName: string;
category: ModelCategory;
creator: ModelCreator;
selectorName: string;
selectorDescription: string;
selectorBadges: string[];
maxGenerationCount: number;
defaultGenerationCount: number;
canEditImages?: boolean;
usesInpaintingMask?: boolean;
tags?: ModelTag[];
}) {
if (args.maxGenerationCount < 1) {
throw new Error("maxGenerationCount must be at least 1");
}
if (args.defaultGenerationCount < 1) {
throw new Error("defaultGenerationCount must be at least 1");
}
if (args.defaultGenerationCount > args.maxGenerationCount) {
throw new Error("defaultGenerationCount must be less than or equal to maxGenerationCount");
}
super(args);
this.maxGenerationCount = args.maxGenerationCount;
this.defaultGenerationCount = args.defaultGenerationCount;
this.canEditImages = args.canEditImages ?? false;
this.usesInpaintingMask = args.usesInpaintingMask ?? false;
}
}
@@ -0,0 +1,82 @@
import { ModelCreator } from "./metadata/ModelCreator.js";
import { ModelCategory, ModelConfig } from "../legacy/ModelConfig.js";
import { ModelTag } from "./metadata/ModelTag.js";
// NB: Do not create instances of this class directly, use subclasses.
export class Model {
// A unique frontend-only string for the model
readonly id: string;
// A unique identifier that Tauri uses for the model (this
// might differ from our backend or other systems)
readonly tauriId: string;
// A long name for the model that might need to be abbreviated.
readonly fullName: string;
// The type of model (image, video, etc.)
// TODO: Not sure that this is used for anything
readonly category: ModelCategory;
// What company made the model.
readonly creator: ModelCreator;
// Name for the selector
readonly selectorName: string;
// Description for the selector
readonly selectorDescription: string;
// Labels for the selector
readonly selectorBadges: string[];
// A list of filterable "capabilities" that can be used to filter models.
readonly tags: ModelTag[];
protected constructor(args: {
id: string;
tauriId: string;
fullName: string;
category: ModelCategory;
creator: ModelCreator;
selectorName: string;
selectorDescription: string;
selectorBadges: string[];
tags?: ModelTag[];
}) {
this.id = args.id;
this.tauriId = args.tauriId;
this.fullName = args.fullName;
this.category = args.category;
this.creator = args.creator;
this.selectorName = args.selectorName;
this.selectorDescription = args.selectorDescription;
this.selectorBadges = args.selectorBadges;
this.tags = args.tags ?? [];
}
toLegacyBadges() : { label: string }[] {
return this.selectorBadges.map((b) => ({ label: b }));
}
// TODO: This is a method to support migration. Kill it after we no longer need it.
toLegacyModelConfig(): ModelConfig {
return {
id: this.id,
label: this.selectorName,
description: this.selectorDescription,
badges: this.toLegacyBadges(),
category: this.category,
info: {
name: this.fullName,
tauri_id: this.tauriId,
creator: this.creator,
},
capabilities: {
maxGenerationCount: 9, // NB: Sentinel value to detect continued use
defaultGenerationCount: 9, // NB: Sentinel value to detect continued use
},
tags: [],
};
}
}
@@ -0,0 +1,30 @@
import { ModelCreator } from "src/index.js";
import { Model } from "./Model.js";
import { ModelCategory } from "../legacy/ModelConfig.js";
import { ModelTag } from "./metadata/ModelTag.js";
export class VideoModel extends Model {
// Whether the model supports image starting frames
readonly startFrame: boolean;
// Whether the model supports image ending frames
readonly endFrame: boolean;
constructor(args: {
id: string;
tauriId: string;
fullName: string;
category: ModelCategory;
creator: ModelCreator;
selectorName: string;
selectorDescription: string;
selectorBadges: string[];
startFrame: boolean;
endFrame: boolean;
tags?: ModelTag[];
}) {
super(args);
this.startFrame = args.startFrame;
this.endFrame = args.endFrame;
}
}
@@ -1,4 +1,6 @@
// TODO: It might be better to have top-level boolean fields on the models.
export enum ModelTag {
// Models that show up in the 2D and 3D editors
InstructiveEdit = "instructiveEdit",
@@ -8,4 +10,4 @@ export enum ModelTag {
// Models that show up in the inpainting editor, but that do not use masking.
NonMaskedInpainting = "nonMaskedInpainting",
}
}

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