diff --git a/crates/_types/buckets/src/public/weight_uploads/original_file.rs b/crates/_types/buckets/src/public/weight_uploads/original_file.rs index 7f0a40a6b7..42f203f849 100644 --- a/crates/_types/buckets/src/public/weight_uploads/original_file.rs +++ b/crates/_types/buckets/src/public/weight_uploads/original_file.rs @@ -2,24 +2,20 @@ use std::path::PathBuf; use crockford::crockford_entropy_lower; -use crate::public::media_uploads::directory::MediaUploadDirectory; +use crate::public::weight_uploads::directory::WeightUploadDirectory; use crate::public::public_path::PublicPath; const ORIGINAL_FILE_BASENAME : &str = "original_upload.bin"; -// TODO: Generate these from a macro. - -/// The original user upload file. -/// It may have derivative files (down samples, crops, etc.) that live alongside it. #[derive(Clone)] -pub struct MediaUploadOriginalFilePath { - directory: MediaUploadDirectory, +pub struct WeightUploadOriginalFilePath { + directory: WeightUploadDirectory, full_object_path: String, } -impl PublicPath for MediaUploadOriginalFilePath {} +impl PublicPath for WeightUploadOriginalFilePath {} -impl MediaUploadOriginalFilePath { +impl WeightUploadOriginalFilePath { pub fn generate_new() -> Self { let entropy = crockford_entropy_lower(32); @@ -28,7 +24,7 @@ impl MediaUploadOriginalFilePath { pub fn from_object_hash(hash: &str) -> Self { // TODO: Path construction could be cleaner. - let directory = MediaUploadDirectory::from_object_hash(hash); + let directory = WeightUploadDirectory::from_object_hash(hash); let full_object_path = format!("{}/{}", directory.get_directory_path_str(), ORIGINAL_FILE_BASENAME); Self { directory, @@ -44,7 +40,7 @@ impl MediaUploadOriginalFilePath { PathBuf::from(&self.full_object_path) } - pub fn get_directory(&self) -> &MediaUploadDirectory { + pub fn get_directory(&self) -> &WeightUploadDirectory { &self.directory } @@ -60,37 +56,36 @@ impl MediaUploadOriginalFilePath { #[cfg(test)] mod tests { use std::path::PathBuf; - - use crate::public::media_uploads::original_file::MediaUploadOriginalFilePath; + use crate::public::weight_uploads::original_file::WeightUploadOriginalFilePath; #[test] pub fn generate_new_entropy() { - let file = MediaUploadOriginalFilePath::generate_new(); + let file = WeightUploadOriginalFilePath::generate_new(); assert_eq!(file.get_object_hash().len(), 32); assert_eq!(file.get_directory().get_object_hash().len(), 32); } #[test] pub fn get_full_object_path_str() { - let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); - assert_eq!(file.get_full_object_path_str(), "/media_upload/a/b/c/d/e/abcdefghijk/original_upload.bin"); + let file = WeightUploadOriginalFilePath::from_object_hash("abcdefghijk"); + assert_eq!(file.get_full_object_path_str(), "/weight_upload/a/b/c/d/e/abcdefghijk/original_upload.bin"); } #[test] pub fn to_full_object_pathbuf() { - let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); - assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/media_upload/a/b/c/d/e/abcdefghijk/original_upload.bin")); + let file = WeightUploadOriginalFilePath::from_object_hash("abcdefghijk"); + assert_eq!(file.to_full_object_pathbuf(), PathBuf::from("/weight_upload/a/b/c/d/e/abcdefghijk/original_upload.bin")); } #[test] pub fn get_full_object_path_str_short_name() { - let file = MediaUploadOriginalFilePath::from_object_hash("foo"); - assert_eq!(file.get_full_object_path_str(), "/media_upload/f/o/foo/original_upload.bin"); + let file = WeightUploadOriginalFilePath::from_object_hash("foo"); + assert_eq!(file.get_full_object_path_str(), "/weight_upload/f/o/foo/original_upload.bin"); } #[test] pub fn get_full_object_path_str_starts_with_directory() { - let file = MediaUploadOriginalFilePath::from_object_hash("abcdefghijk"); + let file = WeightUploadOriginalFilePath::from_object_hash("abcdefghijk"); let full_path = file.get_full_object_path_str(); assert!(full_path.starts_with(file.get_directory().get_directory_path_str())) } @@ -98,7 +93,7 @@ mod tests { #[test] pub fn get_object_hash() { let hash = "abcdefghijk"; - let file = MediaUploadOriginalFilePath::from_object_hash(hash); + let file = WeightUploadOriginalFilePath::from_object_hash(hash); assert_eq!(file.get_object_hash(), hash); } } diff --git a/crates/_types/enums/src/by_table/media_uploads/media_upload_type.rs b/crates/_types/enums/src/by_table/media_uploads/media_upload_type.rs index 91f06ed283..42fd07ce41 100644 --- a/crates/_types/enums/src/by_table/media_uploads/media_upload_type.rs +++ b/crates/_types/enums/src/by_table/media_uploads/media_upload_type.rs @@ -20,6 +20,9 @@ pub enum MediaUploadType { /// Video files: mp4, etc. Video, + + // Binary files: safetensors ... weights, etc. + Binary, } // TODO(bt, 2022-12-21): This desperately needs MySQL integration tests! @@ -33,6 +36,7 @@ impl MediaUploadType { Self::Audio => "audio", Self::Image => "image", Self::Video => "video", + Self::Binary => "binary", } } @@ -41,6 +45,7 @@ impl MediaUploadType { "audio" => Ok(Self::Audio), "image" => Ok(Self::Image), "video" => Ok(Self::Video), + "binary" => Ok(Self::Binary), _ => Err(format!("invalid value: {:?}", value)), } } @@ -56,6 +61,7 @@ mod tests { assert_serialization(MediaUploadType::Audio, "audio"); assert_serialization(MediaUploadType::Image, "image"); assert_serialization(MediaUploadType::Video, "video"); + assert_serialization(MediaUploadType::Binary, "binary"); } #[test] @@ -63,6 +69,7 @@ mod tests { assert_eq!(MediaUploadType::Audio.to_str(), "audio"); assert_eq!(MediaUploadType::Image.to_str(), "image"); assert_eq!(MediaUploadType::Video.to_str(), "video"); + assert_eq!(MediaUploadType::Binary.to_str(), "binary"); } #[test] @@ -70,6 +77,7 @@ mod tests { assert_eq!(MediaUploadType::from_str("audio").unwrap(), MediaUploadType::Audio); assert_eq!(MediaUploadType::from_str("image").unwrap(), MediaUploadType::Image); assert_eq!(MediaUploadType::from_str("video").unwrap(), MediaUploadType::Video); + assert_eq!(MediaUploadType::from_str("binary").unwrap(), MediaUploadType::Binary); assert!(MediaUploadType::from_str("foo").is_err()); } } diff --git a/crates/lib/mysql_queries/src/queries/model_weights/create_weight.rs b/crates/lib/mysql_queries/src/queries/model_weights/create_weight.rs index d21f8a789e..4827edd907 100644 --- a/crates/lib/mysql_queries/src/queries/model_weights/create_weight.rs +++ b/crates/lib/mysql_queries/src/queries/model_weights/create_weight.rs @@ -7,6 +7,7 @@ use enums::by_table::model_weights::{ use log::warn; use enums::common::visibility::Visibility; use tokens::tokens::{ users::UserToken, model_weights::ModelWeightToken }; + pub struct CreateModelWeightsArgs<'a> { pub token: &'a ModelWeightToken, pub weights_type: WeightsType, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs index 1db5217ba1..6b9194135f 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/create_weight.rs @@ -3,17 +3,19 @@ use std::sync::Arc; use actix_multipart::Multipart; use actix_web::{HttpRequest, HttpResponse, web}; +use enums::by_table::model_weights::weights_category::WeightsCategory; +use enums::by_table::model_weights::weights_types::WeightsType; use mysql_queries::queries::model_weights::create_weight::{self, create_weight}; use once_cell::sync::Lazy; -use tokens::tokens::media_uploads::MediaUploadToken; +use tokens::tokens::model_weights::ModelWeightToken; use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError; use crate::server_state::ServerState; use log::{error, info, warn}; -use buckets::public::media_uploads::original_file::MediaUploadOriginalFilePath; +use buckets::public::weight_uploads::original_file::WeightUploadOriginalFilePath; use enums::by_table::media_uploads::media_upload_source::MediaUploadSource; use enums::by_table::media_uploads::media_upload_type::MediaUploadType; use enums::common::visibility::Visibility; @@ -23,17 +25,14 @@ use media::decode_basic_audio_info::decode_basic_audio_bytes_info; use mimetypes::mimetype_for_bytes::get_mimetype_for_bytes; use mysql_queries::queries::media_uploads::get_media_upload_by_uuid::get_media_upload_by_uuid_with_connection; use mysql_queries::queries::media_uploads::insert_media_upload::{Args, insert_media_upload}; -use tokens::tokens::media_uploads::MediaUploadToken; use crate::http_server::endpoints::media_uploads::common::drain_multipart_request::{drain_multipart_request, MediaSource}; -use crate::http_server::endpoints::media_uploads::common::upload_error::UploadError; -use crate::server_state::ServerState; use crate::validations::validate_idempotency_token_format::validate_idempotency_token_format; #[derive(Serialize)] pub struct UploadWeightsSuccessResponse { pub success: bool, - pub upload_token: MediaUploadToken, + pub upload_token: ModelWeightToken, } static ALLOWED_MIME_TYPES : Lazy> = Lazy::new(|| { @@ -53,8 +52,10 @@ pub async fn upload_weights_handler( &server_state, multipart_payload, &ALLOWED_MIME_TYPES).await?; + + // send a body with the required payload along side the multipart upload - let weight_upload_token = response.to_media_token(); + let weight_upload_token = response.to_model_token(); let response: UploadWeightsSuccessResponse = UploadWeightsSuccessResponse { success: true, @@ -69,18 +70,17 @@ pub async fn upload_weights_handler( .body(body)); } - pub enum SuccessCase { WeightAlreadyUploaded { - existing_upload_token: WeightUploadToken, + existing_upload_token: ModelWeightToken, }, WeightSuccessfullyUploaded { - upload_token: WeightUploadToken, + upload_token: ModelWeightToken, } } impl SuccessCase { - pub fn to_media_token(self) -> WeightUploadToken { + pub fn to_model_token(self) -> ModelWeightToken { match self { SuccessCase::WeightAlreadyUploaded { existing_upload_token } => existing_upload_token, SuccessCase::WeightSuccessfullyUploaded { upload_token } => upload_token, @@ -190,6 +190,7 @@ pub async fn handle_weight_upload( .as_ref() .map(|bytes| get_mimetype_for_bytes(bytes.as_ref())) .flatten(); + let bytes = match upload_weights_request.file_bytes { None => return Err(UploadError::BadInput("missing file contents".to_string())), @@ -219,38 +220,14 @@ pub async fn handle_weight_upload( .collect::(); return Err(UploadError::BadInput(format!("unpermitted mime type: {}", &filtered_mimetype))); } - - // NB: .aiff (audio/aiff) isn't supported by Symphonia: - // It contains uncompressed PCM-encoded audio similar to wav. - // See: https://github.com/pdeljanov/Symphonia/issues/75 - // NB: The following formats are not supported by Symphonia and - // do not have any open issues filed. They may simply be too old: - // - .wma (audio/x-ms-wma) - // - .avi (video/x-msvideo) media_upload_type = match mimetype { - // Audio - "audio/aac" /* .aac */ => Some(MediaUploadType::Audio), + "binary/octet-stream" /* .bin */ => Some(MediaUploadType::Binary), _ => None, }; - - - let media_upload_type = match media_upload_type { - Some(m) => m, - None => { - warn!("Invalid mimetype: {:?}", maybe_mimetype); - return Err(UploadError::BadInput(format!("unknown mimetype: {:?}", maybe_mimetype))) - }, - }; - - // probably don't need this - let media_upload_source = match upload_weights_request.media_source { - MediaSource::Unknown => MediaUploadSource::Unknown, - MediaSource::UserFile => MediaUploadSource::File, - MediaSource::UserDeviceApi => MediaUploadSource::DeviceApi, - }; + } // This is fine for now we just don't expose the path! - let public_upload_path = MediaUploadOriginalFilePath::generate_new(); + let public_upload_path = WeightUploadOriginalFilePath::generate_new(); info!("Uploading media to bucket path: {}", public_upload_path.get_full_object_path_str()); @@ -266,18 +243,18 @@ pub async fn handle_weight_upload( })?; let token = create_weight(CreateModelWeightsArgs { - token: &MediaUploadToken::generate(), - weights_type: media_upload_type, - weights_category: media_upload_category, + token: &ModelWeightToken::generate(), + weights_type: WeightsType::LoRA, // to do parse multi part for this + weights_category: WeightsCategory::ImageGeneration, title: upload_weights_request.file_name.as_deref(), - maybe_thumbnail_token: None, - description_markdown: None, - description_rendered_html: None, - creator_user_token: maybe_user_token.as_deref(), + maybe_thumbnail_token: &ThumbnailToken::generate(), + description_markdown: "", + description_rendered_html: "", + creator_user_token: maybe_user_token.as_str(), creator_ip_address: &ip_address, creator_set_visibility, - maybe_last_update_user_token: None, - original_download_url: None, + maybe_last_update_user_token: "", + original_download_url: "", original_filename: upload_weights_request.file_name.as_deref(), file_size_bytes: file_size_bytes as i32, file_checksum_sha2: hash, @@ -290,44 +267,13 @@ pub async fn handle_weight_upload( maybe_cached_user_ratings_ratio: None, version: 0, mysql_pool: &server_state.mysql_pool, - }) - .await - .map_err(|err| { - warn!("New generic download creation DB error: {:?}", err); - UploadError::ServerError - })?; - }).await?; - - - - let (token, record_id) = insert_media_upload(Args{ - uuid_idempotency_token: &uuid_idempotency_token, - media_type: media_upload_type, - media_source: media_upload_source, - maybe_original_filename: upload_weights_request.file_name.as_deref(), - original_file_size_bytes: file_size_bytes as u64, - maybe_original_duration_millis: maybe_duration_millis, - maybe_original_mime_type: maybe_mimetype, - maybe_original_audio_encoding: maybe_codec_name.as_deref(), - maybe_original_video_encoding: None, - maybe_original_frame_width: None, // TODO - maybe_original_frame_height: None, // TODO - checksum_sha2: &hash, - public_upload_path: &public_upload_path, - maybe_extra_file_modification_info: None, - maybe_creator_user_token: maybe_user_token.as_ref(), - maybe_creator_anonymous_visitor_token: None, - creator_ip_address: &ip_address, - creator_set_visibility, - mysql_pool: &server_state.mysql_pool, - }) - .await - .map_err(|err| { - warn!("New generic download creation DB error: {:?}", err); + }).await + .map_err(|e| { + error!("Error creating weights upload: {:?}", e); UploadError::ServerError })?; - info!("new media upload id: {} token: {:?}", record_id, &token); + info!("new weights upload id: {} token: {:?}", record_id, &token); Ok(SuccessCase::WeightSuccessfullyUploaded { upload_token: token, diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs new file mode 100644 index 0000000000..c87a9d82c5 --- /dev/null +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/drain_multipart_weights_request.rs @@ -0,0 +1,100 @@ +use actix_multipart::Multipart; +use actix_web::web::BytesMut; +use futures::TryStreamExt; +use log::warn; + +use errors::AnyhowResult; +use tokens::tokens::weights::ModelWeightToken; + +use crate::http_server::web_utils::read_multipart_field_bytes::{checked_read_multipart_bytes, read_multipart_field_as_text}; + +pub struct WeightUploadRequest { + pub uuid_idempotency_token: Option, + pub model_weights_token: Option, + pub file_name: Option, + pub file_bytes: Option, + pub media_source: MediaSource, +} + +/// Where the frontend tells us the file came from. +#[derive(Eq, PartialEq, Clone, Copy)] +pub enum MediaSource { + Unknown, + UserFile, + /// Eg. the web audio API in Javascript. + UserDeviceApi, +} + +/// Pull common parts out of multipart media HTTP requests, typically for handling file uploads. +pub async fn drain_multipart_request(mut multipart_payload: Multipart) -> AnyhowResult { + let mut uuid_idempotency_token = None; + let mut model_weight_token = None; + let mut file_bytes = None; + let mut file_name = None; + let mut media_source = None; + + while let Ok(Some(mut field)) = multipart_payload.try_next().await { + let mut field_name = None; + let mut field_filename = None; + + if let content_disposition = field.content_disposition() { + field_name = content_disposition.get_name() + .map(|s| s.to_string()); + field_filename = content_disposition.get_filename() // NB: Only used for the file bytes. + .map(|s| s.to_string()); + } + + match field_name.as_deref() { + Some("uuid_idempotency_token") => { + uuid_idempotency_token = read_multipart_field_as_text(&mut field).await + .map_err(|e| { + warn!("Error reading idempotency token: {:}", e); + e + })?; + }, + Some("dataset_token") => { + model_weight_token = read_multipart_field_as_text(&mut field).await + .map_err(|e| { + warn!("Error reading dataset token: {:}", e); + e + })?; + }, + Some("file") => { + file_name = field_filename.clone(); + file_bytes = checked_read_multipart_bytes(&mut field).await + .map_err(|e| { + warn!("Error reading audio upload: {:}", e); + e + })?; + }, + Some("source") => { + media_source = read_multipart_field_as_text(&mut field).await + .map_err(|e| { + warn!("Error reading source: {:}", e); + e + })?; + }, + _ => continue, + } + } + + let media_source = match media_source.as_deref() { + Some("device") => MediaSource::UserDeviceApi, + Some("file") => MediaSource::UserFile, + _ => { + if file_name.as_deref() == Some("blob") { + MediaSource::UserDeviceApi + } else { + MediaSource::Unknown + } + }, + }; + + Ok(UploadSampleRequest { + uuid_idempotency_token, + dataset_token, + file_name, + file_bytes, + media_source, + }) +} diff --git a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs index b9e2bd05c2..ba1f88eb2a 100644 --- a/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs +++ b/crates/service/web/storyteller_web/src/http_server/endpoints/weights/mod.rs @@ -3,4 +3,5 @@ pub mod delete_weight; pub mod update_weight; pub mod list_weights_by_user; pub mod list_availible_weights; -pub mod create_weight; \ No newline at end of file +pub mod create_weight; +pub mod drain_multipart_weights_request; \ No newline at end of file