image-gen: compiles and maybe even runs

Cargo.lock

@@ -5793,6 +5793,7 @@ dependencies = [
  "my-workspace-hack",
  "mysql_queries",
  "once_cell",
+ "openai_sora_client",
  "opentelemetry",
  "opentelemetry-otlp",
  "opentelemetry_sdk",

crates/_types/enums/src/by_table/media_files/media_file_origin_model_type.rs

@@ -57,6 +57,9 @@ pub enum MediaFileOriginModelType {
   #[serde(rename = "studio")]
   StorytellerStudio,
 
+  #[serde(rename = "studio_ig")]
+  StorytellerStudioImageGen,
+
   #[serde(rename = "vst")]
   VideoStyleTransfer,
 
@@ -93,6 +96,7 @@ impl MediaFileOriginModelType {
       Self::MocapNet => "mocap_net",
       Self::StyleTTS2 => "styletts2",
       Self::StableDiffusion15 => "stable_diffusion_1_5",
+      Self::StorytellerStudioImageGen => "studio_ig",
       Self::GptSovits => "gpt_sovits",
       Self::StorytellerStudio => "studio",
       Self::VideoStyleTransfer => "vst",
@@ -115,6 +119,7 @@ impl MediaFileOriginModelType {
       "mocap_net" => Ok(Self::MocapNet),
       "styletts2" => Ok(Self::StyleTTS2),
       "stable_diffusion_1_5" => Ok(Self::StableDiffusion15),
+      "studio_ig" => Ok(Self::StorytellerStudioImageGen),
       "gpt_sovits" => Ok(Self::GptSovits),
       "studio" => Ok(Self::StorytellerStudio),
       "vst" => Ok(Self::VideoStyleTransfer),
@@ -142,6 +147,7 @@ impl MediaFileOriginModelType {
       Self::StableDiffusion15,
       Self::GptSovits,
       Self::StorytellerStudio,
+      Self::StorytellerStudioImageGen,
       Self::VideoStyleTransfer,
       Self::ComfyUi,
       Self::VallEX,

crates/lib/api_clients/openai_sora_client/src/image_gen/image_gen_status.rs

@@ -0,0 +1,349 @@
+use std::{fs::File, io::Write, path::Path};
+
+use crate::credentials::SoraCredentials;
+use errors::AnyhowResult;
+use reqwest::Url;
+use serde_derive::Deserialize;
+
+const SORA_STATUS_URL: &str = "https://sora.com/backend/video_gen";
+
+/// This user agent is tied to the sentinel generation. If we need to change it, we may need to change sentinel generation too.
+const USER_AGENT: &str = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/127.0.0.0 Safari/537.36";
+
+#[derive(Debug, Clone, PartialEq, Eq, Deserialize)]
+pub enum TaskStatus {
+  Queued,
+  Running,
+  Succeeded,
+  Failed,
+  Unknown(String),
+}
+
+impl TaskStatus {
+  pub fn from_str(s: &str) -> Self {
+    match s {
+      "queued" => TaskStatus::Queued,
+      "running" => TaskStatus::Running,
+      "succeeded" => TaskStatus::Succeeded,
+      "failed" => TaskStatus::Failed,
+      _ => TaskStatus::Unknown(s.to_string()),
+    }
+  }
+}
+
+#[derive(Deserialize, Clone)]
+#[serde(rename_all = "snake_case")]
+pub struct VideoGenStatusResponse {
+  pub task_responses: Vec<TaskResponse>,
+  pub last_id: String,
+  pub has_more: bool,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct TaskResponse {
+  pub id: String,
+  pub user: String,
+  pub created_at: String,
+  pub status: String,
+  pub progress_pct: Option<f64>,
+  pub progress_pos_in_queue: Option<i32>,
+  pub estimated_queue_wait_time: Option<i32>,
+  pub queue_status_message: Option<String>,
+  pub priority: i32,
+  pub r#type: String,
+  pub prompt: String,
+  pub n_variants: i32,
+  pub n_frames: i32,
+  pub height: i32,
+  pub width: i32,
+  pub model: Option<String>,
+  pub operation: String,
+  pub inpaint_items: Option<Vec<InpaintItem>>,
+  pub preset_id: Option<String>,
+  pub caption: Option<String>,
+  pub actions: Option<serde_json::Value>,
+  pub interpolation: Option<serde_json::Value>,
+  pub sdedit: Option<serde_json::Value>,
+  pub remix_config: Option<serde_json::Value>,
+  pub quality: Option<String>,
+  pub generations: Vec<Generation>,
+  pub num_unsafe_generations: i32,
+  pub title: String,
+  pub moderation_result: ModerationResult,
+  pub failure_reason: Option<String>,
+  pub needs_user_review: bool,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct Generation {
+  pub id: String,
+  pub task_id: String,
+  pub created_at: String,
+  pub deleted_at: Option<String>,
+  pub url: String,
+  pub seed: i64,
+  pub can_download: bool,
+  pub download_status: String,
+  pub is_favorite: bool,
+  pub is_liked: bool,
+  pub is_public: bool,
+  pub is_archived: bool,
+  pub is_featured: bool,
+  pub has_feedback: bool,
+  pub like_count: i32,
+  pub cloudflare_metadata: Option<serde_json::Value>,
+  pub cf_thumbnail_url: Option<String>,
+  pub encodings: Encodings,
+  pub width: i32,
+  pub height: i32,
+  pub n_frames: i32,
+  pub prompt: String,
+  pub title: String,
+  pub actions: Option<serde_json::Value>,
+  pub inpaint_items: Option<Vec<InpaintItem>>,
+  pub interpolation: Option<serde_json::Value>,
+  pub sdedit: Option<serde_json::Value>,
+  pub operation: String,
+  pub model: Option<String>,
+  pub preset_id: Option<String>,
+  pub user: User,
+  pub moderation_result: ModerationResult,
+  pub paragen_status: Option<serde_json::Value>,
+  pub task_type: String,
+  pub remix_config: Option<serde_json::Value>,
+  pub quality: Option<String>,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct Encodings {
+  pub source: Source,
+  pub md: Option<serde_json::Value>,
+  pub ld: Option<serde_json::Value>,
+  pub thumbnail: Thumbnail,
+  pub spritesheet: Option<serde_json::Value>,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct Source {
+  pub path: String,
+  pub size: Option<i32>,
+  pub width: Option<i32>,
+  pub height: Option<i32>,
+  pub duration_secs: Option<f64>,
+  pub ssim: Option<f64>,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct Thumbnail {
+  pub path: String,
+  pub size: Option<i32>,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct User {
+  pub id: String,
+  pub username: String,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct InpaintItem {
+  pub crop_bounds: Option<String>,
+  pub r#type: String,
+  pub preset_id: Option<String>,
+  pub generation_id: Option<String>,
+  pub upload_media_id: String,
+  pub frame_index: i32,
+  pub source_start_frame: i32,
+}
+
+#[derive(Deserialize, Clone, Debug)]
+#[serde(rename_all = "snake_case")]
+pub struct ModerationResult {
+  pub r#type: String,
+  pub results_by_frame_index: serde_json::Value,
+  pub code: Option<String>,
+  pub is_output_rejection: bool,
+  pub task_id: String,
+}
+
+pub struct StatusRequest {
+  pub limit: Option<u32>,
+  pub before: Option<String>,
+}
+
+impl StatusRequest {
+  pub fn new(limit: Option<u32>, before: Option<String>) -> Self {
+    let limit = limit.unwrap_or(100);
+    Self { limit: Some(limit), before }
+  }
+}
+
+/// Gets the status of image generation tasks from Sora API
+pub async fn get_image_gen_status(status_request: &StatusRequest, credentials: &SoraCredentials) -> AnyhowResult<VideoGenStatusResponse> {
+  let client = reqwest::Client::new();
+
+  let mut url = reqwest::Url::parse(SORA_STATUS_URL)?;
+
+  // Add query parameters
+  if let Some(limit) = status_request.limit {
+    url.query_pairs_mut().append_pair("limit", &limit.to_string());
+  }
+
+  if let Some(before) = &status_request.before {
+    url.query_pairs_mut().append_pair("before", &before);
+  }
+
+  let http_request = client.get(url).header("User-Agent", USER_AGENT).header("Cookie", &credentials.cookie).header("Authorization", credentials.authorization_header_value()).header("Content-Type", "application/json");
+
+  let http_request = http_request.build()?;
+
+  let response = client.execute(http_request).await?;
+
+  let status = response.status();
+
+  let response_body = &response.text().await?;
+  println!("response_body: {}", response_body);
+
+  if status != reqwest::StatusCode::OK {
+    return Err(anyhow::anyhow!("The status request failed; status = {:?}, message = {:?}", status, response_body));
+  }
+
+  let response: VideoGenStatusResponse = serde_json::from_str(response_body)?;
+
+  Ok(response)
+}
+
+pub async fn wait_for_image_gen_status(task_id: &String, credentials: &SoraCredentials, retry_limit: Option<u32>) -> AnyhowResult<TaskResponse> {
+  let status_request = StatusRequest {
+    limit: None,
+    before: None,
+    // before: Some(task_id.clone()),
+  };
+  let retry_limit = retry_limit.unwrap_or(10);
+
+  for _ in 0..retry_limit {
+    tokio::time::sleep(std::time::Duration::from_secs(10)).await;
+    let status_response = get_image_gen_status(&status_request, credentials).await?;
+    match find_task_response_by_id(&status_response, task_id.clone()) {
+      Some(task_response) => match TaskStatus::from_str(&task_response.status) {
+        TaskStatus::Succeeded => {
+          return Ok(task_response.clone());
+        },
+        TaskStatus::Failed => {
+          return Err(anyhow::anyhow!("Task failed"));
+        },
+        TaskStatus::Unknown(status) => {
+          println!("Unknown task status: {}", status);
+          continue;
+        },
+        TaskStatus::Queued => {
+          println!("Task is queued");
+          continue;
+        },
+        TaskStatus::Running => {
+          println!("Task is running");
+          continue;
+        },
+      },
+      None => {
+        return Err(anyhow::anyhow!("Task not found"));
+      },
+    }
+  }
+
+  Err(anyhow::anyhow!("Task not found"))
+}
+
+pub async fn save_generations_to_dir(generations: &[Generation], dir: &str) -> AnyhowResult<()> {
+  let dir = Path::new(dir);
+  if !dir.exists() {
+    std::fs::create_dir_all(dir)?;
+  }
+  for generation in generations {
+    let response = reqwest::get(generation.url.clone()).await?;
+    let image_bytes = response.bytes().await?;
+
+    let url = Url::parse(&generation.url)?;
+    let ext = url.path().split(".").last().unwrap_or("png");
+    let path = Path::new(dir).join(format!("{}.{}", generation.id, ext));
+
+    let mut file = File::create(path)?;
+    file.write_all(&image_bytes)?;
+
+    let thumbnail_url = generation.encodings.thumbnail.path.clone();
+    let thumbnail_response = reqwest::get(thumbnail_url).await?;
+    let thumbnail_bytes = thumbnail_response.bytes().await?;
+    let thumbnail_path = Path::new(dir).join(format!("{}_thumbnail.{}", generation.id, ext));
+    let mut thumbnail_file = File::create(thumbnail_path)?;
+    thumbnail_file.write_all(&thumbnail_bytes)?;
+  }
+
+  Ok(())
+}
+
+fn find_task_response_by_id(status_response: &VideoGenStatusResponse, task_id: String) -> Option<&TaskResponse> {
+  status_response.task_responses.iter().find(|task| task.id == task_id)
+}
+
+#[cfg(test)]
+mod tests {
+  use crate::credentials::SoraCredentials;
+  use crate::image_gen::image_gen_status::{get_image_gen_status, save_generations_to_dir, wait_for_image_gen_status, StatusRequest, VideoGenStatusResponse};
+  use errors::AnyhowResult;
+  use std::fs::read_to_string;
+  use testing::test_file_path::test_file_path;
+
+  #[ignore]
+  #[tokio::test]
+  pub async fn test_deserialize_status_response() -> AnyhowResult<()> {
+    let response = read_to_string(test_file_path("test_data/image_gen/test_status_response_long.json")?)?;
+    let _response: VideoGenStatusResponse = serde_json::from_str(&response)?;
+    Ok(())
+  }
+
+  #[ignore] // You can manually run "ignore" tests in the IDE, but they won't run in CI.
+  #[tokio::test]
+  pub async fn manual_test() -> AnyhowResult<()> {
+    let sentinel = read_to_string(test_file_path("test_data/temp/sentinel.txt")?)?;
+    let sentinel = sentinel.trim().to_string();
+
+    let cookie = read_to_string(test_file_path("test_data/temp/cookie.txt")?)?;
+    let cookie = cookie.trim().to_string();
+
+    let bearer = read_to_string(test_file_path("test_data/temp/bearer.txt")?)?;
+    let bearer = bearer.trim().to_string();
+
+    let creds = SoraCredentials { bearer_token: bearer, cookie, sentinel: Some(sentinel) };
+    // Get the task status for a specific task
+    // let response = get_image_gen_status(&StatusRequest::new(None, Some("task_01jr9yvpfyetx9r7qvvx38scna".to_string())), &creds).await?;
+
+    let response = get_image_gen_status(&StatusRequest::new(Some(50), None), &creds).await?;
+    println!("task_id: {}", response.task_responses[0].id);
+
+    let task_id = response.task_responses[0].id.clone();
+    let task_response = wait_for_image_gen_status(
+      &task_id,
+      &creds,
+      Some(10)
+    ).await?;
+
+    assert!(task_response.status == "succeeded");
+    Ok(())
+  }
+
+  #[ignore]
+  #[tokio::test]
+  pub async fn test_save_generations_to_dir() -> AnyhowResult<()> {
+    let response = read_to_string(test_file_path("test_data/image_gen/test_status_response_long.json")?)?;
+    let response: VideoGenStatusResponse = serde_json::from_str(&response)?;
+    save_generations_to_dir(&response.task_responses[0].generations, test_file_path("test_data/image_gen/generations")?.to_str().unwrap()).await?;
+    Ok(())
+  }
+}

crates/lib/api_clients/openai_sora_client/src/image_gen/mod.rs

@@ -2,3 +2,5 @@ pub (crate) mod image_gen_http_request;
 pub mod common;
 pub mod sora_image_gen_remix;
 pub mod sora_image_gen_simple;
+pub mod image_gen_status;
+pub use image_gen_status::{get_image_gen_status, save_generations_to_dir, wait_for_image_gen_status, Generation, StatusRequest, TaskResponse, VideoGenStatusResponse};

crates/lib/api_clients/openai_sora_client/src/image_gen/test_status_response.json

@@ -0,0 +1,172 @@
+{
+  "task_responses": [
+      {
+          "id": "task_01jr7xwwjyem5t268ym01wn9qs",
+          "user": "user-eOFctNDVUgUV6aaW513mcCFq",
+          "created_at": "2025-04-07T10:49:18.686690Z",
+          "status": "succeeded",
+          "progress_pct": 0.6789999999999999,
+          "progress_pos_in_queue": null,
+          "estimated_queue_wait_time": null,
+          "queue_status_message": null,
+          "priority": 1,
+          "type": "image_gen",
+          "prompt": "Can we make an alien planet skyscape where the clouds were giant whales. Maybe the planet is a gas giant",
+          "n_variants": 2,
+          "n_frames": 1,
+          "height": 480,
+          "width": 854,
+          "model": null,
+          "operation": "simple_compose",
+          "inpaint_items": [],
+          "preset_id": null,
+          "caption": null,
+          "actions": null,
+          "interpolation": null,
+          "sdedit": null,
+          "remix_config": null,
+          "quality": "standard",
+          "generations": [
+              {
+                  "id": "gen_01jr7xycyxfptaqtcq65bz8f46",
+                  "task_id": "task_01jr7xwwjyem5t268ym01wn9qs",
+                  "created_at": "2025-04-07T10:50:08.231958Z",
+                  "deleted_at": null,
+                  "url": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fsrc_1.png?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=iulvTinTA8YVVFDIEjoPKqZT%2BL%2BZvs5aT0mhjQg6YKw%3D&az=oaivgprodscus",
+                  "seed": 12345,
+                  "can_download": true,
+                  "download_status": "ready",
+                  "is_favorite": false,
+                  "is_liked": false,
+                  "is_public": true,
+                  "is_archived": false,
+                  "is_featured": false,
+                  "has_feedback": false,
+                  "like_count": 0,
+                  "cloudflare_metadata": null,
+                  "cf_thumbnail_url": null,
+                  "encodings": {
+                      "source": {
+                          "path": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fimg_1.webp?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=NU1Nx7GtchzkYdeb6z5TEbm9Y31jIUcTwPdqBl65T7s%3D&az=oaivgprodscus",
+                          "size": null,
+                          "width": null,
+                          "height": null,
+                          "duration_secs": null,
+                          "ssim": null
+                      },
+                      "md": null,
+                      "ld": null,
+                      "thumbnail": {
+                          "path": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fimg_1_thumb.webp?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=AQzzZYs0tYGBWwnFM5I03Wg6K8VZiNjLDMalRjxokeY%3D&az=oaivgprodscus",
+                          "size": null
+                      },
+                      "spritesheet": null
+                  },
+                  "width": 1536,
+                  "height": 1024,
+                  "n_frames": 1,
+                  "prompt": "Can we make an alien planet skyscape where the clouds were giant whales. Maybe the planet is a gas giant",
+                  "title": "Whale Cloud Skyscape",
+                  "actions": null,
+                  "inpaint_items": null,
+                  "interpolation": null,
+                  "sdedit": null,
+                  "operation": "simple_compose",
+                  "model": null,
+                  "preset_id": null,
+                  "user": {
+                      "id": "user-eOFctNDVUgUV6aaW513mcCFq",
+                      "username": "kasisnu"
+                  },
+                  "moderation_result": {
+                      "type": "passed",
+                      "results_by_frame_index": {},
+                      "code": null,
+                      "is_output_rejection": false,
+                      "task_id": "task_01jr7xwwjyem5t268ym01wn9qs"
+                  },
+                  "paragen_status": null,
+                  "task_type": "image_gen",
+                  "remix_config": null,
+                  "quality": "standard"
+              },
+              {
+                  "id": "gen_01jr7xycz4fbzrf6kjj70hpfb7",
+                  "task_id": "task_01jr7xwwjyem5t268ym01wn9qs",
+                  "created_at": "2025-04-07T10:50:08.233213Z",
+                  "deleted_at": null,
+                  "url": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fsrc_0.png?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=b%2Frn10WEfVBLqLao9k1UgFrS6QqJis8Co81tSx8oI7M%3D&az=oaivgprodscus",
+                  "seed": 12345,
+                  "can_download": true,
+                  "download_status": "ready",
+                  "is_favorite": false,
+                  "is_liked": false,
+                  "is_public": true,
+                  "is_archived": false,
+                  "is_featured": false,
+                  "has_feedback": false,
+                  "like_count": 0,
+                  "cloudflare_metadata": null,
+                  "cf_thumbnail_url": null,
+                  "encodings": {
+                      "source": {
+                          "path": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fimg_0.webp?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=o1Fdh1mh1O0P6jmnkcJlf4F7FWb3h%2BKSVRtVIXqjEJE%3D&az=oaivgprodscus",
+                          "size": null,
+                          "width": null,
+                          "height": null,
+                          "duration_secs": null,
+                          "ssim": null
+                      },
+                      "md": null,
+                      "ld": null,
+                      "thumbnail": {
+                          "path": "https://videos.openai.com/vg-assets/assets%2Ftask_01jr7xwwjyem5t268ym01wn9qs%2Fimg_0_thumb.webp?st=2025-04-07T09%3A33%3A44Z&se=2025-04-13T10%3A33%3A44Z&sks=b&skt=2025-04-07T09%3A33%3A44Z&ske=2025-04-13T10%3A33%3A44Z&sktid=a48cca56-e6da-484e-a814-9c849652bcb3&skoid=aa5ddad1-c91a-4f0a-9aca-e20682cc8969&skv=2019-02-02&sv=2018-11-09&sr=b&sp=r&spr=https%2Chttp&sig=2l%2BemWaC8q5o%2BIhxI5t%2BexjGe3F8ipPnPym0IIygvrA%3D&az=oaivgprodscus",
+                          "size": null
+                      },
+                      "spritesheet": null
+                  },
+                  "width": 1536,
+                  "height": 1024,
+                  "n_frames": 1,
+                  "prompt": "Can we make an alien planet skyscape where the clouds were giant whales. Maybe the planet is a gas giant",
+                  "title": "Whale Cloud Skyscape",
+                  "actions": null,
+                  "inpaint_items": null,
+                  "interpolation": null,
+                  "sdedit": null,
+                  "operation": "simple_compose",
+                  "model": null,
+                  "preset_id": null,
+                  "user": {
+                      "id": "user-eOFctNDVUgUV6aaW513mcCFq",
+                      "username": "kasisnu"
+                  },
+                  "moderation_result": {
+                      "type": "passed",
+                      "results_by_frame_index": {},
+                      "code": null,
+                      "is_output_rejection": false,
+                      "task_id": "task_01jr7xwwjyem5t268ym01wn9qs"
+                  },
+                  "paragen_status": null,
+                  "task_type": "image_gen",
+                  "remix_config": null,
+                  "quality": "standard"
+              }
+          ],
+          "num_unsafe_generations": 0,
+          "title": "Whale Cloud Skyscape",
+          "moderation_result": {
+              "type": "passed",
+              "results_by_frame_index": {},
+              "code": null,
+              "is_output_rejection": false,
+              "task_id": "task_01jr7xwwjyem5t268ym01wn9qs"
+          },
+          "failure_reason": null,
+          "needs_user_review": false
+      }
+  ],
+  "last_id": "task_01jr7xwwjyem5t268ym01wn9qs",
+  "has_more": false
+}

crates/lib/mysql_queries/src/payloads/generic_inference_args/inner_payloads/sora_image_gen_args.rs

@@ -20,4 +20,8 @@ pub struct SoraImageGenArgs {
   #[serde(rename = "su")]
   #[serde(skip_serializing_if = "Option::is_none")]
   pub maybe_sora_media_upload_tokens: Option<Vec<String>>,
+
+  #[serde(rename = "st")]
+  #[serde(skip_serializing_if = "Option::is_none")]
+  pub maybe_sora_task_id: Option<String>,
 }

crates/lib/redis_common/src/lib.rs

@@ -3,4 +3,4 @@
 pub mod payloads;
 pub mod redis_cache_keys;
 pub mod redis_keys;
-pub mod shared_constants;
+pub mod shared_constants;

crates/service/job/inference_job/Cargo.toml

@@ -40,6 +40,7 @@ memory_caching = { path = "../../../lib/caching/memory_caching" }
 migration = { path = "../../../lib/migration" }
 mimetypes = { path = "../../../lib/files/mimetypes" }
 mysql_queries = { path = "../../../lib/mysql_queries" }
+openai_sora_client = { path = "../../../lib/api_clients/openai_sora_client" }
 primitives = { path = "../../../lib/data/primitives" }
 redis_common = { path = "../../../lib/redis_common" }
 redis_schema = { path = "../../../lib/redis_schema" }

crates/service/job/inference_job/src/job/job_loop/process_single_job.rs

@@ -1,3 +1,4 @@
+use crate::job::job_types::image_generation::v2::process_single_sora_job::process_single_sora_job;
 use std::time::Instant;
 
 use anyhow::anyhow;
@@ -20,6 +21,7 @@ use crate::job::job_types::f5_tts::process_single_f5_tts_job::process_single_f5_
 use crate::job::job_types::format_conversion::process_single_format_conversion_job::process_single_format_conversion_job;
 use crate::job::job_types::gpt_sovits::process_single_gpt_sovits_job::process_single_gpt_sovits_job;
 use crate::job::job_types::image_generation::process_single_ig_job::process_single_ig_job;
+use crate::job::job_types::image_generation::v2::process_single_sora_job;
 use crate::job::job_types::lipsync::process_single_lipsync_job::process_single_lipsync_job;
 use crate::job::job_types::mocap::process_single_mc_job::process_single_mc_job;
 use crate::job::job_types::render_engine_scene::process_render_engine_scene::process_single_render_engine_scene_job;
@@ -265,6 +267,7 @@ fn can_use_new_dispatch(job: &AvailableInferenceJob) -> bool {
     InferenceJobType::RvcV2 => true,
     InferenceJobType::SeedVc => true,
     InferenceJobType::VideoRender => true,
+    InferenceJobType::ImageGenApi => true,
     _ => false,
   }
 }
@@ -297,6 +300,10 @@ async fn new_dispatch(
     InferenceJobType::RvcV2 => {
       dispatch_rvc_v2_job(job_dependencies, job).await?
     },
+    InferenceJobType::ImageGenApi => {
+      process_single_ig_job(job_dependencies, job).await?
+    }
+
     // NB: Make sure to add the job to `can_use_new_dispatch`.
     _ => {
       return Err(ProcessSingleJobError::InvalidJob(

crates/service/job/inference_job/src/job/job_types/image_generation/mod.rs

@@ -1,2 +1,3 @@
 pub mod process_single_ig_job;
-pub mod sd;
+pub mod sd;
+pub mod v2;

crates/service/job/inference_job/src/job/job_types/image_generation/process_single_ig_job.rs

@@ -1,3 +1,4 @@
+use enums::by_table::generic_inference_jobs::inference_job_type::InferenceJobType;
 use enums::by_table::generic_inference_jobs::inference_model_type::InferenceModelType;
 use errors::anyhow;
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
@@ -6,29 +7,24 @@ use crate::job::job_loop::job_success_result::JobSuccessResult;
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 use crate::job::job_types::image_generation::sd::process_job::process_job_selection;
 use crate::job::job_types::image_generation::sd::process_job::StableDiffusionProcessArgs;
+use crate::job::job_types::image_generation::v2::process_single_sora_job::{process_single_sora_job, ImageGenerationProcessArgs};
 use crate::state::job_dependencies::JobDependencies;
 
 pub async fn process_single_ig_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result<JobSuccessResult, ProcessSingleJobError> {
-    match job.maybe_model_type {
-        Some(InferenceModelType::StableDiffusion) => {
-            dispatch_sd_job(
-            job_dependencies,
-            job
-          ).await
-        },
-        Some(other_model_type) => {
-          Err(ProcessSingleJobError::Other(anyhow!("Wrong model type for SD: {:?}", other_model_type)))
-        }
-        None => {
-          Err(ProcessSingleJobError::Other(anyhow!("SD model type not set")))
-        }
-      }
+  match job.job_type {
+    InferenceJobType::ImageGenApi => {
+      let args = ImageGenerationProcessArgs { job_dependencies, job };
+      process_single_sora_job(&args).await
+    },
+    _ => match job.maybe_model_type {
+      Some(InferenceModelType::StableDiffusion) => dispatch_sd_job(job_dependencies, job).await,
+      Some(other_model_type) => Err(ProcessSingleJobError::Other(anyhow!("Wrong model type for SD: {:?}", other_model_type))),
+      None => Err(ProcessSingleJobError::Other(anyhow!("SD model type not set"))),
+    },
+  }
 }
 
 pub async fn dispatch_sd_job(job_dependencies: &JobDependencies, job: &AvailableInferenceJob) -> Result<JobSuccessResult, ProcessSingleJobError> {
-    let args = StableDiffusionProcessArgs {
-        job_dependencies,
-        job,
-    };
-    process_job_selection(args).await
-}
+  let args = StableDiffusionProcessArgs { job_dependencies, job };
+  process_job_selection(args).await
+}

crates/service/job/inference_job/src/job/job_types/image_generation/v2/mod.rs

@@ -0,0 +1 @@
+pub mod process_single_sora_job;

crates/service/job/inference_job/src/job/job_types/image_generation/v2/process_single_sora_job.rs

@@ -0,0 +1,344 @@
+use std::time::Instant;
+
+use anyhow::anyhow;
+use bucket_paths::legacy::remote_file_manager_paths::media_descriptor::MediaImagePngDescriptor;
+
+use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
+use enums::by_table::generic_synthetic_ids::id_category::IdCategory;
+use enums::by_table::media_files::media_file_class::MediaFileClass;
+use enums::by_table::media_files::media_file_origin_category::MediaFileOriginCategory;
+use enums::by_table::media_files::media_file_origin_model_type::MediaFileOriginModelType;
+use enums::by_table::media_files::media_file_origin_product_category::MediaFileOriginProductCategory;
+use enums::by_table::media_files::media_file_type::MediaFileType;
+use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs;
+use mysql_queries::queries::media_files::create::generic_insert::insert_media_file_generic_from_job::{insert_media_file_generic_from_job, InsertFromJobArgs};
+
+use crate::job::job_loop::job_success_result::{JobSuccessResult, ResultEntity};
+use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
+use crate::state::job_dependencies::JobDependencies;
+use crate::util::extractors::get_polymorphic_args_from_job::get_polymorphic_args_from_job;
+use cloud_storage::remote_file_manager::remote_cloud_file_manager::RemoteCloudFileClient;
+use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity;
+use errors::AnyhowResult;
+use mysql_queries::queries::batch_generations::insert_batch_generation_records::{insert_batch_generation_records, InsertBatchArgs};
+use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
+use openai_sora_client::credentials::SoraCredentials;
+use openai_sora_client::image_gen::image_gen_status::TaskStatus;
+use openai_sora_client::image_gen::wait_for_image_gen_status;
+use r2d2_redis::r2d2::PooledConnection;
+use r2d2_redis::redis::Commands;
+use r2d2_redis::RedisConnectionManager;
+use std::collections::HashMap;
+use tokens::tokens::batch_generations::BatchGenerationToken;
+
+/// Redis key
+const SORA_SECRET_REDIS_KEY : &str = "sora_secret";
+
+// Fields within the HKEY
+const BEARER_SUBKEY : &str = "bearer";
+const COOKIE_SUBKEY : &str = "cookie";
+const SENTINEL_SUBKEY : &str = "sentinel";
+
+
+/// Sora credentials stored in Redis
+
+
+#[derive(Clone,Copy,Debug)]
+pub enum RedisSoraCredentialSubkey {
+  Bearer,
+  Cookie,
+  Sentinel,
+}
+
+
+impl RedisSoraCredentialSubkey {
+  pub fn to_str(&self) -> &'static str {
+    match self {
+      RedisSoraCredentialSubkey::Bearer => BEARER_SUBKEY,
+      RedisSoraCredentialSubkey::Cookie => COOKIE_SUBKEY,
+      RedisSoraCredentialSubkey::Sentinel => SENTINEL_SUBKEY,
+    }
+  }
+}
+
+pub fn get_sora_credentials(
+  redis: &mut PooledConnection<RedisConnectionManager>
+) -> AnyhowResult<SoraCredentials> {
+
+  let values : HashMap<String, String> = redis.hgetall(SORA_SECRET_REDIS_KEY)
+    .map_err(|e| anyhow!("Failed to get Sora credentials from Redis: {}", e))?;
+
+  let bearer = values.get(BEARER_SUBKEY);
+  let cookie = values.get(COOKIE_SUBKEY);
+  let sentinel = values.get(SENTINEL_SUBKEY);
+
+  match (bearer, cookie, sentinel) {
+    (Some(b), Some(c), Some(s)) => {
+      Ok(
+        SoraCredentials {
+          bearer_token: b.to_string(),
+          cookie: c.to_string(),
+          sentinel: Some(s.to_string()),
+        }
+      )
+    }
+    _ => Err(anyhow!("redis sora credential values not present")),
+  }
+}
+
+
+pub struct ImageGenerationProcessArgs<'a> {
+  pub job_dependencies: &'a JobDependencies,
+  pub job: &'a AvailableInferenceJob,
+}
+
+pub struct ImageGenerationArgs {
+  pub number_of_samples: u32,
+  pub sora_task_id: String,
+}
+
+pub async fn ig_args_from_job(args: &ImageGenerationProcessArgs<'_>) -> Result<ImageGenerationArgs, ProcessSingleJobError> {
+  let polymorphic_args = get_polymorphic_args_from_job(&args.job)?;
+
+  let job_args = match polymorphic_args {
+    PolymorphicInferenceArgs::Sg(args) => args,
+    _ => {
+      return Err(
+        ProcessSingleJobError::from_anyhow_error(anyhow!("wrong inner args for job!"))
+      );
+    }
+  };
+
+  let number_of_samples = job_args.maybe_number_of_samples.unwrap_or(1);
+  let sora_task_id = job_args.maybe_sora_task_id.clone().ok_or_else(|| ProcessSingleJobError::from_anyhow_error(anyhow!("Missing Sora Task ID")))?;
+
+
+  Ok(ImageGenerationArgs {
+    number_of_samples,
+    sora_task_id,
+  })
+}
+
+pub async fn process_single_sora_job(args: &ImageGenerationProcessArgs<'_>) -> Result<JobSuccessResult, ProcessSingleJobError> {
+  let job = args.job;
+  let deps = args.job_dependencies;
+  let mysql_pool = &deps.db.mysql_pool;
+
+  let ig_args = ig_args_from_job(&args).await?;
+  // let ig_deps: &crate::job::job_types::image_generation::v2::stable_diffusion_dependencies::StableDiffusionDependencies = match
+  // &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies
+  // {
+  //   None => {
+  //     return Err(ProcessSingleJobError::Other(anyhow!("Missing Job Specific Dependencies")));
+  //   }
+  //   Some(val) => { val }
+  // };
+
+  let _job_progress_reporter = args.job_dependencies.clients.job_progress_reporter.new_generic_inference(job.inference_job_token.as_str()).map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+
+  //==================== TEMP DIR ==================== //
+  let work_temp_dir = format!("temp_image_generation_{}", job.id.0);
+
+  //NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
+  let work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work.new_tempdir(&work_temp_dir).map_err(|e| ProcessSingleJobError::from_io_error(e))?;
+
+  // thread::sleep(seconds) to check the directory
+
+  let stderr_output_file = work_temp_dir.path().join("sd_err.txt");
+  let stdout_output_file = work_temp_dir.path().join("sd_out.txt");
+
+
+
+
+  let inference_start_time = Instant::now();
+
+  // Use the sora client to wait for the task to complete and get the resulting images.
+
+  // let sora_credentials = get_sora_credentials_from_request(http_request);
+
+  let redis_pool_dep = deps
+    .db
+    .maybe_keepalive_redis_pool.clone();
+
+  let redis_pool = redis_pool_dep
+    .ok_or_else(|| ProcessSingleJobError::Other(anyhow!("failed to get redis pool")))?;
+
+  let sora_credentials = get_sora_credentials(&mut redis_pool.get().map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?)?;
+
+  let sora_task_id = ig_args.sora_task_id.clone();
+  let output_path = work_temp_dir.path().join("output");
+
+
+  let sora_task_response = wait_for_image_gen_status(
+    &sora_task_id,
+    &sora_credentials,
+    None,
+  ).await.map_err(|e| ProcessSingleJobError::Other(anyhow!(e)))?;
+  // let match TaskStatus::from_str(&task_response.status) {
+  //   TaskStatus::Succeeded => {
+  //   },
+
+
+  let generations = match TaskStatus::from_str(&sora_task_response.status) {
+    TaskStatus::Succeeded => {
+      let generations = sora_task_response.generations;
+      if generations.is_empty() {
+        return Err(ProcessSingleJobError::Other(anyhow!("No generations found in task response")));
+      }
+      generations
+    }
+    TaskStatus::Failed => {
+      return Err(ProcessSingleJobError::Other(anyhow!("Task failed: {:?}", sora_task_response)));
+    }
+    _ => {
+      return Err(ProcessSingleJobError::Other(anyhow!("Task not completed yet: {:?}", sora_task_response)));
+    }
+  };
+
+  // hack to check the directory before clean up.
+  //   let thirtyMinutes = 1800;
+  //   thread::sleep(Duration::from_secs(thirtyMinutes));
+  // upload media and create a record.
+
+  let inference_duration = Instant::now().duration_since(inference_start_time);
+
+  // let mut entries = vec![];
+  // let inputs = MediaFileExtraInfo::S(StableDiffusionExtraInfo {
+  //   prompt: Some(prompt.clone()),
+  //   cfg_scale: Some(sd_args.maybe_cfg_scale.unwrap_or(7)),
+  //   negative_prompt: sd_args.maybe_n_prompt,
+  //   lora_model_weight_token: Some(lora_token),
+  //   lora_name: Some(lora_name),
+  //   sampler: Some(sd_args.maybe_sampler.unwrap_or(String::from("Euler a"))),
+  //   width: Some(sd_args.maybe_width.unwrap_or(512)),
+  //   height: Some(sd_args.maybe_height.unwrap_or(512)),
+  //   seed: Some(sd_args.maybe_seed.unwrap_or(1)),
+  //   number_of_samples: Some(number_of_samples),
+  // });
+
+  let batch_token = BatchGenerationToken::generate();
+  // let prompt_token = PromptToken::generate();
+  let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await?;
+  let mut entries = vec![];
+
+  let mut maybe_first_media_file_token = None;
+
+  for (i, generation) in generations.iter().enumerate() {
+    let path = output_path.clone();
+
+
+    let file_path = format!("{}/{}.png", output_path.to_str().unwrap(), generation.id.clone());
+
+    println!("Upload File Path:{}", file_path);
+
+    let metadata = remote_cloud_file_client.upload_file(
+      Box::new(MediaImagePngDescriptor {}),
+      file_path.as_ref(),
+    ).await?;
+
+    let bucket_details = match metadata.bucket_details {
+      Some(val) => { val }
+      None => {
+        return Err(
+          ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
+        );
+      }
+    };
+
+  // extra_file_modification_info: todo!(), // JSON ENCODED STRUCT
+
+  // TODO(ks: 2025-04-08): This should insert the media file token and maybe additional metadata about the generation
+  let (media_file_token, _id) = insert_media_file_generic_from_job(InsertFromJobArgs {
+    pool: mysql_pool,
+    job,
+    media_class: MediaFileClass::Image,
+    media_type: MediaFileType::Image, // TODO(bt,2024-04-30): This should be a specific type of image
+    origin_category: MediaFileOriginCategory::Inference,
+    origin_product_category: MediaFileOriginProductCategory::ImageGeneration,
+    maybe_origin_model_type: Some(MediaFileOriginModelType::StorytellerStudioImageGen),
+    maybe_origin_model_token: None,
+    maybe_origin_filename: Some(file_path),
+    maybe_batch_token: Some(&batch_token),
+    maybe_prompt_token: None,
+    maybe_mime_type: Some(metadata.mimetype.as_ref()),
+    file_size_bytes: metadata.file_size_bytes,
+    maybe_duration_millis: Some(inference_duration.as_millis() as u64),
+    maybe_audio_encoding: None,
+    maybe_video_encoding: None,
+    maybe_frame_width: Some(1024),
+    maybe_frame_height: Some(1024),
+    checksum_sha2: metadata.sha256_checksum.as_str(),
+    maybe_text_transcript: None,
+    public_bucket_directory_hash: bucket_details.object_hash.as_str(),
+    maybe_public_bucket_prefix: Some(bucket_details.prefix.as_str()),
+    maybe_public_bucket_extension: Some(bucket_details.suffix.as_str()),
+    maybe_extra_media_info: None,
+    maybe_creator_file_synthetic_id_category: IdCategory::MediaFile,
+    maybe_creator_category_synthetic_id_category: IdCategory::ModelWeights,
+    maybe_mod_user_token: None,
+    is_generated_on_prem: args.job_dependencies.job.info.container.is_on_prem,
+    generated_by_worker: Some(&args.job_dependencies.job.info.container.hostname),
+    generated_by_cluster: Some(&args.job_dependencies.job.info.container.cluster_name),
+    maybe_title: None,
+    maybe_scene_source_media_file_token: None,
+    is_intermediate_system_file: false,
+  })
+  .await?;
+
+  if maybe_first_media_file_token.is_none() {
+    maybe_first_media_file_token = Some(media_file_token.clone());
+  }
+
+  let batch_generation_entity: BatchGenerationEntity = BatchGenerationEntity::MediaFile(
+    media_file_token
+  );
+
+    entries.push(batch_generation_entity);
+  }
+
+  // TODO(bt,2024-02-22): This transaction should wrap everything
+  let mut transaction = mysql_pool.begin().await.map_err(|e| ProcessSingleJobError::from_anyhow_error(anyhow!(e)))?;
+
+  let batch_token_result = insert_batch_generation_records(InsertBatchArgs {
+   entries,
+   maybe_existing_batch_token: Some(&batch_token),
+   transaction: &mut transaction,
+  }).await;
+
+  let _batch_token = match batch_token_result {
+   Ok(v) => { v.to_string() }
+   Err(_err) => {
+     return Err(
+       ProcessSingleJobError::from_anyhow_error(
+         anyhow!("No batch token? something has failed.")
+       )
+     );
+   }
+  };
+
+  // NB: Don't fail the job if the query fails.
+  // let prompt_result = insert_prompt(InsertPromptArgs {
+  //   maybe_apriori_prompt_token: Some(&prompt_token),
+  //   prompt_type: PromptType::StableDiffusion,
+  //   maybe_creator_user_token: job.maybe_creator_user_token_typed.as_ref(),
+  //   maybe_positive_prompt: Some(&positive_prompt),
+  //   maybe_negative_prompt: maybe_negative_prompt.as_deref(),
+  //   maybe_other_args: None, // TODO(bt,2024-02-22): Support other arguments
+  //   creator_ip_address: &job.creator_ip_address,
+  //   mysql_executor: &mut *transaction,
+  //   phantom: Default::default(),
+  // }).await;
+
+  // match prompt_result {
+  //   Ok(_token) => {}
+  //   Err(err) => {
+  //     error!("No prompt result token? something has failed: {:?}", err);
+  //     return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No prompt result token? something has failed.")));
+  //   }
+  // }
+
+  let _r = transaction.commit().await.map_err(|e| ProcessSingleJobError::from_anyhow_error(anyhow!(e)))?;
+
+  // TODO(bt,2024-02-12): Return the batch token instead. (We're not ready for that.)
+  Ok(JobSuccessResult { inference_duration, maybe_result_entity: maybe_first_media_file_token.map(|media_file_token| ResultEntity { entity_type: InferenceResultType::MediaFile, entity_token: media_file_token.to_string() }) })
+}

crates/service/web/storyteller_web/src/http_server/endpoints/image_studio/prompt/enqueue_studio_image_generation_handler.rs

@@ -340,6 +340,7 @@ pub async fn enqueue_studio_image_generation_handler(
     maybe_additional_media_file_tokens: Some(additional_images),
     maybe_number_of_samples: Some(number_of_samples),
     maybe_sora_media_upload_tokens: Some(sora_media_tokens),
+    maybe_sora_task_id: Some(response.task_id),
   };
 
   // create the inference args here