Builds on inference for batch creation of entries.

crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs

@@ -26,10 +26,6 @@ pub struct CreateModelWeightsArgs<'a> {
     pub public_bucket_hash: String,
     pub maybe_public_bucket_prefix: Option<String>,
     pub maybe_public_bucket_extension: Option<String>,
-    //pub cached_user_ratings_total_count: u32,
-    //pub cached_user_ratings_positive_count: u32,
-    //pub cached_user_ratings_negative_count: u32,
-    //pub maybe_cached_user_ratings_ratio: Option<f32>,
     pub version: i32,
     pub mysql_pool: &'a MySqlPool,
 }
@@ -78,10 +74,6 @@ pub async fn create_weight(args: CreateModelWeightsArgs<'_>) -> AnyhowResult<Mod
             args.public_bucket_hash,
             args.maybe_public_bucket_prefix,
             args.maybe_public_bucket_extension,
-            //args.cached_user_ratings_total_count,
-            //args.cached_user_ratings_positive_count,
-            //args.cached_user_ratings_negative_count,
-            //args.maybe_cached_user_ratings_ratio,
             args.version
         )
         .execute(args.mysql_pool).await;

crates/service/job/inference_job/src/job/job_types/image_generation/sd/process_job.rs

@@ -1,7 +1,9 @@
 use std::thread;
-use std::time::{Duration, Instant};
+use std::time::{ Duration, Instant };
 
 use anyhow::anyhow;
+use composite_identifiers::by_table::batch_generations::batch_generation_entity::BatchGenerationEntity;
+use enums::by_table::generic_inference_jobs::inference_result_type::InferenceResultType;
 use log::info;
 
 use cloud_storage::remote_file_manager::media_descriptor::MediaImagePngDescriptor;
@@ -14,14 +16,25 @@ use filesys::path_to_string::path_to_string;
 use google_drive_common::google_drive_download_command::GoogleDriveDownloadCommand;
 use mysql_queries::payloads::generic_inference_args::generic_inference_args::PolymorphicInferenceArgs;
 use mysql_queries::payloads::generic_inference_args::image_generation_payload::StableDiffusionArgs;
+use mysql_queries::queries::batch_generations::insert_batch_generation_records::{InsertBatchArgs, self};
+use mysql_queries::queries::batch_generations::insert_batch_generation_records::insert_batch_generation_records;
 use mysql_queries::queries::generic_inference::job::list_available_generic_inference_jobs::AvailableInferenceJob;
-use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{insert_media_file_from_file_upload, InsertMediaFileFromUploadArgs, UploadType};
-use mysql_queries::queries::model_weights::create::create_weight::{create_weight, CreateModelWeightsArgs};
+use mysql_queries::queries::media_files::create::insert_media_file_from_file_upload::{
+    insert_media_file_from_file_upload,
+    InsertMediaFileFromUploadArgs,
+    UploadType,
+};
+use mysql_queries::queries::model_weights::create::create_weight::{
+    create_weight,
+    CreateModelWeightsArgs,
+};
 use mysql_queries::queries::model_weights::get_weight::get_weight_by_token;
+use tokens::tokens::anonymous_visitor_tracking::AnonymousVisitorTrackingToken;
+use tokens::tokens::media_files::MediaFileToken;
 use tokens::tokens::model_weights::ModelWeightToken;
 use tokens::tokens::users::UserToken;
 
-use crate::job::job_loop::job_success_result::JobSuccessResult;
+use crate::job::job_loop::job_success_result::{ JobSuccessResult, ResultEntity };
 use crate::job::job_loop::process_single_job_error::ProcessSingleJobError;
 use crate::job::job_types::image_generation::sd::sd_inference_command::InferenceArgs;
 use crate::job_dependencies::JobDependencies;
@@ -31,16 +44,14 @@ pub struct StableDiffusionProcessArgs<'a> {
     pub job: &'a AvailableInferenceJob,
 }
 
-
 // run inference
 // insert record into the db with the inference job token complete.
