Merge remote-tracking branch 'origin/master' into feature/comfyui

Cargo.lock

@@ -2149,6 +2149,7 @@ dependencies = [
  "errors",
  "log",
  "my-workspace-hack",
+ "regex",
  "serde",
  "serde_json",
  "tokens",

crates/cli/elasticsearch-cli/src/main.rs

@@ -14,6 +14,7 @@ use sqlx::mysql::MySqlPoolOptions;
 use config::shared_constants::DEFAULT_RUST_LOG;
 use elasticsearch_schema::searches::search_model_weights::search_model_weights;
 use elasticsearch_schema::searches::search_tts_models::search_tts_models;
+use enums::by_table::model_weights::weights_types::WeightsType;
 use errors::AnyhowResult;
 
 use crate::cli_args::{Action, Environment, parse_cli_args};
@@ -55,9 +56,10 @@ pub async fn main() -> AnyhowResult<()> {
     }
     Action::SearchModelWeights => {
       info!("Searching model weights...");
-      let results = search_model_weights(&elasticsearch, "zel", Some("en")).await?;
+      //let results = search_model_weights(&elasticsearch, "zel", Some("en"), Some(WeightsType::SoVitsSvc)).await?;
+      let results = search_model_weights(&elasticsearch, "zel", None, Some(WeightsType::RvcV2), None).await?;
       for result in results {
-        println!("Result: {:?}", result);
+        println!("Result: {:#?}", result);
       }
     }
   }

crates/lib/elasticsearch_schema/Cargo.toml

@@ -27,6 +27,7 @@ chrono = { version = "0.4.22", features = ["serde"] }
 serde = { version = "1.0.152", features = ["derive"] }
 serde_json = "1.0.108"
 my-workspace-hack = { version = "0.1", path = "../../my-workspace-hack" }
+regex = "1.10.2"
 
 [dev-dependencies]
-# None yet
+# None yet

crates/lib/elasticsearch_schema/src/searches/search_model_weights.rs

@@ -1,22 +1,125 @@
 use elasticsearch::{Elasticsearch, SearchParts};
-use log::error;
 use serde_json::{json, Value};
 
-use errors::AnyhowResult;
+use enums::by_table::model_weights::weights_category::WeightsCategory;
+use enums::by_table::model_weights::weights_types::WeightsType;
+use errors::{anyhow, AnyhowResult};
 
 use crate::documents::model_weight_document::{MODEL_WEIGHT_INDEX, ModelWeightDocument};
 
