From cfb4f804902045fda80049e5a140c7acb705b943 Mon Sep 17 00:00:00 2001
From: Brandon Thomas <echelon@gmail.com>
Date: Mon, 26 May 2025 19:59:29 -0400
Subject: [PATCH] upload images from recraft etc.

---
 .../fal_client/src/model/fal_endpoint.rs      | 12 ++++++++
 .../src/utils/queue_status_checker.rs         | 28 +++++++++++++++++++
 2 files changed, 40 insertions(+)

diff --git a/crates/lib/api_clients/fal_client/src/model/fal_endpoint.rs b/crates/lib/api_clients/fal_client/src/model/fal_endpoint.rs
index 170f6fe1ce..1a3b7334a2 100644
--- a/crates/lib/api_clients/fal_client/src/model/fal_endpoint.rs
+++ b/crates/lib/api_clients/fal_client/src/model/fal_endpoint.rs
@@ -1,34 +1,46 @@
 use fal::queue::Queue;
 use serde::de::DeserializeOwned;
 
+// TODO(bt,2025-05-26): A macro to define these + tests.
+
+const FLUX_PRO_ULTRA_TEXT_TO_IMAGE: &str = "fal-ai/flux-pro/v1.1-ultra";
+
 const HUNYUAN_3D_2_BASE: &str = "fal-ai/hunyuan3d/v2";
 
 const KLING16_IMAGE_TO_VIDEO: &str = "fal-ai/kling-video/v1.6/pro/image-to-video";
 
 const MINIMAX_IMAGE_TO_VIDEO: &str = "fal-ai/minimax/image-to-video";
 
+const RECRAFT_V3_TEXT_TO_IMAGE: &str = "fal-ai/recraft-v3";
+
 #[derive(Clone, Debug)]
 pub enum FalEndpoint {
+  FluxProUltraTextToImage,
   Hunyuan3d2Base,
   Kling16ImageToVideo,
   Minimax01ImageToVideo,
+  RecraftV3TextToImage,
   Other(String),
 }
 
 impl FalEndpoint {
   pub fn from_queue_response<R: DeserializeOwned>(queue: &Queue<R>) -> Self {
     match queue.endpoint.as_str() {
+      FLUX_PRO_ULTRA_TEXT_TO_IMAGE => FalEndpoint::FluxProUltraTextToImage,
       HUNYUAN_3D_2_BASE => FalEndpoint::Hunyuan3d2Base,
       KLING16_IMAGE_TO_VIDEO => FalEndpoint::Kling16ImageToVideo,
       MINIMAX_IMAGE_TO_VIDEO => FalEndpoint::Minimax01ImageToVideo,
+      RECRAFT_V3_TEXT_TO_IMAGE => FalEndpoint::RecraftV3TextToImage,
       _ => FalEndpoint::Other(queue.endpoint.clone()),
     }
   }
   pub fn url(&self) -> &str {
     match self {
+      Self::FluxProUltraTextToImage => FLUX_PRO_ULTRA_TEXT_TO_IMAGE,
       Self::Hunyuan3d2Base => HUNYUAN_3D_2_BASE,
       Self::Kling16ImageToVideo => KLING16_IMAGE_TO_VIDEO,
       Self::Minimax01ImageToVideo => MINIMAX_IMAGE_TO_VIDEO,
+      Self::RecraftV3TextToImage => RECRAFT_V3_TEXT_TO_IMAGE,
       Self::Other(url) => url,
     }
   }
diff --git a/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs
index 52efb621ff..76df7e3078 100644
--- a/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs
+++ b/crates/lib/api_clients/fal_client/src/utils/queue_status_checker.rs
@@ -2,9 +2,12 @@ use crate::creds::fal_api_key::FalApiKey;
 use crate::error::fal_error_plus::FalErrorPlus;
 use crate::model::enqueued_request::EnqueuedRequest;
 use crate::model::fal_endpoint::FalEndpoint;
+use anyhow::anyhow;
+use fal::endpoints::fal_ai::flux_pro::v1_1_ultra::Output;
 use fal::endpoints::fal_ai::hunyuan3d::v2::ObjectOutput;
 use fal::endpoints::fal_ai::kling_video::v1_6::pro::effects::I2VOutput;
 use fal::endpoints::fal_ai::minimax::image_to_video::VideoOutput;
+use fal::endpoints::fal_ai::recraft_v3::TextToImageOutput;
 use fal::prelude::QueueResponse;
 use fal::queue::{Queue, QueueStatus};
 use reqwest::Client;
@@ -21,6 +24,11 @@ impl QueueStatusChecker {
 
   pub async fn check_status(&self, request: &EnqueuedRequest) -> Result<QueueStatus, FalErrorPlus> {
     match &request.fal_endpoint {
+      FalEndpoint::FluxProUltraTextToImage => {
+        let queue = self.make_queue::<Output>(&request);
+        let status = queue.status(false).await?;
+        Ok(status)
+      }
       FalEndpoint::Hunyuan3d2Base => {
         let queue = self.make_queue::<ObjectOutput>(&request);
         let status = queue.status(false).await?;
@@ -36,6 +44,11 @@ impl QueueStatusChecker {
         let status = queue.status(false).await?;
         Ok(status)
       }
+      FalEndpoint::RecraftV3TextToImage => {
+        let queue = self.make_queue::<TextToImageOutput>(&request);
+        let status = queue.status(false).await?;
+        Ok(status)
+      }
       FalEndpoint::Other(endpoint) => {
         Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint)))
       }
@@ -44,6 +57,14 @@ impl QueueStatusChecker {
   
   pub async fn get_download_url(&self, request: &EnqueuedRequest) -> Result<String, FalErrorPlus> {
     match &request.fal_endpoint {
+      FalEndpoint::FluxProUltraTextToImage => {
+        let queue = self.make_queue::<Output>(&request);
+        let response = queue.response().await?;
+        // TODO(bt,2025-05-26): Handle multiple images
+        Ok(response.images.first()
+            .ok_or_else(|| FalErrorPlus::AnyhowError(anyhow!("No images returned")))?
+            .url.clone())
+      }
       FalEndpoint::Hunyuan3d2Base => {
         let queue = self.make_queue::<ObjectOutput>(&request);
         let response = queue.response().await?;
@@ -59,6 +80,13 @@ impl QueueStatusChecker {
         let response = queue.response().await?;
         Ok(response.video.url)
       }
+      FalEndpoint::RecraftV3TextToImage => {
+        let queue = self.make_queue::<TextToImageOutput>(&request);
+        let response = queue.response().await?;
+        Ok(response.images.first()
+            .ok_or_else(|| FalErrorPlus::AnyhowError(anyhow!("No images returned")))?
+            .url.clone())
+      }
       FalEndpoint::Other(endpoint) => {
         Err(FalErrorPlus::UnhandledEndpoint(format!("Unsupported endpoint: {}", endpoint)))
       }