diff --git a/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs b/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs index ff6f6f5c9e..8c010b970e 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/create/create_weight.rs +++ b/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, pub maybe_public_bucket_extension: Option, - //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, pub version: i32, pub mysql_pool: &'a MySqlPool, } @@ -78,10 +74,6 @@ pub async fn create_weight(args: CreateModelWeightsArgs<'_>) -> AnyhowResult { 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 { 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 { 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 { - - // 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 { Ok(JobSuccessResult { maybe_result_entity: None, inference_duration: Duration::from_secs(0), }) } -pub async fn process_job_lora(_args: &StableDiffusionProcessArgs<'_>) -> Result { + +// 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 { + 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 { 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; + + 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. +// */ diff --git a/crates/service/web/storyteller_web/src/routes.rs b/crates/service/web/storyteller_web/src/routes.rs index d806a10c83..4b24ff780b 100644 --- a/crates/service/web/storyteller_web/src/routes.rs +++ b/crates/service/web/storyteller_web/src/routes.rs @@ -1333,4 +1333,4 @@ fn add_weights_routes(app: App) -> App .route("/list", web::get().to(list_available_weights_handler)) .route("/list_featured", web::get().to(list_featured_weights_handler)) ) -} +} \ No newline at end of file