seedance omni execution wip

2026-10-27 03:27:04 +00:00 · 2026-04-14 11:40:36 -04:00
parent 7e36149d13
commit 93bec65601
7 changed files with 426 additions and 68 deletions
@@ -54,8 +54,11 @@ pub struct DistilledVideoRequest {
  /// Cost estimate as computed by the Artcraft provider — this is what we bill on.
  pub cost: VideoGenerationCostEstimate,

-  /// Execution plan as computed by the Fal provider — what we hand to the router.
+  /// Execution plan as computed by the execution provider — what we hand to the router.
  plan: VideoGenerationPlan<'static>,
+
+  /// The provider used for execution (Fal, Seedance2Pro, etc.).
+  pub execution_provider: Provider,
 }

 impl DistilledVideoRequest {
@@ -76,6 +79,7 @@ impl DistilledVideoRequest {
 pub fn distill_video_request(
  request: &OmniGenVideoCostAndGenerateRequest,
  media_file_hydration_map: Option<&HashMap<MediaFileToken, Url>>,
+  execution_provider: Provider,
 ) -> Result<DistilledVideoRequest, AdvancedCommonWebError> {
  // 1. Convert the raw API request into a router request.
  let initial = hydrate_to_router_request(request)?;
@@ -155,7 +159,7 @@ pub fn distill_video_request(

  let request_static: GenerateVideoRequest<'static> = GenerateVideoRequest {
    model: initial.model,
-    provider: Provider::Fal,
+    provider: execution_provider,
    prompt: prompt_static,
    negative_prompt: negative_prompt_static,
    start_frame: start_frame_static,
@@ -210,6 +214,7 @@ pub fn distill_video_request(
    request: request_static,
    cost,
    plan,
+    execution_provider,
  })
 }

@@ -0,0 +1,305 @@
+//! Execution helpers for the omni-gen video generation handler.
+//!
+//! Routes generation to the appropriate provider (Fal or Seedance2Pro/Kinovi)
+//! based on the distilled request's `execution_provider`.
+
+use std::collections::HashMap;
+
+use log::{info, warn};
+use url::Url;
+
+use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
+use artcraft_router::api::provider::Provider;
+use artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse;
+use enums::common::generation::common_generation_mode::CommonGenerationMode;
+use seedance2pro_client::creds::seedance2pro_session::Seedance2ProSession;
+use seedance2pro_client::requests::generate_video::generate_video::{
+  generate_video, BatchCount, GenerateVideoArgs, ModelType, Resolution,
+};
+use seedance2pro_client::requests::prepare_file_upload::prepare_file_upload::{
+  prepare_file_upload, PrepareFileUploadArgs,
+};
+use seedance2pro_client::requests::upload_file::upload_file::{upload_file, UploadFileArgs};
+use tokens::tokens::media_files::MediaFileToken;
+use url_utils::extension::extract_extension_from_url::{
+  extract_extension_from_url, ExtractExtensions,
+};
+
+use crate::http_server::common_responses::advanced_common_web_error::AdvancedCommonWebError;
+use crate::state::server_state::ServerState;
+use crate::util::http_download_url_to_bytes::http_download_url_to_bytes;
+
+use super::distill_video_request::DistilledVideoRequest;
+
+/// Result of a successful generation, regardless of provider.
+pub struct GenerationResult {
+  /// The external job ID used to track the generation (first order_id for Seedance).
+  pub external_job_id: String,
+
+  /// Whether this is a Seedance2Pro generation (changes DB insertion path).
+  pub is_seedance2pro: bool,
+
+  /// For Seedance2Pro batch jobs, the list of all order IDs.
+  pub maybe_seedance_order_ids: Option<Vec<String>>,
+
+  /// The generation mode (Text, Keyframe, Reference).
+  pub generation_mode: CommonGenerationMode,
+}
+
+/// Execute video generation via the appropriate provider.
+pub async fn execute_generation(
+  distilled: &DistilledVideoRequest,
+  request: &OmniGenVideoCostAndGenerateRequest,
+  server_state: &ServerState,
+  media_file_hydration_map: Option<&HashMap<MediaFileToken, Url>>,
+) -> Result<GenerationResult, AdvancedCommonWebError> {
+  match distilled.execution_provider {
+    Provider::Seedance2Pro => {
+      execute_generation_kinovi(
+        distilled,
+        request,
+        server_state,
+        media_file_hydration_map,
+      ).await
+    }
+    _ => {
+      execute_generation_fal(
+        distilled,
+        request,
+        server_state,
+      ).await
+    }
+  }
+}
+
+/// Execute generation via Fal (the existing path).
+async fn execute_generation_fal(
+  distilled: &DistilledVideoRequest,
+  request: &OmniGenVideoCostAndGenerateRequest,
+  server_state: &ServerState,
+) -> Result<GenerationResult, AdvancedCommonWebError> {
+  let fal_client = artcraft_router::client::router_fal_client::RouterFalClient::new(
+    server_state.fal.api_key.clone(),
+    server_state.fal.webhook_url.clone(),
+  );
+  let router_client = artcraft_router::client::router_client::RouterClient::Fal(fal_client);
+
+  let generation_response = distilled.plan().generate_video(&router_client)
+    .await
+    .map_err(|e| {
+      warn!("Video generation failed (Fal): {:?}", e);
+      AdvancedCommonWebError::from_error(e)
+    })?;
+
+  let external_job_id = match &generation_response {
+    GenerateVideoResponse::Artcraft(p) => p.inference_job_token.as_str().to_string(),
+    GenerateVideoResponse::Muapi(p) => p.request_id.as_str().to_string(),
+    GenerateVideoResponse::Seedance2Pro(p) => p.order_id.clone(),
+    GenerateVideoResponse::Fal(p) => p.request_id.clone().unwrap_or_default(),
+  };
+
+  let generation_mode = if request.start_frame_image_media_token.is_some() {
+    CommonGenerationMode::Keyframe
+  } else {
+    CommonGenerationMode::Text
+  };
+
+  Ok(GenerationResult {
+    external_job_id,
+    is_seedance2pro: false,
+    maybe_seedance_order_ids: None,
+    generation_mode,
+  })
+}
+
+/// Execute generation via Seedance2Pro/Kinovi.
+///
+/// Downloads all referenced media from our CDN, re-uploads to Seedance2Pro's
+/// CDN, then calls the Seedance2Pro generate endpoint.
+async fn execute_generation_kinovi(
+  distilled: &DistilledVideoRequest,
+  request: &OmniGenVideoCostAndGenerateRequest,
+  server_state: &ServerState,
+  media_file_hydration_map: Option<&HashMap<MediaFileToken, Url>>,
+) -> Result<GenerationResult, AdvancedCommonWebError> {
+  let session = Seedance2ProSession::from_cookies_string(
+    server_state.seedance2pro.cookies.clone()
+  );
+
+  let empty_map = HashMap::new();
+  let hydration_map = media_file_hydration_map.unwrap_or(&empty_map);
+
+  // Upload each media reference to Seedance2Pro CDN.
+  let start_frame_url = upload_token_to_seedance2pro(
+    &session, hydration_map, request.start_frame_image_media_token.as_ref(),
+  ).await?;
+
+  let end_frame_url = upload_token_to_seedance2pro(
+    &session, hydration_map, request.end_frame_image_media_token.as_ref(),
+  ).await?;
+
+  let reference_image_urls = upload_tokens_to_seedance2pro(
+    &session, hydration_map, request.reference_image_media_tokens.as_deref(),
+  ).await?;
+
+  let reference_video_urls = upload_tokens_to_seedance2pro(
+    &session, hydration_map, request.reference_video_media_tokens.as_deref(),
+  ).await?;
+
+  let reference_audio_urls = upload_tokens_to_seedance2pro(
+    &session, hydration_map, request.reference_audio_media_tokens.as_deref(),
+  ).await?;
+
+  // Determine generation mode.
+  let is_keyframe = request.start_frame_image_media_token.is_some()
+      || request.end_frame_image_media_token.is_some();
+  let is_reference = request.reference_image_media_tokens.as_ref().map_or(false, |t| !t.is_empty())
+      || request.reference_video_media_tokens.as_ref().map_or(false, |t| !t.is_empty())
+      || request.reference_audio_media_tokens.as_ref().map_or(false, |t| !t.is_empty());
+
+  let generation_mode = if is_keyframe {
+    CommonGenerationMode::Keyframe
+  } else if is_reference {
+    CommonGenerationMode::Reference
+  } else {
+    CommonGenerationMode::Text
+  };
+
+  // Map aspect ratio / duration / batch from the request.
+  let resolution = Resolution::Landscape16x9; // TODO: Map from request.aspect_ratio
+  let duration_seconds = request.duration_seconds.unwrap_or(5).clamp(4, 15) as u8;
+  let batch_count = match request.video_batch_count {
+    Some(2) => BatchCount::Two,
+    Some(4) => BatchCount::Four,
+    _ => BatchCount::One,
+  };
+
+  let prompt = request.prompt.clone().unwrap_or_default();
+
+  let video_gen_args = GenerateVideoArgs {
+    session: &session,
+    model_type: ModelType::Seedance2Pro,
+    prompt,
+    resolution,
+    duration_seconds,
+    batch_count,
+    start_frame_url,
+    end_frame_url,
+    reference_image_urls,
+    reference_video_urls,
+    reference_audio_urls,
+    character_ids: None, // TODO: resolve_kinovi_character_ids when character support is added
+    use_face_blur_hack: None,
+    host_override: None,
+  };
+
+  let gen_response = generate_video(video_gen_args).await
+    .map_err(|err| {
+      warn!("Seedance2Pro generate_video failed: {:?}", err);
+      AdvancedCommonWebError::from_error(err)
+    })?;
+
+  info!(
+    "Seedance2Pro generation: task_id={}, order_id={}",
+    gen_response.task_id, gen_response.order_id
+  );
+
+  let order_ids: Vec<String> = match gen_response.order_ids {
+    Some(ids) if !ids.is_empty() => ids,
+    _ => vec![gen_response.order_id.clone()],
+  };
+
+  Ok(GenerationResult {
+    external_job_id: gen_response.order_id,
+    is_seedance2pro: true,
+    maybe_seedance_order_ids: Some(order_ids),
+    generation_mode,
+  })
+}
+
+// --- Upload helpers ---
+
+/// Upload a single media token (if present) to Seedance2Pro CDN.
+async fn upload_token_to_seedance2pro(
+  session: &Seedance2ProSession,
+  hydration_map: &HashMap<MediaFileToken, Url>,
+  maybe_token: Option<&MediaFileToken>,
+) -> Result<Option<String>, AdvancedCommonWebError> {
+  let token = match maybe_token {
+    None => return Ok(None),
+    Some(t) => t,
+  };
+
+  let url = hydration_map.get(token).ok_or_else(|| {
+    AdvancedCommonWebError::BadInputWithSimpleMessage(
+      format!("Media token not found in hydration map: {:?}", token))
+  })?;
+
+  let uploaded = upload_url_to_seedance2pro(session, url).await?;
+  Ok(Some(uploaded))
+}
+
+/// Upload a list of media tokens (if present) to Seedance2Pro CDN.
+async fn upload_tokens_to_seedance2pro(
+  session: &Seedance2ProSession,
+  hydration_map: &HashMap<MediaFileToken, Url>,
+  maybe_tokens: Option<&[MediaFileToken]>,
+) -> Result<Option<Vec<String>>, AdvancedCommonWebError> {
+  let tokens = match maybe_tokens {
+    None => return Ok(None),
+    Some(tokens) if tokens.is_empty() => return Ok(None),
+    Some(tokens) => tokens,
+  };
+
+  let mut urls = Vec::with_capacity(tokens.len());
+  for token in tokens {
+    let url = hydration_map.get(token).ok_or_else(|| {
+      AdvancedCommonWebError::BadInputWithSimpleMessage(
+        format!("Media token not found in hydration map: {:?}", token))
+    })?;
+    urls.push(upload_url_to_seedance2pro(session, url).await?);
+  }
+  Ok(Some(urls))
+}
+
+/// Download from our CDN and re-upload to Seedance2Pro.
+async fn upload_url_to_seedance2pro(
+  session: &Seedance2ProSession,
+  our_cdn_url: &Url,
+) -> Result<String, AdvancedCommonWebError> {
+  let extension = extract_extension_from_url(our_cdn_url, &ExtractExtensions::All)
+      .map(|ext| ext.without_period().to_string())
+      .unwrap_or_else(|| "png".to_string());
+
+  let file_bytes = http_download_url_to_bytes(our_cdn_url.as_str())
+      .await
+      .map_err(|err| {
+        warn!("Error downloading media from CDN for Seedance2Pro upload: {:?}", err);
+        AdvancedCommonWebError::from_error(err)
+      })?
+      .to_vec();
+
+  let prepare_result = prepare_file_upload(PrepareFileUploadArgs {
+    session,
+    extension,
+    host_override: None,
+  })
+      .await
+      .map_err(|err| {
+        warn!("Error preparing seedance2pro file upload: {:?}", err);
+        AdvancedCommonWebError::from_error(err)
+      })?;
+
+  let upload_result = upload_file(UploadFileArgs {
+    upload_url: prepare_result.upload_url,
+    file_bytes,
+    host_override: None,
+  })
+      .await
+      .map_err(|err| {
+        warn!("Error uploading file to seedance2pro: {:?}", err);
+        AdvancedCommonWebError::from_error(err)
+      })?;
+
+  Ok(upload_result.public_url)
+}
@@ -1,6 +1,7 @@
 pub mod omni_gen_video_generate_handler;
 pub mod distill_video_request;
 pub(crate) mod distill_helper;
+mod execute_generation;

 #[cfg(test)]
 mod tests;
@@ -9,9 +9,10 @@ use url::Url;

 use artcraft_api_defs::omni_gen::cost_and_generate_requests::omni_gen_video_cost_and_generate_request::OmniGenVideoCostAndGenerateRequest;
 use artcraft_api_defs::omni_gen::generate_response::omni_gen_video_generate_response::OmniGenVideoGenerateResponse;
+use artcraft_router::api::provider::Provider;
+use enums::common::generation::common_video_model::CommonVideoModel;
 use enums::by_table::prompt_context_items::prompt_context_semantic_type::PromptContextSemanticType;
 use enums::by_table::prompts::prompt_type::PromptType;
-use enums::common::generation::common_generation_mode::CommonGenerationMode;
 use enums::common::generation::common_model_type::CommonModelType;
 use enums::common::generation_provider::GenerationProvider;
 use enums::common::visibility::Visibility;
@@ -21,6 +22,10 @@ use mysql_queries::queries::generic_inference::fal::insert_generic_inference_job
  insert_generic_inference_job_for_fal_queue_with_apriori_job_token,
  InsertGenericInferenceForFalWithAprioriJobTokenArgs,
 };
+use mysql_queries::queries::generic_inference::seedance2pro::insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token::{
+  insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token,
+  InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs,
+};
 use mysql_queries::queries::idepotency_tokens::insert_idempotency_token::insert_idempotency_token;
 use mysql_queries::queries::prompt_context_items::insert_batch_prompt_context_items::{
  insert_batch_prompt_context_items, InsertBatchArgs, PromptContextItem,
@@ -37,6 +42,7 @@ use crate::state::server_state::ServerState;
 use crate::util::lookup::lookup_image_urls_as_map::lookup_image_urls_as_map;

 use super::distill_video_request::distill_video_request;
+use super::execute_generation::execute_generation;

 /// Generate a video using the omni-gen unified endpoint.
 #[utoipa::path(
@@ -149,12 +155,20 @@ pub async fn omni_gen_video_generate_handler(
    }
  };

+  // ==================== DETERMINE PROVIDER ==================== //
+
+  let execution_provider = match request.model {
+    Some(CommonVideoModel::Seedance2p0) => Provider::Seedance2Pro,
+    _ => Provider::Fal,
+  };
+
  // ==================== DISTILL ==================== //

-  let distilled = distill_video_request(&request, media_file_hydration_map.as_ref())?;
+  let distilled = distill_video_request(&request, media_file_hydration_map.as_ref(), execution_provider)?;

  info!(">>> Distilled cost: {:?}", distilled.cost);
  info!(">>> Distilled plan: {:?}", distilled.plan());
+  info!(">>> Execution provider: {:?}", distilled.execution_provider);

  // ==================== BILLING ==================== //

@@ -175,33 +189,12 @@ pub async fn omni_gen_video_generate_handler(

  // ==================== EXECUTE GENERATION ==================== //

-  let fal_client = artcraft_router::client::router_fal_client::RouterFalClient::new(
-    server_state.fal.api_key.clone(),
-    server_state.fal.webhook_url.clone(),
-  );
-  let router_client = artcraft_router::client::router_client::RouterClient::Fal(fal_client);
-
-  let generation_response = distilled.plan().generate_video(&router_client)
-    .await
-    .map_err(|e| {
-      warn!("Video generation failed: {:?}", e);
-      AdvancedCommonWebError::from_error(e)
-    })?;
-
-  let external_job_id = match &generation_response {
-    artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Artcraft(p) => {
-      p.inference_job_token.as_str().to_string()
-    }
-    artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Muapi(p) => {
-      p.request_id.as_str().to_string()
-    }
-    artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Seedance2Pro(p) => {
-      p.order_id.clone()
-    }
-    artcraft_router::generate::generate_video::generate_video_response::GenerateVideoResponse::Fal(p) => {
-      p.request_id.clone().unwrap_or_default()
-    }
-  };
+  let gen_result = execute_generation(
+    &distilled,
+    &request,
+    &server_state,
+    media_file_hydration_map.as_ref(),
+  ).await?;

  // ==================== DB TRANSACTION ==================== //

@@ -217,12 +210,6 @@ pub async fn omni_gen_video_generate_handler(

  // -- Prompt --

-  let generation_mode = if request.start_frame_image_media_token.is_some() {
-    CommonGenerationMode::Keyframe
-  } else {
-    CommonGenerationMode::Text
-  };
-
  let prompt_result = insert_prompt(InsertPromptArgs {
    maybe_apriori_prompt_token: None,
    prompt_type: PromptType::ArtcraftApp,
@@ -232,7 +219,7 @@ pub async fn omni_gen_video_generate_handler(
    maybe_positive_prompt: request.prompt.as_deref(),
    maybe_negative_prompt: request.negative_prompt.as_deref(),
    maybe_other_args: None,
-    maybe_generation_mode: Some(generation_mode),
+    maybe_generation_mode: Some(gen_result.generation_mode),
    maybe_aspect_ratio: None,
    maybe_resolution: None,
    maybe_batch_count: request.video_batch_count.map(|c| c as u8),
@@ -301,31 +288,88 @@ pub async fn omni_gen_video_generate_handler(

  // -- Inference job --

-  let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(
-    InsertGenericInferenceForFalWithAprioriJobTokenArgs {
-      apriori_job_token: &apriori_job_token,
-      uuid_idempotency_token: &idempotency_token,
-      maybe_external_third_party_id: &external_job_id,
-      fal_category: FalCategory::VideoGeneration,
-      maybe_inference_args: None,
-      maybe_prompt_token: prompt_token.as_ref(),
-      maybe_creator_user_token: Some(user_token),
-      maybe_avt_token: maybe_avt_token.as_ref(),
-      creator_ip_address: &ip_address,
-      creator_set_visibility: Visibility::Public,
-      mysql_executor: &mut *transaction,
-      starting_job_status_override: None,
-      maybe_frontend_failure_category: None,
-      maybe_failure_reason: None,
-      phantom: Default::default(),
-    }
-  ).await;
+  let job_token = if gen_result.is_seedance2pro {
+    // Seedance2Pro path: insert one job per order_id (batch support).
+    let fallback_ids = vec![gen_result.external_job_id.clone()];
+    let order_ids = gen_result.maybe_seedance_order_ids
+        .as_deref()
+        .unwrap_or(&fallback_ids);

-  let job_token = match db_result {
-    Ok(token) => token,
-    Err(err) => {
-      warn!("Error inserting inference job: {:?}", err);
-      return Err(AdvancedCommonWebError::from_error(err));
+    let mut all_job_tokens: Vec<InferenceJobToken> = Vec::with_capacity(order_ids.len());
+
+    for (i, order_id) in order_ids.iter().enumerate() {
+      let job_token = if i == 0 {
+        apriori_job_token.clone()
+      } else {
+        InferenceJobToken::generate()
+      };
+
+      let idempotency_str = if i == 0 {
+        idempotency_token.clone()
+      } else {
+        format!("{}-batch-{}", idempotency_token, i)
+      };
+
+      let db_result = insert_generic_inference_job_for_seedance2pro_queue_with_apriori_job_token(
+        InsertGenericInferenceForSeedance2ProWithAprioriJobTokenArgs {
+          apriori_job_token: &job_token,
+          uuid_idempotency_token: &idempotency_str,
+          maybe_external_third_party_id: order_id,
+          maybe_inference_args: None,
+          maybe_prompt_token: prompt_token.as_ref(),
+          maybe_wallet_ledger_entry_token: None,
+          maybe_creator_user_token: Some(user_token),
+          maybe_avt_token: maybe_avt_token.as_ref(),
+          creator_ip_address: &ip_address,
+          creator_set_visibility: Visibility::Public,
+          mysql_executor: &mut *transaction,
+          phantom: Default::default(),
+        }
+      ).await;
+
+      match db_result {
+        Ok(token) => all_job_tokens.push(token),
+        Err(err) => {
+          warn!("Error inserting seedance2pro inference job (order_id={}): {:?}", order_id, err);
+          if i == 0 {
+            return Err(AdvancedCommonWebError::from_error(err));
+          }
+        }
+      }
+    }
+
+    all_job_tokens.first().cloned().ok_or_else(|| {
+      error!("No inference job token was created");
+      AdvancedCommonWebError::server_error_with_message("No inference job token was created")
+    })?
+  } else {
+    // Fal / other providers path.
+    let db_result = insert_generic_inference_job_for_fal_queue_with_apriori_job_token(
+      InsertGenericInferenceForFalWithAprioriJobTokenArgs {
+        apriori_job_token: &apriori_job_token,
+        uuid_idempotency_token: &idempotency_token,
+        maybe_external_third_party_id: &gen_result.external_job_id,
+        fal_category: FalCategory::VideoGeneration,
+        maybe_inference_args: None,
+        maybe_prompt_token: prompt_token.as_ref(),
+        maybe_creator_user_token: Some(user_token),
+        maybe_avt_token: maybe_avt_token.as_ref(),
+        creator_ip_address: &ip_address,
+        creator_set_visibility: Visibility::Public,
+        mysql_executor: &mut *transaction,
+        starting_job_status_override: None,
+        maybe_frontend_failure_category: None,
+        maybe_failure_reason: None,
+        phantom: Default::default(),
+      }
+    ).await;
+
+    match db_result {
+      Ok(token) => token,
+      Err(err) => {
+        warn!("Error inserting inference job: {:?}", err);
+        return Err(AdvancedCommonWebError::from_error(err));
+      }
    }
  };

@@ -19,6 +19,7 @@ mod tests {
  use tokens::tokens::media_files::MediaFileToken;
  use url::Url;

+  use artcraft_router::api::provider::Provider;
  use crate::http_server::endpoints::omni_gen::generate::video::distill_video_request::{
    distill_video_request, DistilledVideoRequest,
  };
@@ -75,7 +76,7 @@ mod tests {

  fn distill_text(request: &OmniGenVideoCostAndGenerateRequest) -> DistilledVideoRequest {
    let empty: HashMap<MediaFileToken, Url> = HashMap::new();
-    distill_video_request(request, Some(&empty))
+    distill_video_request(request, Some(&empty), Provider::Fal)
      .expect("distill_video_request should succeed for Seedance 1.5 Pro (text)")
  }

@@ -83,7 +84,7 @@ mod tests {
    request: &OmniGenVideoCostAndGenerateRequest,
    hydration: &HashMap<MediaFileToken, Url>,
  ) -> DistilledVideoRequest {
-    distill_video_request(request, Some(hydration))
+    distill_video_request(request, Some(hydration), Provider::Fal)
      .expect("distill_video_request should succeed for Seedance 1.5 Pro (image)")
  }

@@ -20,6 +20,7 @@ mod tests {
  use tokens::tokens::media_files::MediaFileToken;
  use url::Url;

+  use artcraft_router::api::provider::Provider;
  use crate::http_server::endpoints::omni_gen::generate::video::distill_video_request::{
    distill_video_request, DistilledVideoRequest,
  };
@@ -62,7 +63,7 @@ mod tests {

  fn distill_text(request: &OmniGenVideoCostAndGenerateRequest) -> DistilledVideoRequest {
    let empty: HashMap<MediaFileToken, Url> = HashMap::new();
-    distill_video_request(request, Some(&empty))
+    distill_video_request(request, Some(&empty), Provider::Fal)
      .expect("distill_video_request should succeed for Veo 2 (text)")
  }

@@ -70,7 +71,7 @@ mod tests {
    request: &OmniGenVideoCostAndGenerateRequest,
    hydration: &HashMap<MediaFileToken, Url>,
  ) -> DistilledVideoRequest {
-    distill_video_request(request, Some(hydration))
+    distill_video_request(request, Some(hydration), Provider::Fal)
      .expect("distill_video_request should succeed for Veo 2 (image)")
  }

@@ -24,6 +24,7 @@ mod tests {
  use tokens::tokens::media_files::MediaFileToken;
  use url::Url;

+  use artcraft_router::api::provider::Provider;
  use crate::http_server::endpoints::omni_gen::generate::video::distill_video_request::{
    distill_video_request, DistilledVideoRequest,
  };
@@ -68,7 +69,7 @@ mod tests {

  fn distill_text(request: &OmniGenVideoCostAndGenerateRequest) -> DistilledVideoRequest {
    let empty: HashMap<MediaFileToken, Url> = HashMap::new();
-    distill_video_request(request, Some(&empty))
+    distill_video_request(request, Some(&empty), Provider::Fal)
      .expect("distill should succeed for Veo 3 (text)")
  }

@@ -76,7 +77,7 @@ mod tests {
    request: &OmniGenVideoCostAndGenerateRequest,
    hydration: &HashMap<MediaFileToken, Url>,
  ) -> DistilledVideoRequest {
-    distill_video_request(request, Some(hydration))
+    distill_video_request(request, Some(hydration), Provider::Fal)
      .expect("distill should succeed for Veo 3 (image)")
  }