[cm] Updating non realtime wrapper with class typing

Author: christhetree

neutone_sdk/non_realtime_wrapper.py

@@ -6,7 +6,7 @@
 import torch as tr
 from torch import Tensor, nn
 
-from neutone_sdk import NeutoneModel, constants, NeutoneParameter, NeutoneParameterType
+from neutone_sdk import NeutoneModel, constants, NeutoneParameterType
 from neutone_sdk.utils import validate_waveform
 
 logging.basicConfig()
@@ -49,6 +49,16 @@ class NonRealtimeBase(NeutoneModel):
         NeutoneParameterType.CATEGORICAL,
         NeutoneParameterType.TEXT,
     }
+    # TorchScript typing does not support instance attributes, so we need to type them
+    # as class attributes. This is required for supporting models with no parameters.
+    # (https://github.com/pytorch/pytorch/issues/51041#issuecomment-767061194)
+    cont_param_names: List[str]
+    cont_param_indices: List[int]
+    cat_param_names: List[str]
+    cat_param_indices: List[int]
+    cat_param_n_values: Dict[str, int]
+    text_param_max_n_chars: List[int]
+    text_param_default_values: List[str]
 
     def __init__(self, model: nn.Module, use_debug_mode: bool = True) -> None:
         """
@@ -73,8 +83,6 @@ def __init__(self, model: nn.Module, use_debug_mode: bool = True) -> None:
         self.text_param_max_n_chars = []
         self.text_param_default_values = []
 
-        numerical_default_values = []
-
         # We have to keep track of this manually since text params are separate
         numerical_param_idx = 0
         for p in self.get_neutone_parameters():
@@ -86,33 +94,18 @@ def __init__(self, model: nn.Module, use_debug_mode: bool = True) -> None:
                 self.n_cont_params += 1
                 self.cont_param_names.append(p.name)
                 self.cont_param_indices.append(numerical_param_idx)
-                numerical_default_values.append(p.default_value)
                 numerical_param_idx += 1
             elif p.type == NeutoneParameterType.CATEGORICAL:
                 self.n_cat_params += 1
                 self.cat_param_names.append(p.name)
                 self.cat_param_indices.append(numerical_param_idx)
                 self.cat_param_n_values[p.name] = p.n_values
-                # Convert to float since default value tensor must be all the same type
-                numerical_default_values.append(float(p.default_value))
                 numerical_param_idx += 1
             elif p.type == NeutoneParameterType.TEXT:
                 self.n_text_params += 1
                 self.text_param_max_n_chars.append(p.max_n_chars)
                 self.text_param_default_values.append(p.default_value)
 
-        # This is needed for TorchScript typing since it doesn't allow empty lists etc.
-        if not self.n_cont_params:
-            self.cont_param_names.append("__torchscript_typing")
-            self.cont_param_indices.append(-1)
-        if not self.n_cat_params:
-            self.cat_param_names.append("__torchscript_typing")
-            self.cat_param_indices.append(-1)
-            self.cat_param_n_values["__torchscript_typing"] = -1
-        if not self.n_text_params:
-            self.text_param_max_n_chars.append(-1)
-            self.text_param_default_values.append("__torchscript_typing")
-
         self.n_numerical_params = self.n_cont_params + self.n_cat_params
 
         assert self.n_numerical_params <= constants.NEUTONE_GEN_N_NUMERICAL_PARAMS, (
@@ -126,63 +119,38 @@ def __init__(self, model: nn.Module, use_debug_mode: bool = True) -> None:
         if self.n_text_params:
             self.has_text_param = True
 
-        all_neutone_parameters = self._get_all_neutone_parameters()
-        assert len(all_neutone_parameters) == self._get_max_n_params()
-
         # This overrides the base class definitions to remove the text param or extra
         # base param since it is handled separately in the UI.
         # TODO(cm): this if statement will be removed once we get rid of the extra
         # core methods we don't need anymore
         if self.has_text_param:
             self.neutone_parameter_names = [
                 p.name
-                for p in all_neutone_parameters
+                for p in self.get_neutone_parameters()
                 if p.type != NeutoneParameterType.TEXT
             ]
             self.neutone_parameter_descriptions = [
                 p.description
-                for p in all_neutone_parameters
+                for p in self.get_neutone_parameters()
                 if p.type != NeutoneParameterType.TEXT
             ]
             self.neutone_parameter_types = [
                 p.type.value
-                for p in all_neutone_parameters
+                for p in self.get_neutone_parameters()
                 if p.type != NeutoneParameterType.TEXT
             ]
             self.neutone_parameter_used = [
                 p.used
-                for p in all_neutone_parameters
+                for p in self.get_neutone_parameters()
                 if p.type != NeutoneParameterType.TEXT
             ]
-        else:
-            self.neutone_parameter_names = self.neutone_parameter_names[
-                : constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
-            ]
-            self.neutone_parameter_descriptions = self.neutone_parameter_descriptions[
-                : constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
-            ]
-            self.neutone_parameter_types = self.neutone_parameter_types[
-                : constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
-            ]
-            self.neutone_parameter_used = self.neutone_parameter_used[
-                : constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
-            ]
 
         # TODO(cm): this statement will also be removed once core is refactored
         assert (
             len(self.get_default_param_names())
             == constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
         )
 
-        numerical_param_default_values_t = Tensor(numerical_default_values)
-        if self.n_numerical_params > 0:
-            numerical_param_default_values_t = (
-                numerical_param_default_values_t.unsqueeze(1)
-            )
-        assert numerical_param_default_values_t.size(0) == self.n_numerical_params
-        # TODO(cm): rename once moved from core
-        self.register_buffer("default_param_values", numerical_param_default_values_t)
-
         assert all(
             1 <= n <= 2 for n in self.get_audio_in_channels()
         ), "Input audio channels must be mono or stereo"
@@ -212,23 +180,23 @@ def _get_max_n_params(self) -> int:
             + constants.NEUTONE_GEN_N_TEXT_PARAMS
         )
 
-    def _get_all_neutone_parameters(self) -> List[NeutoneParameter]:
+    def _create_default_param_values(self) -> Tensor:
         """