-
 pub async fn sd_args_from_job(
     args: &StableDiffusionProcessArgs<'_>
 ) -> Result<StableDiffusionArgs, ProcessSingleJobError> {
     let inference_args = args.job.maybe_inference_args
         .as_ref()
-        .map(|args| args.args.as_ref())
+        .map(|args: &mysql_queries::payloads::generic_inference_args::generic_inference_args::GenericInferenceArgs| args.args.as_ref())
         .flatten();
     let polymorphic_args = match inference_args {
         Some(args) => args,
@@ -64,37 +75,157 @@ pub async fn sd_args_from_job(
 
 // store the prompt and cluster them today.
 
-pub async fn process_job_selection (args: StableDiffusionProcessArgs<'_>
+pub async fn process_job_selection(
+    args: StableDiffusionProcessArgs<'_>
 ) -> Result<JobSuccessResult, ProcessSingleJobError> {
     let sd_args = sd_args_from_job(&args).await?;
+
     if sd_args.type_of_inference == "inference" {
         process_job_inference(&args).await
-    }
-    else if sd_args.type_of_inference == "lora" {
+    } else if sd_args.type_of_inference == "lora" {
         process_job_lora(&args).await
-    }
-    else if sd_args.type_of_inference == "checkpoint"{
+    } else if sd_args.type_of_inference == "model" {
         process_job_inference(&args).await
-    }
-    else {
+    } else {
         Err(ProcessSingleJobError::Other(anyhow!("inference type doesn't exist!")))
     }
 }
 
-pub async fn process_job_checkpoint(_args: &StableDiffusionProcessArgs<'_>) -> Result<JobSuccessResult, ProcessSingleJobError> {
-
-    // insert model record in create model record
-
-    // let maybe_result_entity = ResultEntity {
-    //     entity_type: InferenceResultType::UploadModel,
-    //     entity_token:
-    // };
+pub async fn process_job_model(
+    args: &StableDiffusionProcessArgs<'_>
+) -> Result<JobSuccessResult, ProcessSingleJobError> {
     Ok(JobSuccessResult {
         maybe_result_entity: None,
         inference_duration: Duration::from_secs(0),
     })
 }
-pub async fn process_job_lora(_args: &StableDiffusionProcessArgs<'_>) -> Result<JobSuccessResult, ProcessSingleJobError> {
+
+// let pool = setup().await;
+
+// let mut transaction = pool.begin().await.unwrap();
+
+// let entries = vec![
+//   BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_foo")),
+//   BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_bar")),
+//   BatchGenerationEntity::MediaFile(MediaFileToken::new_from_str("media_baz")),
+// ];
+
+// let r = insert_batch_generation_records(InsertBatchArgs {
+//   entries,
+//   transaction: &mut transaction,
+// }).await;
+
+pub async fn process_job_lora(
+    args: &StableDiffusionProcessArgs<'_>
+) -> Result<JobSuccessResult, ProcessSingleJobError> {
+    let job = args.job;
+    let deps = args.job_dependencies;
+    let mysql_pool = &deps.db.mysql_pool;
+
+    let sd_args = sd_args_from_job(&args).await?;
+
+    let sd_deps = 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 creator_ip_address = &job.creator_ip_address;
+    let creator_user_token: UserToken;
+    let anon_user_token: Option<&AnonymousVisitorTrackingToken>;
+    // The parameters will be updated on another screen perhaps?
+    // so right now it will fill with the availible  values.
+    let model_weight_token = &ModelWeightToken::generate();
+
+    let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0);
+    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))?;
+
+    let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors");
+    let lora_path = work_temp_dir.path().join("lora.safetensors");
+    let vae_path = work_temp_dir.path().join("vae.safetensors");
+    let output_path = work_temp_dir.path().join("output");
+    let g_drive_path = work_temp_dir.path().join("gdrive");
+
+    info!("Paths to download to:");
+    info!("sd_checkpoint_path:{}", sd_checkpoint_path.display());
+    info!("lora_path:{}", lora_path.display());
+    info!("vae_path:{}", vae_path.display());
+    info!("output_path:{}", output_path.display());
+    info!("tmp_google_drive_path:{}", g_drive_path.display());
+
+    let download_url = match sd_args.maybe_lora_upload_path {
+        Some(val) => { val }
+        None => { "".to_string() }
+    };
+
+    if download_url.len() == 0 {
+        return Err(
+            ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing"))
+        );
+    }
+
+    let file_name = "download.safetensors";
+    let download_script = easyenv::get_env_string_or_default(
+        "DOWNLOAD_SCRIPT",
+        "./scripts/download_internet_file.py"
+    );
+
+    let google_drive_downloader = GoogleDriveDownloadCommand::new_production(&download_script);
+    info!("Downloading {}", download_url);
+
+    let _download_filename = match
+        google_drive_downloader.download_file_with_file_name(
+            &download_url,
+            &work_temp_dir,
+            file_name
+        ).await
+    {
+        Ok(filename) => filename,
+        Err(_e) => {
+            return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download")));
+        }
+    };
+    let download_file_path = work_temp_dir.path().join(file_name);
+    info!("File Retrieved at {}", download_file_path.display());
+
+    // create_weight(CreateModelWeightsArgs {
+    //     token: model_weight_token,
+    //     weights_type: WeightsType::StableDiffusion15,
+    //     weights_category: WeightsCategory::ImageGeneration,
+    //     title: "".to_string(),
+    //     maybe_description_markdown: None,
+    //     maybe_description_rendered_html: None,
+    //     creator_user_token: Some(&creator_user_token),
+    //     creator_ip_address,
+    //     creator_set_visibility: Default::default(),
+    //     maybe_last_update_user_token: None,
+    //     original_download_url: None,
+    //     original_filename: None,
+    //     file_size_bytes: 0,
+    //     file_checksum_sha2: "".to_string(),
+    //     public_bucket_hash: "".to_string(),
+    //     maybe_public_bucket_prefix: None,
+    //     maybe_public_bucket_extension: None,
+    //     version: 0,
+    //     mysql_pool,
+    // }).await?;
+
+    Ok(JobSuccessResult {
+        maybe_result_entity: Some(ResultEntity {
+            entity_type: InferenceResultType::UploadModel,
+            entity_token: model_weight_token.to_string(),
+        }),
+        inference_duration: Duration::from_secs(0),
+    })
+}
+
+pub async fn process_inference(
+    _args: &StableDiffusionProcessArgs<'_>
+) -> Result<JobSuccessResult, ProcessSingleJobError> {
     Ok(JobSuccessResult {
         maybe_result_entity: None,
         inference_duration: Duration::from_secs(0),
@@ -108,22 +239,33 @@ pub async fn process_job_inference(
     let deps = args.job_dependencies;
     let mysql_pool = &deps.db.mysql_pool;
 
-
     let sd_args = sd_args_from_job(&args).await?;
-    let sd_deps = match &args.job_dependencies.job.job_specific_dependencies.maybe_stable_diffusion_dependencies {
+    let sd_deps = 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
-        }
+        Some(val) => { val }
     };
 
-    let creator_user_token:UserToken;
+    let creator_user_token: UserToken;
+    let anon_user_token: Option<AnonymousVisitorTrackingToken>;
+
+    match &job.maybe_creator_anonymous_visitor_token {
+        Some(token) => {
+            let anonymous_visitor_token = AnonymousVisitorTrackingToken::new_from_str(token);
+            anon_user_token = Some(anonymous_visitor_token);
+        }
+        None => {
+            anon_user_token = None;
+        }
+    }
+
     match &job.maybe_creator_user_token {
         Some(token) => {
             creator_user_token = UserToken::new_from_str(token);
-        },
+        }
         None => {
             return Err(ProcessSingleJobError::InvalidJob(anyhow!("Missing Creator User Token")));
         }
@@ -137,65 +279,23 @@ pub async fn process_job_inference(
 
     let work_temp_dir = format!("temp_stable_diffusion_inference_{}", job.id.0);
 
-    // thread::sleep(seconds) to check the directory
-
     //NB: TempDir exists until it goes out of scope, at which point it should delete from filesystem.
-    let mut work_temp_dir = args.job_dependencies.fs.scoped_temp_dir_creator_for_work
+    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))?;
 
-
     let sd_checkpoint_path = work_temp_dir.path().join("sd_checkpoint.safetensors");
     let lora_path = work_temp_dir.path().join("lora.safetensors");
     let vae_path = work_temp_dir.path().join("vae.safetensors");
     let output_path = work_temp_dir.path().join("output");
-    let g_drive_path = work_temp_dir.path().join("gdrive");
 
     info!("Paths to download to:");
-    info!("sd_checkpoint_path:{}",sd_checkpoint_path.display());
-    info!("lora_path:{}",lora_path.display());
-    info!("vae_path:{}",vae_path.display());
-    info!("output_path:{}",output_path.display());
-    info!("tmp_google_drive_path:{}",g_drive_path.display());
+    info!("sd_checkpoint_path:{}", sd_checkpoint_path.display());
+    info!("lora_path:{}", lora_path.display());
+    info!("vae_path:{}", vae_path.display());
+    info!("output_path:{}", output_path.display());
 
-    // If it is not inference then it must be loRA or checkpoint downloading
-    if sd_args.type_of_inference != "inference" {
-        let mut download_url= String::from("");
-
-        if sd_args.type_of_inference == "lora" {
-            download_url = match sd_args.maybe_lora_upload_path {
-                Some(val) => { val },
-                None => { "".to_string() }
-            };
-        }
-        else if sd_args.type_of_inference == "checkpoint" {
-            download_url = match sd_args.maybe_upload_path {
-                Some(val) => { val },
-                None => { "".to_string() }
-            }
-        }
-
-        if download_url.len() == 0 {
-            return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download URL Missing")));
-        }
-
-        let file_name = "download.safetensors";
-        let download_script = easyenv::get_env_string_or_default(
-            "DOWNLOAD_SCRIPT",
-            "./scripts/download_internet_file.py");
-        let google_drive_downloader =
-            GoogleDriveDownloadCommand::new_production(&download_script);
-        info!("Downloading {}",download_url);
-
-        let _download_filename = match google_drive_downloader.download_file_with_file_name(&download_url, &work_temp_dir,file_name).await {
-            Ok(filename) => filename,
-            Err(_e) => {
-                return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("Failed to Download")));
-            }
-        };
-        let download_file_path = work_temp_dir.path().join(file_name);
-        info!("File Retrieved at {}",download_file_path.display());
-    }
+    // thread::sleep(seconds) to check the directory
 
     // // Unpack loRA and Checkpoint
     // // run inference by downloading from google drive.
@@ -246,13 +346,16 @@ pub async fn process_job_inference(
         retrieved_sd_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
         retrieved_sd_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
     );
-    remote_cloud_file_client.download_file(details, path_to_string(sd_checkpoint_path.clone())).await?;
+    
+    remote_cloud_file_client.download_file(
+        details,
+        path_to_string(sd_checkpoint_path.clone())
+    ).await?;
 
     match retrieved_loRA_record {
         Some(record) => {
             match record {
                 Some(model_weight_record) => {
-
                     let lora_details = RemoteCloudBucketDetails::new(
                         model_weight_record.public_bucket_hash.clone(),
                         model_weight_record.maybe_public_bucket_prefix
@@ -266,7 +369,6 @@ pub async fn process_job_inference(
                         lora_details,
                         path_to_string(lora_path.clone())
                     ).await?;
-
                 }
                 None => {}
             }
@@ -278,7 +380,7 @@ pub async fn process_job_inference(
     let vae_token = String::from("weight_rb0959wfzjhk3d1k93hr3s0qw");
     let model_weight_vae = ModelWeightToken(vae_token);
 
-    let vae_weight_record =  get_weight_by_token(
+    let vae_weight_record = get_weight_by_token(
         &model_weight_vae,
         false,
         &deps.db.mysql_pool
@@ -287,29 +389,25 @@ pub async fn process_job_inference(
     let vae_weight_record = match vae_weight_record {
         Some(val) => val,
         None => {
-            return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem.")))
+            return Err(
+                ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
+            );
         }
     };
 
-
     let vae_details = RemoteCloudBucketDetails::new(
         vae_weight_record.public_bucket_hash.clone(),
-        vae_weight_record.maybe_public_bucket_prefix
-            .clone()
-            .unwrap_or_else(|| "".to_string()),
-        vae_weight_record.maybe_public_bucket_extension
-            .clone()
-            .unwrap_or_else(|| "".to_string())
+        vae_weight_record.maybe_public_bucket_prefix.clone().unwrap_or_else(|| "".to_string()),
+        vae_weight_record.maybe_public_bucket_extension.clone().unwrap_or_else(|| "".to_string())
     );
 
-    remote_cloud_file_client.download_file(
-        vae_details,
-        path_to_string(vae_path.clone())
-    ).await?;
+    remote_cloud_file_client.download_file(vae_details, path_to_string(vae_path.clone())).await?;
 
     let prompt = match sd_args.maybe_prompt {
         Some(val) => val,
-        None => return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!")))
+        None => {
+            return Err(ProcessSingleJobError::from_anyhow_error(anyhow!("No Prompt provided!")));
+        }
     };
 
     let stderr_output_file = work_temp_dir.path().join("sd_err.txt");
@@ -326,7 +424,7 @@ pub async fn process_job_inference(
         output_file: output_path.clone(),
         stderr_output_file,
         prompt,
-        negative_prompt:sd_args.maybe_n_prompt.unwrap_or_default(),
+        negative_prompt: sd_args.maybe_n_prompt.unwrap_or_default(),
         number_of_samples,
         samplers: sd_args.maybe_sampler.unwrap_or(String::from("Euler a")),
         width: sd_args.maybe_width.unwrap_or(512),
@@ -344,85 +442,81 @@ pub async fn process_job_inference(
     // run a for loop for output images output_0 in the folder then use upload media.
     // pngs
 
-    if sd_args.type_of_inference == "inference" {
-        // gather all the images into a single batch
+    let mut entries = vec![];
 
 
-        for i in 0..sd_args.maybe_batch_count.unwrap_or(1) {
-            let path = output_path.clone();
-            let file_path = format!("{}_{}.png", path_to_string(path), i);
-            println!("Upload File Path:{}", file_path);
-            let metadata = remote_cloud_file_client.upload_file(Box::new(MediaImagePngDescriptor {}), file_path.as_ref()).await?;
+    for i in 0..sd_args.maybe_batch_count.unwrap_or(1) {
+        let path = output_path.clone();
+        let file_path = format!("{}_{}.png", path_to_string(path), i);
+        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.")))
-            };
+        let bucket_details = match metadata.bucket_details {
+            Some(val) => { val }
+            None => {
+                return Err(
+                    ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
+                );
+            }
+        };
 
-            // TODO: Generic Insert with batch id.
-            // insert media file generic should run here.
-            insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs {
-                pool: mysql_pool,
-                maybe_creator_user_token: Some(&creator_user_token),
-                maybe_creator_anonymous_visitor_token: None,
-                creator_ip_address,
-                creator_set_visibility: Default::default(),
-                upload_type: UploadType::Filesystem,
-                media_file_type: MediaFileType::Image,
-                maybe_mime_type: Some(metadata.mimetype.as_ref()),
-                file_size_bytes: metadata.file_size_bytes,
-                duration_millis: inference_duration.as_secs() * 1000,
-                sha256_checksum: metadata.sha256_checksum.as_str(),
-                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()),
-            }).await?;
-        }
-
-        // hack to check the directory before clean up.
-        let thirtyMinutes = 1800;
-        thread::sleep(Duration::from_secs(thirtyMinutes));
-        // upload media and create a record.
-
-        Ok(JobSuccessResult {// TODO: batch token needs to go here
-            maybe_result_entity: None,
-            inference_duration: Duration::from_secs(0),
-        })
-    }
-    else { // must be an upload and need upload token
-
-
-        // The parameters will be updated on another screen perhaps?
-        // so right now it will fill with the availible  values.
-        create_weight(CreateModelWeightsArgs {
-            token: &ModelWeightToken::generate(),
-            weights_type: WeightsType::StableDiffusion15,
-            weights_category: WeightsCategory::ImageGeneration,
-            title: "".to_string(),
-            maybe_description_markdown: None,
-            maybe_description_rendered_html: None,
-            creator_user_token: Some(&creator_user_token),
+        // TODO: Generic Insert with batch id.
+        // insert media file generic should run here.
+        let media_file_token = insert_media_file_from_file_upload(InsertMediaFileFromUploadArgs {
+            pool: mysql_pool,
+            maybe_creator_user_token: Some(&creator_user_token),
+            maybe_creator_anonymous_visitor_token: anon_user_token.as_ref(),
             creator_ip_address,
             creator_set_visibility: Default::default(),
-            maybe_last_update_user_token: None,
-            original_download_url: None,
-            original_filename: None,
-            file_size_bytes: 0,
-            file_checksum_sha2: "".to_string(),
-            public_bucket_hash: "".to_string(),
-            maybe_public_bucket_prefix: None,
-            maybe_public_bucket_extension: None,
-            version: 0,
-            mysql_pool,
+            upload_type: UploadType::Filesystem,
+            media_file_type: MediaFileType::Image,
+            maybe_mime_type: Some(metadata.mimetype.as_ref()),
+            file_size_bytes: metadata.file_size_bytes,
+            duration_millis: inference_duration.as_secs() * 1000,
+            sha256_checksum: metadata.sha256_checksum.as_str(),
+            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()),
         }).await?;
 
-        Ok(JobSuccessResult {// TODO: batch token needs to go here
-            maybe_result_entity: None,
-            inference_duration: Duration::from_secs(0),
-        })
+        
+        let batch_generation_entity: BatchGenerationEntity = BatchGenerationEntity::MediaFile(media_file_token.0);
+        entries.push(batch_generation_entity);
     }
 
+ 
+    let mut transaction = mysql_pool.begin().await.unwrap();
+    let batch_token_result = insert_batch_generation_records(InsertBatchArgs {
+      entries,
+      transaction: &mut transaction,
+    }).await;
 
+    // hack to check the directory before clean up.
+    let thirtyMinutes = 1800;
+    thread::sleep(Duration::from_secs(thirtyMinutes));
+    // upload media and create a record.
+
+    let batch_token = match batch_token_result {
+        Ok(v) => {
+            v.to_string()
+        },
+        Err(err) => {
+            return Err(
+                ProcessSingleJobError::from_anyhow_error(anyhow!("no VAE? thats a problem."))
+            );
+        }
+    };
+
+    Ok(JobSuccessResult { // TODO: batch token needs to go here
+        maybe_result_entity: Some(ResultEntity {
+            entity_type: InferenceResultType::MediaFile,
+            entity_token: batch_token,
+        }),
+        inference_duration:inference_duration,
+    })
 
 }
 
@@ -445,9 +539,7 @@ mod tests {
         let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data");
         let remote_cloud_file_client = RemoteCloudFileClient::get_remote_cloud_file_client().await;
         let remote_cloud_file_client = match remote_cloud_file_client {
-            Ok(res) => {
-                res
-            }
+            Ok(res) => { res }
             Err(_) => {
                 return Err(anyhow!("failed to get remote cloud file client"));
             }
@@ -470,9 +562,42 @@ mod tests {
             suffix: String::from("safetensors"),
         };
 
-        remote_cloud_file_client.download_file(bucket_details1, String::from("./checkpoint")).await?;
+        remote_cloud_file_client.download_file(
+            bucket_details1,
+            String::from("./checkpoint")
+        ).await?;
         remote_cloud_file_client.download_file(bucket_details2, String::from("./loRA")).await?;
 
         Ok(())
     }
 }
+
+// /**
+//  *
+//  *
+//  * @CrossProduct this shows the batch table usage: https://github.com/storytold/storyteller-rust/blob/master/crates/lib/mysql_queries/src/queries/batch_generations/insert_batch_generation_records.rs#L83-L96
+// I wrote a test to show how you can insert multiple media file tokens into the table at once
+// I tested and it works:
+
+// mysql> select * from batch_generations;
+// +----+----------------------------------+-------------+--------------+---------------------+---------------------+
+// | id | token                            | entity_type | entity_token | created_at          | updated_at          |
+// +----+----------------------------------+-------------+--------------+---------------------+---------------------+
+// |  4 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file  | media_foo    | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 |
+// |  5 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file  | media_bar    | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 |
+// |  6 | batch_g_25dgzw53jdwgvesgdb48bam6 | media_file  | media_baz    | 2024-01-17 02:18:52 | 2024-01-17 02:18:52 |
+// |  7 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file  | media_foo    | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 |
+// |  8 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file  | media_bar    | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 |
+// |  9 | batch_g_79a2evktmnqbn9sq3k3eybze | media_file  | media_baz    | 2024-01-17 02:20:32 | 2024-01-17 02:20:32 |
+// +----+----------------------------------+-------------+--------------+---------------------+---------------------+
+// 6 rows in set (0.00 sec)
+
+// (I ran the test twice, which is why there are two separate batches of the same media file tokens)
+// I'll write an endpoint, probably something like:
+
+// GET api.fakeyou.com/v1/batches/{batch_token}
+
+// That will return a list of all the media files in the batch.
+
+// We can create a page around this. ALso, we can have a separate endpoint that looks up all the other files in a batch (if you're starting on a given media file page), so that you can show thumbnails of related generations on a media file page.
+//  */

crates/service/web/storyteller_web/src/routes.rs

@@ -1333,4 +1333,4 @@ fn add_weights_routes<T, B>(app: App<T>) -> App<T>
             .route("/list", web::get().to(list_available_weights_handler))
             .route("/list_featured", web::get().to(list_featured_weights_handler))
     )
-}
+}