From cd76b7c87fcadac1fac05d6ed774bf129f0d53ec Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Wed, 4 Mar 2026 05:35:25 -0500
Subject: [PATCH] fix pricing data handler

---
 .../src/common/generation/common_image_model.rs  | 16 ++++++++++++++++
 .../image/estimate_image_cost_handler.rs         |  8 ++++++++
 2 files changed, 24 insertions(+)

diff --git a/crates/schema/public/enums/src/common/generation/common_image_model.rs b/crates/schema/public/enums/src/common/generation/common_image_model.rs
index 599c131406..a9aa9345db 100644
--- a/crates/schema/public/enums/src/common/generation/common_image_model.rs
+++ b/crates/schema/public/enums/src/common/generation/common_image_model.rs
@@ -5,6 +5,22 @@ use utoipa::ToSchema;
 #[derive(Copy, Clone, Debug, Serialize, Deserialize, ToSchema)]
 #[serde(rename_all = "snake_case")]
 pub enum CommonImageModel {
+  #[serde(rename = "flux_1_dev")]
+  Flux1Dev,
+  #[serde(rename = "flux_1_schnell")]
+  Flux1Schnell,
+  #[serde(rename = "flux_pro_1p1")]
+  FluxPro11,
+  #[serde(rename = "flux_pro_1p1_ultra")]
+  FluxPro11Ultra,
+  #[serde(rename = "gpt_image_1p5")]
+  GptImage1p5,
+  #[serde(rename = "nano_banana")]
+  NanaBanana,
   #[serde(rename = "nano_banana_pro")]
   NanaBananaPro,
+  #[serde(rename = "seedream_4")]
+  Seedream4,
+  #[serde(rename = "seedream_4p5")]
+  Seedream4p5,
 }
diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs
index f54ce474ff..bc68ae0057 100644
--- a/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs
+++ b/crates/service/web/storyteller_web/src/http_server/endpoints/generate/cost_estimate/image/estimate_image_cost_handler.rs
@@ -121,7 +121,15 @@ fn map_provider(
 
 fn map_image_model(model: CommonImageModel) -> Result<RouterImageModel, HandlerError> {
   let router_model = match model {
+    CommonImageModel::Flux1Dev => RouterImageModel::Flux1Dev,
+    CommonImageModel::Flux1Schnell => RouterImageModel::Flux1Schnell,
+    CommonImageModel::FluxPro11 => RouterImageModel::FluxPro11,
+    CommonImageModel::FluxPro11Ultra => RouterImageModel::FluxPro11Ultra,
+    CommonImageModel::GptImage1p5 => RouterImageModel::GptImage1p5,
+    CommonImageModel::NanaBanana => RouterImageModel::NanaBanana,
     CommonImageModel::NanaBananaPro => RouterImageModel::NanaBananaPro,
+    CommonImageModel::Seedream4 => RouterImageModel::Seedream4,
+    CommonImageModel::Seedream4p5 => RouterImageModel::Seedream4p5,
   };
   Ok(router_model)
 }