fix: labeling

feat: load vit model config.json

Cargo.lock

@@ -453,8 +453,8 @@ checksum = "8318a53db07bb3f8dca91a600466bdb3f2eaadeedfdbcf02e1accbad9271ba50"
 
 [[package]]
 name = "candle-core"
-version = "0.7.2"
+version = "0.8.1"
-source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.7.2#3a3c48b14b861da0d275cae44edb1b0191f83bf4"
+source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.8.1#67cab7d6b8279f953b0a8cc5012b135b9743cdc8"
 dependencies = [
  "byteorder",
  "gemm",
@@ -467,14 +467,15 @@ dependencies = [
  "rayon",
  "safetensors",
  "thiserror",
+ "ug",
  "yoke",
  "zip",
 ]
 
 [[package]]
 name = "candle-nn"
-version = "0.7.2"
+version = "0.8.1"
-source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.7.2#3a3c48b14b861da0d275cae44edb1b0191f83bf4"
+source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.8.1#67cab7d6b8279f953b0a8cc5012b135b9743cdc8"
 dependencies = [
  "candle-core",
  "half",
@@ -487,8 +488,8 @@ dependencies = [
 
 [[package]]
 name = "candle-transformers"
-version = "0.7.2"
+version = "0.8.1"
-source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.7.2#3a3c48b14b861da0d275cae44edb1b0191f83bf4"
+source = "git+https://git.v0l.io/huggingface/candle.git?tag=0.8.1#67cab7d6b8279f953b0a8cc5012b135b9743cdc8"
 dependencies = [
  "byteorder",
  "candle-core",
@@ -2149,6 +2150,30 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "num"
+version = "0.4.3"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "35bd024e8b2ff75562e5f34e7f4905839deb4b22955ef5e73d2fea1b9813cb23"
+dependencies = [
+ "num-bigint",
+ "num-complex",
+ "num-integer",
+ "num-iter",
+ "num-rational",
+ "num-traits",
+]
+
+[[package]]
+name = "num-bigint"
+version = "0.4.6"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "a5e44f723f1133c9deac646763579fdb3ac745e418f2a7af9cd0c431da1f20b9"
+dependencies = [
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-bigint-dig"
 version = "0.8.4"
@@ -2202,6 +2227,17 @@ dependencies = [
  "num-traits",
 ]
 
+[[package]]
+name = "num-rational"
+version = "0.4.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "f83d14da390562dca69fc84082e73e548e1ad308d24accdedd2720017cb37824"
+dependencies = [
+ "num-bigint",
+ "num-integer",
+ "num-traits",
+]
+
 [[package]]
 name = "num-traits"
 version = "0.2.19"
@@ -4104,6 +4140,19 @@ version = "0.1.6"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "ed646292ffc8188ef8ea4d1e0e0150fb15a5c2e12ad9b8fc191ae7a8a7f3c4b9"
 
+[[package]]
+name = "ug"
+version = "0.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "c4eef2ebfc18c67a6dbcacd9d8a4d85e0568cc58c82515552382312c2730ea13"
+dependencies = [
+ "half",
+ "num",
+ "serde",
+ "serde_json",
+ "thiserror",
+]
+
 [[package]]
 name = "uncased"
 version = "0.9.10"

Cargo.toml

@@ -47,8 +47,8 @@ clap = { version = "4.5.18", features = ["derive"] }
 
 libc = { version = "0.2.153", optional = true }
 ffmpeg-rs-raw = { git = "https://git.v0l.io/Kieran/ffmpeg-rs-raw.git", rev = "df69b2f05da4279e36ad55086d77b45b2caf5174", optional = true }
-candle-core = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.7.2", optional = true }
+candle-core = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.8.1", optional = true }
-candle-nn = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.7.2", optional = true }
+candle-nn = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.8.1", optional = true }
-candle-transformers = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.7.2", optional = true }
+candle-transformers = { git = "https://git.v0l.io/huggingface/candle.git", tag = "0.8.1", optional = true }
 sqlx-postgres = { version = "0.8.2", optional = true, features = ["chrono", "uuid"] }
 mime2ext = "0.1.53"

config.toml

@@ -17,7 +17,10 @@ public_url = "http://localhost:8000"
 # whitelist = ["63fe6318dc58583cfe16810f86dd09e18bfd76aabc24a0081ce2856f330504ed"]
 
 # Path for ViT(224) image model (https://huggingface.co/google/vit-base-patch16-224)
-# vit_model_path = "model.safetennsors"
+#vit_model_path = "falcon_nsfw.safetensors"
+
+# Path for ViT(224) image model config.json
+#vit_model_config_path = "falcon_nsfw.json"
 
 # Webhook api endpoint
 # webhook_url = "https://api.snort.social/api/v1/media/webhook"

src/filesystem.rs

@@ -109,9 +109,14 @@ impl FileStore {
 
                 #[cfg(feature = "labels")]
                 let labels = if let Some(mp) = &self.settings.vit_model_path {
-                    label_frame(&new_temp.result, mp.clone())?
+                    let config_path = if let Some(c) = &self.settings.vit_model_config_path {
+                        c
+                    } else {
+                        bail!("Missing vit_model_config_path");
+                    };
+                    label_frame(&new_temp.result, mp.clone(), config_path.clone())?
                         .iter()
-                        .map(|l| FileLabel::new(l.clone(), "vit224".to_string()))
+                        .map(|l| FileLabel::new(l.0.clone(), "vit224".to_string()))
                         .collect()
                 } else {
                     vec![]

src/settings.rs

@@ -24,6 +24,9 @@ pub struct Settings {
     /// Path for ViT image model
     pub vit_model_path: Option<PathBuf>,
 
+    /// Path for ViT image model config.json
+    pub vit_model_config_path: Option<PathBuf>,
+
     /// Webhook api endpoint
     pub webhook_url: Option<String>,