-        Returns a list of NeutoneParameters that is equal to the maximum number of
-        parameters that the model can have. Placeholder unused NeutoneParameters are
-        created if the model has less than the maximum number of parameters.
-        This should not be overwritten by SDK users.
+        Creates the default parameter values tensor, which must be 1-dimensional.
+        For NonRealtimeBase models, this is a tensor of the default values for the
+        continuous and categorical (numerical) parameters and ignores the default
+        values for the text parameters since these are handled separately.
         """
-        neutone_parameters = self.get_neutone_parameters()
-        if len(neutone_parameters) < self._get_max_n_params():
-            neutone_parameters += [
-                NeutoneParameter(
-                    name="",
-                    description="",
-                    used=False,
-                )
-            ] * (self._get_max_n_params() - len(neutone_parameters))
-        return neutone_parameters
+        numerical_default_values = []
+        for p in self.get_neutone_parameters():
+            if p.type == NeutoneParameterType.CONTINUOUS:
+                numerical_default_values.append(p.default_value)
+            elif p.type == NeutoneParameterType.CATEGORICAL:
+                # Convert to float to match the type of the continuous parameters
+                numerical_default_values.append(float(p.default_value))
+        assert len(numerical_default_values) == self.n_numerical_params
+        numerical_default_values = tr.tensor(numerical_default_values)
+        return numerical_default_values
 
     @abstractmethod
     def get_audio_in_channels(self) -> List[int]:
@@ -424,11 +392,7 @@ def forward(
         This method should not be overwritten by SDK users.
         """
         if text_params is None:
-            # Needed for TorchScript typing
-            if self.n_text_params:
-                text_params = self.text_param_default_values
-            else:
-                text_params = []
+            text_params = self.text_param_default_values
 
         if self.use_debug_mode:
             assert len(audio_in) == len(self.get_audio_in_channels())
@@ -455,13 +419,7 @@ def forward(
 
         if self.use_debug_mode:
             if numerical_params is not None:
-                assert numerical_params.ndim == 2
-                assert (
-                    self.n_numerical_params
-                    <= numerical_params.size(0)
-                    <= constants.NEUTONE_GEN_N_NUMERICAL_PARAMS
-                )
-                assert numerical_params.size(1) == in_n
+                assert numerical_params.shape == (self.n_numerical_params, in_n)
             if not self.is_one_shot_model() and self.get_native_buffer_sizes():
                 assert (
                     in_n in self.get_native_buffer_sizes()