split image gen

2026-11-02 08:28:18 +00:00 · 2026-05-15 01:03:22 -04:00
1 parent 9c0033c673
commit 03354b73d9
@@ -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.
@@ -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)]
@@ -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,
@@ -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,
+}
@@ -0,0 +1,3 @@
+pub(crate) mod distill_helper;
+pub(crate) mod distill_image_request;
+pub mod run_pipeline_v1;
@@ -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<sqlx::MySql>,
+  pub user_token: &'a UserToken,
+  pub media_file_hydration_map: &'a Option<HashMap<MediaFileToken, Url>>,
+}
+
+pub async fn run_pipeline_v1(
+  args: RunPipelineV1Args<'_>,
+) -> Result<ImagePipelineResult, AdvancedCommonWebError> {
+  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,
+  })
+}
@@ -0,0 +1 @@
+pub mod run_pipeline_v2;
@@ -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<sqlx::MySql>,
+  pub user_token: &'a UserToken,
+  pub media_file_hydration_map: &'a Option<HashMap<MediaFileToken, Url>>,
+}
+
+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<ImagePipelineResult, AdvancedCommonWebError> {
+  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<ImageGenerationDraftOrRequest, AdvancedCommonWebError> {
+  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<MediaFileToken, Url>>,
+) -> Result<GenerateImageRequestBuilder, AdvancedCommonWebError> {
+  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<MediaFileToken, Url>>,
+) -> Result<Option<Vec<String>>, 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<String> = 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<u64, AdvancedCommonWebError> {
+  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<GenerateImageResponse, AdvancedCommonWebError> {
+  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<ImageGenerationRequest, AdvancedCommonWebError> {
+  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<RouterClient, AdvancedCommonWebError> {
+  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),
+      ))
+    },
+  }
+}
@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };

@@ -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,
  };