+pub struct SearchModelWeightsQuery<'a> {
+  pub search_term: &'a str,
+  pub maybe_language_subtag: Option<&'a str>,
+  pub maybe_weights_type: Option<WeightsType>,
+  pub maybe_weights_category: Option<WeightsCategory>,
+}
+
+impl <'a>SearchModelWeightsQuery<'a> {
+
+  fn base_query(&self) -> Value {
+    json!({
+      "query": {
+        "bool": {
+          "must": [
+            {
+              "bool": {
+                "should": [
+                  {
+                    "fuzzy": {
+                      "title": {
+                        "value": self.search_term,
+                        "fuzziness": 2
+                      }
+                    }
+                  },
+                  {
+                    "match": {
+                      "title": {
+                        "query": self.search_term,
+                        "boost": 1
+                      }
+                    }
+                  },
+                  {
+                    "multi_match": {
+                      "query": self.search_term,
+                      "type": "bool_prefix",
+                      "fields": [
+                        "title",
+                        "title._2gram",
+                        "title._3gram"
+                      ],
+                      "boost": 50
+                    }
+                  }
+                ]
+              }
+            }
+          ]
+        }
+      }
+    })
+  }
+
+  pub fn query(&self) -> AnyhowResult<Value> {
+    let mut search_json = self.base_query();
+
+    let must_clause = search_json.pointer_mut("/query/bool/must")
+        .map(|pointer| pointer.as_array_mut())
+        .flatten()
+        .ok_or_else(|| anyhow!("could not get pointer to must clause"))?;
+
+    if let Some(maybe_language_subtag) = &self.maybe_language_subtag {
+      must_clause.push(json!(
+        {
+          "match": {
+            "maybe_ietf_primary_language_subtag": maybe_language_subtag
+          }
+        }
+      ));
+    }
+
+    if let Some(weights_type) = &self.maybe_weights_type {
+      must_clause.push(json!(
+        {
+          "match": {
+            "weights_type": weights_type.to_string(),
+          }
+        }
+      ));
+    }
+
+    if let Some(weights_category) = &self.maybe_weights_category {
+      must_clause.push(json!(
+        {
+          "match": {
+            "weights_category": weights_category.to_string(),
+          }
+        }
+      ));
+    }
+
+    Ok(search_json)
+  }
+}
+
 pub async fn search_model_weights(
   client: &Elasticsearch,
   search_term: &str,
   maybe_language_subtag: Option<&str>,
+  maybe_weights_type: Option<WeightsType>,
+  maybe_weights_category: Option<WeightsCategory>,
 ) -> AnyhowResult<Vec<ModelWeightDocument>> {
 
-  let search_json = match maybe_language_subtag {
-    None => query_model_weights(search_term),
-    Some(language) => query_model_weights_with_required_language(search_term, language)
+  let query = SearchModelWeightsQuery {
+    search_term,
+    maybe_language_subtag,
+    maybe_weights_type,
+    maybe_weights_category,
   };
 
+  let search_json = query.query()?;
+
   let search_response = client
       .search(SearchParts::Index(&[MODEL_WEIGHT_INDEX]))
       .body(search_json)
@@ -44,11 +147,7 @@ pub async fn search_model_weights(
         let maybe_object = hit.get_mut("_source")
             .map(|source| source.take());
         if let Some(value) = maybe_object {
-          let document = serde_json::from_value::<ModelWeightDocument>(value)
-              .map_err(|err| {
-                error!("Error: {:?}", &err);
-                err
-              })?;
+          let document = serde_json::from_value::<ModelWeightDocument>(value)?;
           documents.push(document);
         }
       }
@@ -59,97 +158,164 @@ pub async fn search_model_weights(
   Ok(documents)
 }
 
-fn query_model_weights(search_term: &str) -> Value {
-  json!({
-    "query": {
-      "bool": {
-        "must": [
-          {
-            "bool": {
-              "should": [
-                {
-                  "fuzzy": {
-                    "title": {
-                      "value": search_term,
-                      "fuzziness": 2
-                    }
-                  }
-                },
-                {
-                  "match": {
-                    "title": {
-                      "query": search_term,
-                      "boost": 1
-                    }
-                  }
-                },
-                {
-                  "multi_match": {
-                    "query": search_term,
-                    "type": "bool_prefix",
-                    "fields": [
-                      "title",
-                      "title._2gram",
-                      "title._3gram"
-                    ],
-                    "boost": 50
-                  }
-                }
-              ]
-            }
-          }
-        ]
-      }
-    }
-  })
-}
+#[cfg(test)]
+mod tests {
+  use regex::Regex;
 
-fn query_model_weights_with_required_language(search_term: &str, language_tag: &str) -> Value {
-  json!({
-    "query": {
-      "bool": {
-        "must": [
-          {
-            "match": {
-              "maybe_ietf_primary_language_subtag": language_tag
-            }
-          },
-          {
-            "bool": {
-              "should": [
-                {
-                  "fuzzy": {
-                    "title": {
-                      "value": search_term,
-                      "fuzziness": 2
+  use enums::by_table::model_weights::weights_category::WeightsCategory;
+  use enums::by_table::model_weights::weights_types::WeightsType;
+  use errors::AnyhowResult;
+
+  use crate::searches::search_model_weights::SearchModelWeightsQuery;
+
+  fn query_to_json_string(query: SearchModelWeightsQuery<'_>) -> AnyhowResult<String> {
+    let json = query.query()?;
+    let json = serde_json::to_string(&json)?;
+    Ok(json)
+  }
+
+  fn compact_json(json: &str) -> String {
+    let regex = Regex::new("\\s+").expect("regex should parse");
+    let json = regex.replace_all(&json, "");
+    json.to_string()
+  }
+
+  #[test]
+  fn default_search_only_keyword() {
+    let query = SearchModelWeightsQuery {
+      search_term: "FOO_BAR_BAZ",
+      maybe_language_subtag: None,
+      maybe_weights_type: None,
+      maybe_weights_category: None,
+    };
+
+    let json = query_to_json_string(query).unwrap();
+
+    // NB: Keys in emitted JSON are sorted.
+    let expected_json = r#"
+      {
+        "query": {
+          "bool": {
+            "must": [
+              {
+                "bool": {
+                  "should": [
+                    {
+                      "fuzzy": {
+                        "title": {
+                          "fuzziness": 2,
+                          "value": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "match": {
+                        "title": {
+                          "boost": 1,
+                          "query": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "multi_match": {
+                        "boost": 50,
+                        "fields": [
+                          "title",
+                          "title._2gram",
+                          "title._3gram"
+                        ],
+                        "query": "FOO_BAR_BAZ",
+                        "type": "bool_prefix"
+                      }
                     }
-                  }
-                },
-                {
-                  "match": {
-                    "title": {
-                      "query": search_term,
-                      "boost": 1
-                    }
-                  }
-                },
-                {
-                  "multi_match": {
-                    "query": search_term,
-                    "type": "bool_prefix",
-                    "fields": [
-                      "title",
-                      "title._2gram",
-                      "title._3gram"
-                    ],
-                    "boost": 50
-                  }
+                  ]
                 }
-              ]
-            }
+              }
+            ]
           }
-        ]
+        }
       }
-    }
-  })
+    "#;
+
+    let expected_json = compact_json(&expected_json);
+
+    assert_eq!(&json, &expected_json);
+  }
+
+  #[test]
+  fn search_with_language_and_weights() {
+    let query = SearchModelWeightsQuery {
+      search_term: "FOO_BAR_BAZ",
+      maybe_language_subtag: Some("en"),
+      maybe_weights_type: Some(WeightsType::SoVitsSvc),
+      maybe_weights_category: Some(WeightsCategory::Vocoder),
+    };
+
+    let json = query_to_json_string(query).unwrap();
+
+    // NB: Keys in emitted JSON are sorted.
+    let expected_json = r#"
+      {
+        "query": {
+          "bool": {
+            "must": [
+              {
+                "bool": {
+                  "should": [
+                    {
+                      "fuzzy": {
+                        "title": {
+                          "fuzziness": 2,
+                          "value": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "match": {
+                        "title": {
+                          "boost": 1,
+                          "query": "FOO_BAR_BAZ"
+                        }
+                      }
+                    },
+                    {
+                      "multi_match": {
+                        "boost": 50,
+                        "fields": [
+                          "title",
+                          "title._2gram",
+                          "title._3gram"
+                        ],
+                        "query": "FOO_BAR_BAZ",
+                        "type": "bool_prefix"
+                      }
+                    }
+                  ]
+                }
+              },
+              {
+                "match": {
+                  "maybe_ietf_primary_language_subtag": "en"
+                }
+              },
+              {
+                "match": {
+                  "weights_type": "so_vits_svc"
+                }
+              },
+              {
+                "match": {
+                  "weights_category": "vocoder"
+                }
+              }
+            ]
+          }
+        }
+      }
+    "#;
+
+    let expected_json = compact_json(&expected_json);
+
+    assert_eq!(&json, &expected_json);
+  }
 }

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

@@ -247,7 +247,7 @@ mod tests {
     };
 
     let json = serde_json::ser::to_string(&args).unwrap();
-    assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"vs":"video_media_token","is":"image_media_token","sd":"sd_model_token","lm":"lora_model_token","p":"prompt","ap":"a_prompt","np":"n_prompt","se":1,"mu":"upload_path","lu":"lora_upload_path"}}}"#.to_string());
+    assert_eq!(json, r#"{"cat":"ig","args":{"Ig":{"sd":"sd_model_token","lm":"lora_model_token","w":512,"h":512,"s":"sampler","p":"prompt","np":"n_prompt","se":1,"mu":"upload_path","cf":7,"lu":"lora_upload_path","sa":25,"bc":1,"t":"type_of_inference","de":"Option","na":"Model Name"}}}"#.to_string());
   }
 
   #[test]

crates/lib/mysql_queries/src/queries/model_weights/tests.rs

@@ -293,8 +293,6 @@ mod tests {
             maybe_scoped_weight_type: None,
             maybe_scoped_weight_category: None,
             mysql_pool: &pool,
-            maybe_scoped_weight_type: None,
-            maybe_scoped_weight_category: None,
         }).await?;
         for weight in weights_by_username.records.iter() {
             // print weight

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

@@ -427,6 +427,7 @@ pub async fn process_job_inference(
 
 }
 
+#[ignore]
 #[cfg(test)]
 mod tests {
     use std::path::PathBuf;
@@ -439,6 +440,7 @@ mod tests {
     };
     use errors::AnyhowResult;
 
+    #[ignore]
     #[tokio::test]
     async fn test_seed_weights_files() -> AnyhowResult<()> {
         let seed_path = PathBuf::from("/storyteller/root/custom-seed-tool-data");

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_dataset_samples/list_samples_by_dataset.rs

@@ -86,11 +86,15 @@ pub async fn list_samples_by_dataset_handler(
   };
 
   let dataset_token = ZsVoiceDatasetToken(path.dataset_token.clone());
-  let is_mod = user_session.can_ban_users;
+  let _is_mod = user_session.can_ban_users;
+
+  // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+  // We should focus on visibility-related controls instead.
+  const CAN_SEE_DELETED : bool = false;
 
   let query_results = list_dataset_samples_for_dataset_token(
     &dataset_token,
-    is_mod,
+    CAN_SEE_DELETED,
     &server_state.mysql_pool
   ).await.map_err(|e| {
     warn!("list_samples error: {:?}", e);

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_session.rs

@@ -77,12 +77,16 @@ pub async fn list_datasets_by_session_handler(
       })?;
 
 
-  let is_mod = user_session.can_ban_users;
+  let _is_mod = user_session.can_ban_users;
+
+  // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+  // We should focus on visibility-related controls instead.
+  const CAN_SEE_DELETED : bool = false;
 
   let query_results = list_datasets_by_username(
     &server_state.mysql_pool,
     &user_session.username,
-    is_mod,
+    CAN_SEE_DELETED,
   ).await.map_err(|e| {
     warn!("Error querying for datasets: {:?}", e);
     ListDatasetsByUserError::ServerError

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voice_datasets/list_datasets_by_user.rs

@@ -89,12 +89,16 @@ pub async fn list_datasets_by_user_handler(
 
   let username = path.username.as_ref();
   let creator_user_token = user_session.user_token.clone();
-  let is_mod = user_session.can_ban_users;
+  let _is_mod = user_session.can_ban_users;
+
+  // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+  // We should focus on visibility-related controls instead.
+  const CAN_SEE_DELETED : bool = false;
 
   let query_results = list_datasets_by_username(
     &server_state.mysql_pool,
     &username,
-    is_mod,
+    CAN_SEE_DELETED,
   ).await.map_err(|e| {
     warn!("Error querying for datasets: {:?}", e);
     ListDatasetsByUserError::ServerError

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_available_voices.rs

@@ -110,13 +110,15 @@ pub async fn list_available_voices_handler(
 
     let include_user_hidden = is_mod;
 
+    // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+    const CAN_SEE_DELETED : bool = false;
 
     let mut query_builder = ListVoicesQueryBuilder::new()
         .sort_ascending(sort_ascending)
         .scope_creator_username(None)
         .include_user_hidden(include_user_hidden)
-        .include_user_deleted_results(is_mod)
-        .include_mod_deleted_results(is_mod)
+        .include_user_deleted_results(CAN_SEE_DELETED)
+        .include_mod_deleted_results(CAN_SEE_DELETED)
         .limit(limit)
         .cursor_is_reversed(cursor_is_reversed)
         .offset(cursor);

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_session.rs

@@ -76,12 +76,16 @@ pub async fn list_voices_by_session_handler(
       })?;
 
   let creator_user_token = user_session.user_token.clone();
-  let is_mod = user_session.can_ban_users;
+  let _is_mod = user_session.can_ban_users;
+
+  // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+  // We should focus on visibility-related controls instead.
+  const CAN_SEE_DELETED : bool = false;
 
   let query_results = list_zs_voices_by_username(
     &server_state.mysql_pool,
     &user_session.username,
-    is_mod,
+    CAN_SEE_DELETED,
   ).await.map_err(|e| {
     warn!("Error querying for voices: {:?}", e);
     ListVoicesByUserError::ServerError

crates/service/web/storyteller_web/src/http_server/endpoints/voice_designer/voices/list_voices_by_user.rs

@@ -89,12 +89,16 @@ pub async fn list_voices_by_user_handler(
 
   let username = path.username.as_ref();
   let creator_user_token = user_session.user_token.clone();
-  let is_mod = user_session.can_ban_users;
+  let _is_mod = user_session.can_ban_users;
+
+  // NB(bt,2024-01-18): Showing mods deleted files is actually kind of annoying!
+  // We should focus on visibility-related controls instead.
+  const CAN_SEE_DELETED : bool = false;
 
   let query_results = list_zs_voices_by_username(
     &server_state.mysql_pool,
     &username,
-    is_mod,
+    CAN_SEE_DELETED,
   ).await.map_err(|e| {
     warn!("Error querying for voices: {:?}", e);
     ListVoicesByUserError::ServerError

crates/service/web/storyteller_web/src/http_server/endpoints/weights/search_model_weights_handler.rs

@@ -30,6 +30,9 @@ use crate::server_state::ServerState;
 #[derive(Deserialize, ToSchema)]
 pub struct SearchModelWeightsRequest {
   pub search_term: String,
+  pub weight_type: Option<WeightsType>,
+  pub weight_category: Option<WeightsCategory>,
+  pub ietf_language_subtag: Option<String>,
 }
 
 #[derive(Serialize, Clone, ToSchema)]
@@ -113,7 +116,9 @@ pub async fn search_model_weights_handler(
   let results = search_model_weights(
     &server_state.elasticsearch,
     &request.search_term,
-    None)
+    request.ietf_language_subtag.as_deref(),
+    request.weight_type,
+    request.weight_category)
       .await
       .map_err(|err| {
         error!("Searching error: {:?}", err);

crates/service/web/storyteller_web/src/http_server/web_utils/read_multipart_field_bytes.rs

@@ -6,7 +6,8 @@ use futures::{StreamExt, TryStreamExt};
 use log::warn;
 
 const MIN_BYTES : usize = 10;
-const MAX_BYTES : usize = 1024 * 1024 * 20;
+//const MAX_BYTES : usize = 1024 * 1024 * 20; // 20 Mb
+const MAX_BYTES : usize = 1024 * 1024 * 125; // 125 Mb
 
 /// Read any field from a multipart form as bytes, but do some error checking.
 pub async fn checked_read_multipart_bytes(field: &mut Field) -> anyhow::Result<Option<BytesMut>> {