72 lines
2.2 KiB
Rust
72 lines
2.2 KiB
Rust
use std::sync::Arc;
|
|
|
|
use galaxyui::{Entity, ModelContext, SingletonEntity};
|
|
use input_classifier::{HeuristicClassifier, InputClassifier};
|
|
#[cfg(any(
|
|
feature = "nld_classifier_v1",
|
|
feature = "nld_classifier_v2",
|
|
feature = "nld_classifier_v3"
|
|
))]
|
|
use input_classifier::{OnnxClassifier, OnnxModel};
|
|
|
|
pub struct InputClassifierModel {
|
|
pub classifier: Arc<dyn InputClassifier>,
|
|
}
|
|
|
|
impl InputClassifierModel {
|
|
pub fn new(_ctx: &mut ModelContext<Self>) -> Self {
|
|
#[cfg(feature = "nld_classifier_v1")]
|
|
{
|
|
match OnnxClassifier::new(OnnxModel::BertTinyV1) {
|
|
Ok(classifier) => {
|
|
log::info!("Loaded onnx classifier bert_tiny_v1.onnx");
|
|
return Self {
|
|
classifier: Arc::new(classifier),
|
|
};
|
|
}
|
|
Err(e) => log::warn!("Failed to load onnx classifier bert_tiny_v1.onnx: {e:#}"),
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "nld_classifier_v2")]
|
|
{
|
|
match OnnxClassifier::new(OnnxModel::BertTinyV2) {
|
|
Ok(classifier) => {
|
|
log::info!("Loaded onnx classifier bert_tiny_v2.onnx");
|
|
return Self {
|
|
classifier: Arc::new(classifier),
|
|
};
|
|
}
|
|
Err(e) => log::warn!("Failed to load onnx classifier bert_tiny_v2.onnx: {e:#}"),
|
|
}
|
|
}
|
|
|
|
#[cfg(feature = "nld_classifier_v3")]
|
|
{
|
|
match OnnxClassifier::new(OnnxModel::BertTinyV3) {
|
|
Ok(classifier) => {
|
|
log::info!("Loaded onnx classifier bert_tiny_v3.onnx");
|
|
return Self {
|
|
classifier: Arc::new(classifier),
|
|
};
|
|
}
|
|
Err(e) => log::warn!("Failed to load onnx classifier bert_tiny_v3.onnx: {e:#}"),
|
|
}
|
|
}
|
|
|
|
Self {
|
|
classifier: Arc::new(HeuristicClassifier),
|
|
}
|
|
}
|
|
|
|
pub fn classifier(&self) -> Arc<dyn InputClassifier> {
|
|
self.classifier.clone()
|
|
}
|
|
}
|
|
|
|
impl Entity for InputClassifierModel {
|
|
type Event = ();
|
|
}
|
|
|
|
impl SingletonEntity for InputClassifierModel {}
|