[FEATURE] Implement loading WaveNet model from .nam file (#565)

Implement loading WaveNet model from .nam file

Author: sdatkinson

nam/models/__init__.py

@@ -9,6 +9,7 @@
 from . import base  # noqa F401
 from . import exportable  # noqa F401
 from . import losses  # noqa F401
+from ._from_nam import init_from_nam  # noqa F401
 from .conv_net import ConvNet  # noqa F401
 from .linear import Linear  # noqa F401
 from .recurrent import LSTM  # noqa F401

nam/models/_from_nam.py

@@ -0,0 +1,31 @@
+# File: _from_nam.py
+# Created Date: Tuesday May 27th 2025
+# Author: Steven Atkinson (steven@atkinson.mn)
+
+"""
+Initialize models from .nam files
+"""
+
+import torch as _torch
+
+from .base import BaseNet as _BaseNet
+from .wavenet import WaveNet as _WaveNet
+
+def _init_wavenet(config) -> _WaveNet:
+    return _WaveNet(layers_configs=config["layers"], head_config=config["head"], head_scale=config["head_scale"])
+
+
+def init_from_nam(config) -> _BaseNet:
+    """
+    Taking the contents of a .nam file, initialize a model
+
+    E.g.
+    >>> with open("model.nam", "r") as fp:
+    ...     config = json.load(fp)
+    ...     model = init_from_nam(config)
+    """
+    model = {
+        "WaveNet": _init_wavenet
+    }[config["architecture"]](config["config"])
+    model.import_weights(_torch.Tensor(config["weights"]))
+    return model

nam/models/wavenet.py

@@ -327,6 +327,8 @@ def export_weights(self) -> _np.ndarray:
         return weights.detach().cpu().numpy()
 
     def import_weights(self, weights: _torch.Tensor):
+        if self._head is not None:
+            raise NotImplementedError("Head importing isn't implemented yet.")
         i = 0
         for layer in self._layers:
             i = layer.import_weights(weights, i)

tests/test_nam/test_models/test_from_nam.py

@@ -0,0 +1,91 @@
+"""
+Test loading from a .nam file
+"""
+
+import json as _json
+from pathlib import Path as _Path
+from tempfile import TemporaryDirectory as _TemporaryDirectory
+import pytest as _pytest
+
+from nam.models import _from_nam
+from nam.models.wavenet import WaveNet as _WaveNet
+
+
+@_pytest.mark.parametrize(
+    "factory,kwargs",
+    (
+        # A standard WaveNet
+        (
+            _WaveNet,  # i.e. .__init__()
+            {
+                "layers_configs": [
+                    {
+                        "condition_size": 1,
+                        "input_size": 1,
+                        "channels": 16,
+                        "head_size": 8,
+                        "kernel_size": 3,
+                        "dilations": [
+                            1,
+                            2,
+                            4,
+                            8,
+                            16,
+                            32,
+                            64,
+                            128,
+                            256,
+                            512
+                        ],
+                        "activation": "Tanh",
+                        "gated": False,
+                        "head_bias": False
+                    },
+                    {
+                        "condition_size": 1,
+                        "input_size": 16,
+                        "channels": 8,
+                        "head_size": 1,
+                        "kernel_size": 3,
+                        "dilations": [
+                            1,
+                            2,
+                            4,
+                            8,
+                            16,
+                            32,
+                            64,
+                            128,
+                            256,
+                            512
+                        ],
+                        "activation": "Tanh",
+                        "gated": False,
+                        "head_bias": True
+                    }
+                ],
+                "head_scale": 0.02
+            }
+        ),
+    ),
+)
+def test_load_from_nam(factory, kwargs):
+    """
+    Assert that loading from a .nam file works by saving the model twice
+    """
+    model = factory(**kwargs)
+    with _TemporaryDirectory() as tmpdir:
+        model.export(_Path(tmpdir), basename="model")
+        with open(_Path(tmpdir, "model.nam"), "r") as fp:
+            nam_file_contents = _json.load(fp)
+        model2 = _from_nam.init_from_nam(nam_file_contents)
+        model2.export(_Path(tmpdir), basename="model2")
+        with open(_Path(tmpdir, "model2.nam"), "r") as fp:
+            nam_file_contents2 = _json.load(fp)
+
+        # Metadata isn't preseved. At least creation time will be slightly different
+        # Could improve this to preserve metadata on load
+        nam_file_contents.pop("metadata")
+        nam_file_contents2.pop("metadata")
+        assert nam_file_contents == nam_file_contents2
+