fix: module import fails due to TensorflowPredictEffnetDiscogs not fo…

…und (#4)
okio-ai · Dec 31, 2023 · 8f51cbe · 8f51cbe
1 parent 88907d6
commit 8f51cbe
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 6 deletions.
diff --git a/README.md b/README.md
@@ -19,6 +19,17 @@ Automatic music information retrieval (based on [essentia](https://essentia.upf.
 - Use descriptive features to filter, search and sort your library
 - Extract rich features to annotate datasets for training custom models
 
+## Requirements
+
+This plugin requires `tensorflow`. Please install it [as explained in its official documentation](https://www.tensorflow.org/install/pip):
+
+```bash
+# For GPU users
+pip install tensorflow[and-cuda]
+# For CPU users
+pip install tensorflow
+```
+
 ## Installation
 
 1. [Install nendo](https://github.com/okio-ai/nendo#installation)

diff --git a/pyproject.toml b/pyproject.toml
@@ -35,8 +35,7 @@ classifiers = [
 python = "^3.8,<3.11"
 pydantic = { version = "^2.0.0, <2.5.0" }
 nendo = "^0.1.2"
-essentia = "^2.1b6.dev1034"
-essentia-tensorflow= "^2.1b6.dev1110"
+essentia-tensorflow = "^2.1b6.dev1110"
 tqdm = "^4.65.0"
 
 [tool.poetry.group.lint.dependencies]

diff --git a/src/nendo_plugin_classify_core/plugin.py b/src/nendo_plugin_classify_core/plugin.py
@@ -54,10 +54,10 @@ class NendoClassifyCore(NendoAnalysisPlugin):
 
     nendo_instance: Nendo
     config: NendoConfig = None
-    embedding_model: Any = None
-    mood_model: Any = None
-    genre_model: Any = None
-    instrument_model: Any = None
+    embedding_model: es.TensorflowPredictEffnetDiscogs = None
+    mood_model: es.TensorflowPredict2D = None
+    genre_model: es.TensorflowPredict2D = None
+    instrument_model: es.TensorflowPredict2D = None
 
     def __init__(self, **data: Any):
         """Initialize plugin."""