[cm] Adding categorical parameters

Neutone · Mar 4, 2024 · c9554ab · c9554ab
1 parent 1e672ba
commit c9554ab
Show file tree

Hide file tree

Showing 12 changed files with 127 additions and 52 deletions.
diff --git a/README.md b/README.md
@@ -46,7 +46,7 @@ The Neutone Plugin is available at [https://neutone.space](https://neutone.space
 If you just want to wrap a model without going through a detailed description of what everything does we prepared these examples for you.
 
 - The clipper example shows how to wrap a very simple PyTorch module that does not contain any AI model. Check it out for getting a high level overview of what is needed for wrapping a model. It is available at [examples/example_clipper.py](examples/example_clipper.py).
-- An example with a simple convolutional model based on [Randomized Overdrive Neural Networks](https://csteinmetz1.github.io/ronn/) can be found at [examples/example_overdrive-random.py](examples/example_overdrive-random.py).
+- An example with a simple convolutional model based on [micro-tcn](https://github.com/csteinmetz1/micro-tcn) can be found at [examples/example_overdrive-random.py](examples/example_overdrive-random.py). This does not contain the training code and assumes you have a pretrained `micro-tcn` model.
 - We also have Notebooks for more complicated models showing the entire workflow from training to exporting them using Neutone:
     - [TCN FX Emulation](https://colab.research.google.com/drive/1gHZ-AEoYmfmWrjlKpKkK_SW1xzfxD24-?usp=sharing)
     - [DDSP Timbre Transfer](https://colab.research.google.com/drive/1yPHU6PRWw1lRWZLUxXimIa6chFQ2JdRW?usp=sharing)

diff --git a/examples/example_clipper.py b/examples/example_clipper.py
@@ -8,7 +8,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.utils import save_neutone_model
 
 logging.basicConfig()
@@ -59,9 +59,9 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter("min", "min clip threshold", default_value=0.15),
-            KnobNeutoneParameter("max", "max clip threshold", default_value=0.15),
-            KnobNeutoneParameter("gain", "scale clip threshold", default_value=1.0),
+            ContinuousNeutoneParameter("min", "min clip threshold", default_value=0.15),
+            ContinuousNeutoneParameter("max", "max clip threshold", default_value=0.15),
+            ContinuousNeutoneParameter("gain", "scale clip threshold", default_value=1.0),
         ]
 
     @tr.jit.export

diff --git a/examples/example_clipper_prefilter.py b/examples/example_clipper_prefilter.py
@@ -6,7 +6,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.filters import FIRFilter, FilterType
 from neutone_sdk.utils import save_neutone_model
 
@@ -65,9 +65,9 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter("min", "min clip threshold", default_value=0.15),
-            KnobNeutoneParameter("max", "max clip threshold", default_value=0.15),
-            KnobNeutoneParameter("gain", "scale clip threshold", default_value=1.0),
+            ContinuousNeutoneParameter("min", "min clip threshold", default_value=0.15),
+            ContinuousNeutoneParameter("max", "max clip threshold", default_value=0.15),
+            ContinuousNeutoneParameter("gain", "scale clip threshold", default_value=1.0),
         ]
 
     @tr.jit.export

diff --git a/examples/example_overdrive-random.py b/examples/example_overdrive-random.py
@@ -10,7 +10,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.tcn_1d import FiLM
 from neutone_sdk.utils import save_neutone_model
 
@@ -202,9 +202,9 @@ def get_citation(self) -> str:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter("depth", "Effect Depth", 0.0),
-            KnobNeutoneParameter("P1", "Feature modulation 1", 0.0),
-            KnobNeutoneParameter("P2", "Feature modulation 2", 0.0),
+            ContinuousNeutoneParameter("depth", "Effect Depth", 0.0),
+            ContinuousNeutoneParameter("P1", "Feature modulation 1", 0.0),
+            ContinuousNeutoneParameter("P2", "Feature modulation 2", 0.0),
         ]
 
     @torch.jit.export

diff --git a/examples/example_rave.py b/examples/example_rave.py
@@ -8,7 +8,7 @@
 import torchaudio
 from torch import Tensor
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.audio import (
     AudioSample,
     AudioSamplePair,
@@ -60,20 +60,20 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Chaos", description="Magnitude of latent noise", default_value=0.0
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z edit index",
                 description="Index of latent dimension to edit",
                 default_value=0.0,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z scale",
                 description="Scale of latent variable",
                 default_value=0.5,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z offset",
                 description="Offset of latent variable",
                 default_value=0.5,

diff --git a/examples/example_rave_prefilter.py b/examples/example_rave_prefilter.py
@@ -8,7 +8,7 @@
 import torchaudio
 from torch import Tensor, nn
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.audio import (
     AudioSample,
     AudioSamplePair,
@@ -69,20 +69,20 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Chaos", description="Magnitude of latent noise", default_value=0.0
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z edit index",
                 description="Index of latent dimension to edit",
                 default_value=0.0,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z scale",
                 description="Scale of latent variable",
                 default_value=0.5,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z offset",
                 description="Offset of latent variable",
                 default_value=0.5,

diff --git a/examples/example_rave_v1_prefilter.py b/examples/example_rave_v1_prefilter.py
@@ -8,7 +8,7 @@
 import torchaudio
 from torch import Tensor, nn
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.audio import (
     AudioSample,
     AudioSamplePair,
@@ -67,22 +67,22 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Chaos",
                 description="Magnitude of latent noise",
                 default_value=0.0,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z edit index",
                 description="Index of latent dimension to edit",
                 default_value=0.0,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z scale",
                 description="Scale of latent variable",
                 default_value=0.5,
             ),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 name="Z offset",
                 description="Offset of latent variable",
                 default_value=0.5,

diff --git a/examples/example_spectral_filter.py b/examples/example_spectral_filter.py
@@ -8,7 +8,7 @@
 import torch.nn as nn
 from torch import Tensor
 
-from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, KnobNeutoneParameter
+from neutone_sdk import WaveformToWaveformBase, NeutoneParameter, ContinuousNeutoneParameter
 from neutone_sdk.realtime_stft import RealtimeSTFT
 from neutone_sdk.utils import save_neutone_model
 
@@ -171,11 +171,11 @@ def is_experimental(self) -> bool:
 
     def get_neutone_parameters(self) -> List[NeutoneParameter]:
         return [
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter(
                 "center", "center frequency of the filter", default_value=0.3
             ),
-            KnobNeutoneParameter("width", "width of the filter", default_value=0.5),
-            KnobNeutoneParameter(
+            ContinuousNeutoneParameter("width", "width of the filter", default_value=0.5),
+            ContinuousNeutoneParameter(
                 "amount", "spectral attenuation amount", default_value=0.9
             ),
         ]

diff --git a/neutone_sdk/constants.py b/neutone_sdk/constants.py
@@ -3,6 +3,8 @@
 SDK_VERSION = "1.4.3"
 
 MAX_N_PARAMS = 4
+MAX_N_CATEGORICAL_VALUES = 20
+MAX_N_CATEGORICAL_LABEL_CHARS = 20
 MAX_N_AUDIO_SAMPLES = 3
 
 DEFAULT_DAW_SR = 48000

diff --git a/neutone_sdk/metadata.py b/neutone_sdk/metadata.py
@@ -135,7 +135,7 @@
             "properties": {
                 "name": {"type": "string"},
                 "description": {"type": "string"},
-                "type": {"type": "string", "enum": ["knob"]},
+                "type": {"type": "string", "enum": ["continuous"]},
                 "used": {"type": "string", "enum": ["True", "False"]},
                 "default_value": {"type": "string"},
             },

diff --git a/neutone_sdk/parameter.py b/neutone_sdk/parameter.py
@@ -2,15 +2,18 @@
 import os
 from abc import ABC
 from enum import Enum
-from typing import Dict, Union
+from typing import Dict, List, Union, Optional
+
+from neutone_sdk import constants
 
 logging.basicConfig()
 log = logging.getLogger(__name__)
 log.setLevel(level=os.environ.get("LOGLEVEL", "INFO"))
 
 
 class NeutoneParameterType(Enum):
-    KNOB = "knob"
+    CONTINUOUS = "continuous"
+    CATEGORICAL = "categorical"
     TEXT = "text"
 
 
@@ -19,7 +22,7 @@ class NeutoneParameter(ABC):
     Defines a Neutone Parameter abstract base class.
 
     The name and the description of the parameter will be shown as a tooltip
-    within the UI. This parameter has no functionality.
+    within the UI. This parameter has no functionality and is meant to subclassed.
     """
 
     def __init__(
@@ -37,6 +40,7 @@ def __init__(
         self.type = param_type
 
     def to_metadata_dict(self) -> Dict[str, str]:
+        """Returns a string dictionary containing the metadata of the parameter."""
         return {
             "name": self.name,
             "description": self.description,
@@ -46,13 +50,14 @@ def to_metadata_dict(self) -> Dict[str, str]:
         }
 
 
-class KnobNeutoneParameter(NeutoneParameter):
+class ContinuousNeutoneParameter(NeutoneParameter):
     """
-    Defines a knob Neutone Parameter that the user can use to control a model.
+    Defines a continuous Neutone Parameter that the user can use to control a model.
 
     The name and the description of the parameter will be shown as a tooltip
-    within the UI. `default_value` must be between 0 and 1 and will be used
-    as a default in the plugin when no presets are available.
+    within the UI.
+    `default_value` must be between 0 and 1 and will be used as a default in the plugin
+    when no presets are available.
     """
 
     def __init__(
@@ -63,8 +68,65 @@ def __init__(
             description,
             default_value,
             used,
-            NeutoneParameterType.KNOB,
+            NeutoneParameterType.CONTINUOUS,
         )
+        assert (
+            0.0 <= default_value <= 1.0
+        ), "`default_value` for continuous params must be between 0 and 1"
+
+
+class CategoricalNeutoneParameter(NeutoneParameter):
+    """
+    Defines a categorical Neutone Parameter that the user can use to control a model.
+
+    The name and the description of the parameter will be shown as a tooltip
+    within the UI.
+    `n_values` must be an int greater than or equal to 2 and less than or equal to
+    `constants.MAX_N_CATEGORICAL_VALUES`.
+    `default_value` must be in the range [0, `n_values` - 1].
+    `labels` is a list of strings that will be used as the labels for the parameter.
+
+    The categorical parameter is considered to be a knob parameter.
+    """
+
+    def __init__(
+        self,
+        name: str,
+        description: str,
+        n_values: int,
+        default_value: int,
+        labels: Optional[List[str]] = None,
+        used: bool = True,
+    ):
+        super().__init__(
+            name, description, default_value, used, NeutoneParameterType.CATEGORICAL
+        )
+        assert 2 <= n_values <= constants.MAX_N_CATEGORICAL_VALUES, (
+            f"`n_values` for categorical params must between 2 and "
+            f"{constants.MAX_N_CATEGORICAL_VALUES}"
+        )
+        assert (
+            0 <= default_value <= n_values - 1
+        ), "`default_value` for categorical params must be between 0 and `n_values`-1"
+        self.n_values = n_values
+        if labels is None:
+            labels = [str(idx) for idx in range(n_values)]
+        else:
+            assert len(labels) == self.n_values, "labels must have `n_values` elements"
+        assert all(
+            len(label) < constants.MAX_N_CATEGORICAL_LABEL_CHARS for label in labels
+        ), (
+            f"All labels must have length less than "
+            f"{constants.MAX_N_CATEGORICAL_LABEL_CHARS} characters"
+        )
+        self.labels = labels
+
+    def to_metadata_dict(self) -> Dict[str, str]:
+        """Returns a string dictionary containing the metadata of the parameter."""
+        data = super().to_metadata_dict()
+        data["n_values"] = str(self.n_values)
+        data["labels"] = "\t".join(self.labels)
+        return data
 
 
 class TextNeutoneParameter(NeutoneParameter):

diff --git a/neutone_sdk/wavform_to_wavform.py b/neutone_sdk/wavform_to_wavform.py
@@ -5,8 +5,12 @@
 import torch as tr
 from torch import Tensor, nn
 
-from neutone_sdk import NeutoneModel, constants, NeutoneParameterType, \
-    KnobNeutoneParameter
+from neutone_sdk import (
+    NeutoneModel,
+    constants,
+    NeutoneParameterType,
+    ContinuousNeutoneParameter,
+)
 from neutone_sdk.queues import CircularInplaceTensorQueue
 from neutone_sdk.utils import validate_waveform
 
@@ -55,20 +59,27 @@ def __init__(self, model: nn.Module, use_debug_mode: bool = True) -> None:
         self.agg_params = tr.zeros((self.MAX_N_PARAMS, 1))
 
         assert all(
-            p.type == NeutoneParameterType.KNOB for p in self.get_neutone_parameters()
-        ), "Only knob type parameters are supported in WaveformToWaveformBase models."
+            p.type == NeutoneParameterType.CONTINUOUS
+            for p in self.get_neutone_parameters()
+        ), (
+            "Only continuous type parameters are supported in WaveformToWaveformBase "
+            "models."
+        )
 
         # For compatibility with the current plugin, we fill in missing params
         # TODO(cm): remove once plugin metadata parsing is implemented
         for idx in range(self.n_neutone_parameters, self.MAX_N_PARAMS):
-            tmp_p = KnobNeutoneParameter(
-                name="", description="", default_value=0.0, used=False,
+            unused_p = ContinuousNeutoneParameter(
+                name="",
+                description="",
+                default_value=0.0,
+                used=False,
             )
-            self.neutone_parameters_metadata[f"p{idx + 1}"] = tmp_p.to_metadata_dict()
-            self.neutone_parameter_names.append(tmp_p.name)
-            self.neutone_parameter_descriptions.append(tmp_p.description)
-            self.neutone_parameter_types.append(tmp_p.type.value)
-            self.neutone_parameter_used.append(tmp_p.used)
+            self.neutone_parameters_metadata[f"p{idx+1}"] = unused_p.to_metadata_dict()
+            self.neutone_parameter_names.append(unused_p.name)
+            self.neutone_parameter_descriptions.append(unused_p.description)
+            self.neutone_parameter_types.append(unused_p.type.value)
+            self.neutone_parameter_used.append(unused_p.used)
 
     def _get_max_n_params(self) -> int:
         """