From 03354b73d9ae960ab90f2d9bdd1654522931dfaa Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Fri, 15 May 2026 01:03:22 -0400 Subject: [PATCH] split image gen --- .../cost/image/omni_gen_image_cost_handler.rs | 2 +- .../endpoints/omni_gen/generate/image/mod.rs | 5 +- .../image/omni_gen_image_generate_handler.rs | 70 +++--- .../generate/image/pipeline_result.rs | 8 + .../hydrate_to_router_request.rs | 0 .../{ => pipeline_v1}/distill_helper/mod.rs | 0 .../distill_helper/request_to_costs.rs | 0 .../distill_helper/request_to_plan.rs | 0 .../distill_image_request.rs | 0 .../generate/image/pipeline_v1/mod.rs | 3 + .../image/pipeline_v1/run_pipeline_v1.rs | 73 +++++++ .../generate/image/pipeline_v2/mod.rs | 1 + .../image/pipeline_v2/run_pipeline_v2.rs | 203 ++++++++++++++++++ .../generate/image/tests/test_flux_1_dev.rs | 2 +- .../image/tests/test_flux_1_schnell.rs | 2 +- .../generate/image/tests/test_flux_pro_1p1.rs | 2 +- .../image/tests/test_flux_pro_1p1_ultra.rs | 2 +- .../generate/image/tests/test_gpt_image_1.rs | 2 +- .../image/tests/test_gpt_image_1p5.rs | 2 +- .../generate/image/tests/test_gpt_image_2.rs | 2 +- .../generate/image/tests/test_nano_banana.rs | 2 +- .../image/tests/test_nano_banana_2.rs | 2 +- .../image/tests/test_nano_banana_pro.rs | 2 +- .../generate/image/tests/test_seedream_4.rs | 2 +- .../generate/image/tests/test_seedream_4p5.rs | 2 +- .../image/tests/test_seedream_5_lite.rs | 2 +- 26 files changed, 333 insertions(+), 58 deletions(-) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{ => pipeline_v1}/distill_helper/hydrate_to_router_request.rs (100%) rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{ => pipeline_v1}/distill_helper/mod.rs (100%) rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{ => pipeline_v1}/distill_helper/request_to_costs.rs (100%) rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{ => pipeline_v1}/distill_helper/request_to_plan.rs (100%) rename crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/{ => pipeline_v1}/distill_image_request.rs (100%) create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs create mode 100644 crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs index 81ea5691d6..1ec511bbea 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/cost/image/omni_gen_image_cost_handler.rs @@ -8,7 +8,7 @@ use artcraft_router::api::provider::Provider; use log::warn; use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; -use crate::http_server::endpoints::omni_gen::generate::image::distill_helper::hydrate_to_router_request::hydrate_to_router_request; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_helper::hydrate_to_router_request::hydrate_to_router_request; use crate::state::server_state::ServerState; /// Estimate the cost of an image generation. diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs index 8fdd478417..27181efccc 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/mod.rs @@ -1,6 +1,7 @@ pub mod omni_gen_image_generate_handler; -pub(crate) mod distill_helper; -pub(crate) mod distill_image_request; +pub(crate) mod pipeline_result; +pub mod pipeline_v1; +pub mod pipeline_v2; mod resolve_media_tokens; #[cfg(test)] diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs index 513d546958..2b430d79f3 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/omni_gen_image_generate_handler.rs @@ -9,8 +9,6 @@ use url::Url; use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; use artcraft_api_defs::omni_gen::generate_response::omni_gen_image_generate_response::OmniGenImageGenerateResponse; -use artcraft_router::client::router_client::RouterClient; -use artcraft_router::client::router_fal_client::RouterFalClient; use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; use enums::by_table::debug_logs::debug_log_type::DebugLogType; use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType; @@ -31,19 +29,18 @@ use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_it insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem, }; use mysql_queries::queries::prompts::insert_prompt::{insert_prompt, InsertPromptArgs}; -use tokens::tokens::generic_inference_jobs::InferenceJobToken; use tokens::tokens::media_files::MediaFileToken; use tokens::tokens::non_unique::debug_logs_event_token::DebugLogEventToken; -use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; use crate::http_server::endpoints::generate::common::payments_error_test::payments_error_test; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_helper::hydrate_to_router_request::hydrate_to_router_request; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::run_pipeline_v1::{run_pipeline_v1, RunPipelineV1Args}; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v2::run_pipeline_v2::{run_pipeline_v2, should_use_pipeline_v2, RunPipelineV2Args}; use crate::http_server::validations::validate_idempotency_token_format::validate_idempotency_token_format; use crate::state::server_state::ServerState; use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map; -use super::distill_image_request::distill_image_request; - /// Generate an image using the omni-gen unified endpoint. #[utoipa::path( post, @@ -138,26 +135,9 @@ pub async fn omni_gen_image_generate_handler( _ => None, }; - // ==================== DISTILL ==================== // + // ==================== HYDRATE ROUTER REQUEST ==================== // - let distilled = distill_image_request(&request, media_file_hydration_map.as_ref())?; - - // ==================== BILLING ==================== // - - let cost = distilled.cost.cost_in_credits.unwrap_or(0); - - info!("Charging wallet: {} credits", cost); - - let apriori_job_token = InferenceJobToken::generate(); - - if cost > 0 { - attempt_wallet_deduction_else_common_web_error( - user_token, - Some(apriori_job_token.as_str()), - cost, - &mut mysql_connection, - ).await?; - } + let router_builder = hydrate_to_router_request(&request)?; // ==================== DEBUG LOG: HTTP REQUEST ==================== // @@ -172,25 +152,31 @@ pub async fn omni_gen_image_generate_handler( warn!("Failed to insert HTTP request debug log: {:?}", err); } - // ==================== EXECUTE GENERATION ==================== // + // ==================== PIPELINE DISPATCH ==================== // - let fal_client = RouterFalClient::new( - server_state.fal.api_key.clone(), - server_state.fal.webhook_url.clone(), - ); - - let router_client = RouterClient::Fal(fal_client); - - let generation_response = distilled.plan().generate_image(&router_client) - .await - .map_err(|e| { - warn!("Image generation failed: {:?}", e); - AdvancedCommonWebError::from_error(e) - })?; + let pipeline_result = if should_use_pipeline_v2(&router_builder) { + info!("Using image pipeline v2"); + run_pipeline_v2(RunPipelineV2Args { + router_builder: &router_builder, + server_state: &server_state, + mysql_connection: &mut mysql_connection, + user_token, + media_file_hydration_map: &media_file_hydration_map, + }).await? + } else { + info!("Using image pipeline v1"); + run_pipeline_v1(RunPipelineV1Args { + request: &request, + server_state: &server_state, + mysql_connection: &mut mysql_connection, + user_token, + media_file_hydration_map: &media_file_hydration_map, + }).await? + }; // ==================== DEBUG LOG: FAL REQUEST ==================== // - if let GenerateImageResponse::Fal(ref fal_payload) = generation_response { + if let GenerateImageResponse::Fal(ref fal_payload) = pipeline_result.response { if let Some(ref outbound_request) = fal_payload.maybe_outbound_request { if let Err(err) = insert_debug_log(InsertDebugLogArgs { apriori_debug_log_event_token: Some(&debug_log_event_token), @@ -205,7 +191,7 @@ pub async fn omni_gen_image_generate_handler( } } - let external_job_id = match &generation_response { + let external_job_id = match &pipeline_result.response { GenerateImageResponse::Artcraft(p) => { p.inference_job_token.as_str().to_string() } @@ -291,7 +277,7 @@ pub async fn omni_gen_image_generate_handler( let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token( InsertGenericInferenceForFalWithAprioriJobTokenArgs { - apriori_job_token: &apriori_job_token, + apriori_job_token: &pipeline_result.apriori_job_token, uuid_idempotency_token: &idempotency_token, maybe_external_third_party_id: &external_job_id, fal_category: FalCategory::ImageGeneration, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs new file mode 100644 index 0000000000..80a094b166 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_result.rs @@ -0,0 +1,8 @@ +use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; + +/// Both image pipelines produce this, then the shared handler suffix writes DB rows. +pub struct ImagePipelineResult { + pub apriori_job_token: InferenceJobToken, + pub response: GenerateImageResponse, +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/hydrate_to_router_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/hydrate_to_router_request.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/hydrate_to_router_request.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/hydrate_to_router_request.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/mod.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/mod.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/mod.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/request_to_costs.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_costs.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/request_to_costs.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/request_to_plan.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_helper/request_to_plan.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_helper/request_to_plan.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_image_request.rs similarity index 100% rename from crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/distill_image_request.rs rename to crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/distill_image_request.rs diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs new file mode 100644 index 0000000000..4927f82fd1 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/mod.rs @@ -0,0 +1,3 @@ +pub(crate) mod distill_helper; +pub(crate) mod distill_image_request; +pub mod run_pipeline_v1; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs new file mode 100644 index 0000000000..b94e52532b --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v1/run_pipeline_v1.rs @@ -0,0 +1,73 @@ +use std::collections::HashMap; + +use log::{info, warn}; +use url::Url; + +use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_image_cost_and_generate_request::OmniGenImageCostAndGenerateRequest; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::client::router_fal_client::RouterFalClient; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_result::ImagePipelineResult; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::distill_image_request; +use crate::state::server_state::ServerState; + +pub struct RunPipelineV1Args<'a> { + pub request: &'a OmniGenImageCostAndGenerateRequest, + pub server_state: &'a ServerState, + pub mysql_connection: &'a mut sqlx::pool::PoolConnection, + pub user_token: &'a UserToken, + pub media_file_hydration_map: &'a Option>, +} + +pub async fn run_pipeline_v1( + args: RunPipelineV1Args<'_>, +) -> Result { + let RunPipelineV1Args { + request, + server_state, + mysql_connection, + user_token, + media_file_hydration_map, + } = args; + + let distilled = distill_image_request(request, media_file_hydration_map.as_ref())?; + + let cost = distilled.cost.cost_in_credits.unwrap_or(0); + + info!("Charging wallet: {} credits", cost); + + let apriori_job_token = InferenceJobToken::generate(); + + if cost > 0 { + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + mysql_connection, + ).await?; + } + + let fal_client = RouterFalClient::new( + server_state.fal.api_key.clone(), + server_state.fal.webhook_url.clone(), + ); + + let router_client = RouterClient::Fal(fal_client); + + let response = distilled.plan().generate_image(&router_client) + .await + .map_err(|e| { + warn!("Image generation failed: {:?}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(ImagePipelineResult { + apriori_job_token, + response, + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs new file mode 100644 index 0000000000..f207b21293 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/mod.rs @@ -0,0 +1 @@ +pub mod run_pipeline_v2; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs new file mode 100644 index 0000000000..23271c02fe --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/pipeline_v2/run_pipeline_v2.rs @@ -0,0 +1,203 @@ +use std::collections::HashMap; + +use log::{info, warn}; +use url::Url; + +use artcraft_router::api::image_list_ref::ImageListRef; +use artcraft_router::api::provider::Provider; +use artcraft_router::client::router_client::RouterClient; +use artcraft_router::client::router_fal_webhook_optional_client::RouterFalWebhookOptionalClient; +use artcraft_router::generate::generate_image::generate_image_request_builder::GenerateImageRequestBuilder; +use artcraft_router::generate::generate_image::generate_image_response::GenerateImageResponse; +use artcraft_router::generate::generate_image_v2::image_generation_draft_context::ImageGenerationDraftContext; +use artcraft_router::generate::generate_image_v2::image_generation_draft_or_request::ImageGenerationDraftOrRequest; +use artcraft_router::generate::generate_image_v2::image_generation_request::ImageGenerationRequest; +use tokens::tokens::generic_inference_jobs::InferenceJobToken; +use tokens::tokens::media_files::MediaFileToken; +use tokens::tokens::users::UserToken; + +use crate::billing::wallets::attempt_wallet_deduction::attempt_wallet_deduction_else_common_web_error; +use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError; +use crate::http_server::endpoints::omni_gen::generate::image::pipeline_result::ImagePipelineResult; +use crate::state::server_state::ServerState; + +pub struct RunPipelineV2Args<'a> { + pub router_builder: &'a GenerateImageRequestBuilder, + pub server_state: &'a ServerState, + pub mysql_connection: &'a mut sqlx::pool::PoolConnection, + pub user_token: &'a UserToken, + pub media_file_hydration_map: &'a Option>, +} + +pub fn should_use_pipeline_v2(router_builder: &GenerateImageRequestBuilder) -> bool { + let mut execution_builder = router_builder.clone(); + execution_builder.provider = Provider::Fal; + execution_builder.use_new_builder() +} + +pub async fn run_pipeline_v2( + args: RunPipelineV2Args<'_>, +) -> Result { + let RunPipelineV2Args { + router_builder, + server_state, + mysql_connection, + user_token, + media_file_hydration_map, + } = args; + + let hydrated_builder = apply_hydrated_media_inputs( + router_builder, + media_file_hydration_map.as_ref(), + )?; + + let draft_or_request = build_execution_request(&hydrated_builder)?; + let cost = estimate_cost_in_credits(&hydrated_builder)?; + + info!("Charging wallet: {} credits", cost); + + let apriori_job_token = InferenceJobToken::generate(); + + if cost > 0 { + attempt_wallet_deduction_else_common_web_error( + user_token, + Some(apriori_job_token.as_str()), + cost, + mysql_connection, + ).await?; + } + + let response = finalize_and_generate(draft_or_request, server_state).await?; + + Ok(ImagePipelineResult { + apriori_job_token, + response, + }) +} + +fn build_execution_request( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + let mut execution_builder = router_builder.clone(); + execution_builder.provider = Provider::Fal; + + execution_builder.build2().map_err(|e| { + warn!("Failed to build2 for image v2 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + }) +} + +fn apply_hydrated_media_inputs( + router_builder: &GenerateImageRequestBuilder, + media_file_hydration_map: Option<&HashMap>, +) -> Result { + let mut hydrated_builder = router_builder.clone(); + let image_input_urls = build_image_input_urls( + hydrated_builder.image_inputs.as_ref(), + media_file_hydration_map, + )?; + + if let Some(urls) = image_input_urls { + hydrated_builder.image_inputs = Some(ImageListRef::Urls(urls)); + } + + Ok(hydrated_builder) +} + +fn build_image_input_urls( + image_inputs: Option<&ImageListRef>, + hydration_map: Option<&HashMap>, +) -> Result>, AdvancedCommonWebError> { + let tokens = match image_inputs { + Some(ImageListRef::MediaFileTokens(tokens)) if !tokens.is_empty() => tokens, + _ => return Ok(None), + }; + + let map = hydration_map.ok_or_else(|| { + AdvancedCommonWebError::BadInputWithSimpleMessage( + "image_media_tokens supplied but no hydration map was provided".to_string(), + ) + })?; + + let mut urls: Vec = Vec::with_capacity(tokens.len()); + for token in tokens { + match map.get(token) { + Some(url) => urls.push(url.to_string()), + None => { + return Err(AdvancedCommonWebError::BadInputWithSimpleMessage(format!( + "Image media token not found in hydration map: {:?}", + token + ))); + }, + } + } + + Ok(Some(urls)) +} + +fn estimate_cost_in_credits( + router_builder: &GenerateImageRequestBuilder, +) -> Result { + let mut cost_builder = router_builder.clone(); + cost_builder.provider = Provider::Artcraft; + + let cost_plan = cost_builder.build().map_err(|e| { + warn!("Failed to build image cost plan for v2 pipeline: {}", e); + AdvancedCommonWebError::from_error(e) + })?; + + Ok(cost_plan.estimate_costs().cost_in_credits.unwrap_or(0)) +} + +async fn finalize_and_generate( + draft_or_request: ImageGenerationDraftOrRequest, + server_state: &ServerState, +) -> Result { + let provider = draft_or_request.get_provider(); + let client = build_router_client(provider, server_state)?; + + let request = finalize_request(draft_or_request).await?; + + request.send_request(&client) + .await + .map_err(|err| { + warn!("v2 image generation failed: {:?}", err); + AdvancedCommonWebError::from_error(err) + }) +} + +async fn finalize_request( + draft_or_request: ImageGenerationDraftOrRequest, +) -> Result { + match draft_or_request { + ImageGenerationDraftOrRequest::Request(request) => Ok(request), + ImageGenerationDraftOrRequest::Draft(draft) => { + draft.finalize(ImageGenerationDraftContext::default()) + .await + .map_err(|err| { + warn!("Failed to finalize image v2 draft: {:?}", err); + AdvancedCommonWebError::from_error(err) + }) + } + } +} + +fn build_router_client( + provider: Provider, + server_state: &ServerState, +) -> Result { + match provider { + Provider::Fal => { + let fal_client = RouterFalWebhookOptionalClient::new_with_webhook( + server_state.fal.api_key.clone(), + server_state.fal.webhook_url.clone(), + ); + Ok(RouterClient::FalWebhookOptional(fal_client)) + }, + other => { + Err(AdvancedCommonWebError::server_error_with_message( + &format!("Unsupported provider for image v2 generation: {:?}", other), + )) + }, + } +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs index 8a937aea11..4c682a8171 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_dev.rs @@ -17,7 +17,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs index 80fb1b96cb..e362f55e8b 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_1_schnell.rs @@ -18,7 +18,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs index be69f5e709..8687850e85 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1.rs @@ -17,7 +17,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs index 666f7974e9..e8a621080e 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_flux_pro_1p1_ultra.rs @@ -17,7 +17,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs index bc9875d44e..4362e056ab 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1.rs @@ -31,7 +31,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs index 6b8b255caa..d924fe0b35 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_1p5.rs @@ -23,7 +23,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs index 68dff3f116..5a0a5554ab 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_gpt_image_2.rs @@ -27,7 +27,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs index fef4177b39..1053a64782 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana.rs @@ -21,7 +21,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs index 473e0839b4..5e05165aa9 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_2.rs @@ -24,7 +24,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs index 9fa822c33e..3efe8e8b63 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_nano_banana_pro.rs @@ -23,7 +23,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs index 6f5b7f056a..cd39101002 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4.rs @@ -21,7 +21,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs index d03f0a4a0b..9e58784155 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_4p5.rs @@ -20,7 +20,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, }; diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs index d92ec7bb03..8018412b66 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/omni_gen/generate/image/tests/test_seedream_5_lite.rs @@ -21,7 +21,7 @@ mod tests { use tokens::tokens::media_files::MediaFileToken; use url::Url; - use crate::http_server::endpoints::omni_gen::generate::image::distill_image_request::{ + use crate::http_server::endpoints::omni_gen::generate::image::pipeline_v1::distill_image_request::{ distill_image_request, DistilledImageRequest, };