Repository: ktonal/mimikit (language: python)

Issue:
Feature/scripts



Discussion / hints:

Fix patch:
diff --git a/.github/workflows/ci-pipeline.yml b/.github/workflows/ci-pipeline.yml
index 065c710..a7b1ef4 100644
--- a/.github/workflows/ci-pipeline.yml
+++ b/.github/workflows/ci-pipeline.yml
@@ -3,9 +3,6 @@ name: Upload Python Package
 on:
   workflow_dispatch:
   push:
-#    branches-ignore: [master, develop]
-  pull_request:
-    branches: [master, develop]
 
 jobs:
   ci-pipeline:
@@ -15,7 +12,7 @@ jobs:
     - name: Set up Python
       uses: actions/setup-python@v2
       with:
-        python-version: '3.6'
+        python-version: '3.7'
     - name: Install dependencies
       run: |
         sudo apt-get install -y libsndfile-dev
diff --git a/.gitignore b/.gitignore
index 5b2f5cf..0bf091e 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,4 +3,6 @@
 *.ipynb_checkpoints*
 dist/*
 build/*
+docs/_build/
 *egg-info*
+.idea
\ No newline at end of file
diff --git a/README.md b/README.md
index 325b086..7e4325b 100644
--- a/README.md
+++ b/README.md
@@ -1,18 +1,13 @@
 # mimikit
 
-`mimikit` is an open-source python package for audio and music data
-modeling with neural networks.
+The MusIc ModelIng toolKIT (MIMIKIT) is a python package that does Machine Learning with music data.
 
-It contains models out-of-the-box and simple utilities for dealing with audio data.
-The aim of mimikit is double : it will lower the barrier to deep-learning for non-developers while providing tools to developers that enable fast experimentation.
-
------
-
-## Usage 
-
-Check out the [mimikit-notebooks](https://github.com/ktonal/mimikit-notebooks) for examples of client code.
+The goal of `mimikit` is to enable you to use referenced and experimental algorithms on data you provide.
 
+`mimikit` is still in early development, details and documentation are on their way.
 
 ## License
 
 mimikit is distributed under the terms of the [GNU General Public License v3.0](https://choosealicense.com/licenses/gpl-3.0/)
+
+
diff --git a/docs/Makefile b/docs/Makefile
new file mode 100644
index 0000000..d4bb2cb
--- /dev/null
+++ b/docs/Makefile
@@ -0,0 +1,20 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line, and also
+# from the environment for the first two.
+SPHINXOPTS    ?=
+SPHINXBUILD   ?= sphinx-build
+SOURCEDIR     = .
+BUILDDIR      = _build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/docs/audios.rst b/docs/audios.rst
new file mode 100644
index 0000000..5fd6b94
--- /dev/null
+++ b/docs/audios.rst
@@ -0,0 +1,8 @@
+audios
+======
+
+this package exposes helper classes for processing, interacting & modeling audio data
+
+.. automodule:: mimikit.audios.fmodules
+
+.. automodule:: mimikit.audios.features
diff --git a/docs/conf.py b/docs/conf.py
new file mode 100644
index 0000000..f90abf6
--- /dev/null
+++ b/docs/conf.py
@@ -0,0 +1,79 @@
+# Configuration file for the Sphinx documentation builder.
+#
+# This file only contains a selection of the most common options. For a full
+# list see the documentation:
+# https://www.sphinx-doc.org/en/master/usage/configuration.html
+
+# -- Path setup --------------------------------------------------------------
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+sys.path.insert(0, os.path.abspath('../mimikit'))
+
+
+# -- Project information -----------------------------------------------------
+
+project = 'mimikit'
+copyright = '2021, k-tonal'
+author = 'k-tonal'
+
+# The full version, including alpha/beta/rc tags
+release = 'v0.1.6'
+
+
+# -- General configuration ---------------------------------------------------
+
+import sphinx_rtd_theme
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.doctest',
+    "sphinx.ext.viewcode",
+    'sphinx_rtd_theme',
+    "numpydoc",
+    # "sphinx.ext.napoleon",
+    # "sphinx.ext.autosummary"
+]
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path.
+exclude_patterns = []
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = 'sphinx_rtd_theme'
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ['_static']
+
+autoclass_content = "class"
+
+# autodoc_default_options = ['no-undoc-members', 'no-inherited-members']
+autodoc_mock_imports = ['soundfile', 'jsonschema.compat']
+autodoc_typehints = 'description'
+autodoc_class_signature = 'separated'
+autodoc_typehints_description_target = 'documented'
+autosummary_generate = True
+
+numpydoc_show_class_members = False
+
+napoleon_use_ivar = False
+
+napoleon_use_admonition_for_examples = False
\ No newline at end of file
diff --git a/docs/data.rst b/docs/data.rst
new file mode 100644
index 0000000..33739fb
--- /dev/null
+++ b/docs/data.rst
@@ -0,0 +1,37 @@
+data
+====
+
+.. py:currentmodule:: mimikit.data.database
+
+.. autoclass:: Database
+    :members:
+
+.. autoclass:: FeatureProxy
+    :members:
+
+.. py:currentmodule:: mimikit.data.datamodule
+
+.. autoclass:: Getter
+    :members:
+    :special-members: __call__
+
+
+.. autoclass:: AsSlice
+    :members:
+
+.. autoclass:: AsFramedSlice
+    :members:
+
+.. autoclass:: Input
+    :members:
+
+.. autoclass:: Target
+    :members:
+
+.. autoclass:: DefaultDataset
+    :members:
+
+.. autoclass:: DataModule
+    :no-inherited-members:
+
+
diff --git a/docs/index.rst b/docs/index.rst
new file mode 100644
index 0000000..f83edfd
--- /dev/null
+++ b/docs/index.rst
@@ -0,0 +1,29 @@
+.. mimikit documentation master file, created by
+   sphinx-quickstart on Fri Jan  8 04:24:18 2021.
+   You can adapt this file completely to your liking, but it should at least
+   contain the root `toctree` directive.
+
+mimikit
+=======
+
+.. toctree::
+   :maxdepth: 2
+   :caption: Contents:
+
+.. toctree::
+   :maxdepth: 1
+   :caption: User Guides
+
+
+.. toctree::
+   :maxdepth: 1
+   :caption: API Reference
+
+   audios
+   data
+   models
+   networks
+
+* :ref:`genindex`
+* :ref:`modindex`
+* :ref:`search`
diff --git a/docs/make.bat b/docs/make.bat
new file mode 100644
index 0000000..2119f51
--- /dev/null
+++ b/docs/make.bat
@@ -0,0 +1,35 @@
+@ECHO OFF
+
+pushd %~dp0
+
+REM Command file for Sphinx documentation
+
+if "%SPHINXBUILD%" == "" (
+	set SPHINXBUILD=sphinx-build
+)
+set SOURCEDIR=.
+set BUILDDIR=_build
+
+if "%1" == "" goto help
+
+%SPHINXBUILD% >NUL 2>NUL
+if errorlevel 9009 (
+	echo.
+	echo.The 'sphinx-build' command was not found. Make sure you have Sphinx
+	echo.installed, then set the SPHINXBUILD environment variable to point
+	echo.to the full path of the 'sphinx-build' executable. Alternatively you
+	echo.may add the Sphinx directory to PATH.
+	echo.
+	echo.If you don't have Sphinx installed, grab it from
+	echo.http://sphinx-doc.org/
+	exit /b 1
+)
+
+%SPHINXBUILD% -M %1 %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+goto end
+
+:help
+%SPHINXBUILD% -M help %SOURCEDIR% %BUILDDIR% %SPHINXOPTS% %O%
+
+:end
+popd
diff --git a/docs/models.rst b/docs/models.rst
new file mode 100644
index 0000000..0ef4df7
--- /dev/null
+++ b/docs/models.rst
@@ -0,0 +1,13 @@
+models
+======
+
+exposes well-known and experimental models as :class:`pytorch_lightning.LightningModule`
+
+.. py:currentmodule:: mimikit.models
+
+.. autoclass:: SampleRNN
+    :no-inherited-members:
+
+
+.. autoclass:: FreqNet
+    :no-inherited-members:
diff --git a/docs/networks.rst b/docs/networks.rst
new file mode 100644
index 0000000..7828df9
--- /dev/null
+++ b/docs/networks.rst
@@ -0,0 +1,15 @@
+networks
+========
+
+exposes networks as pure :class:`torch.nn.Module`.
+
+.. py:currentmodule:: mimikit.networks
+
+.. autoclass:: FreqNetNetwork
+    :no-inherited-members:
+
+.. autoclass:: SampleRNNNetwork
+    :no-inherited-members:
+
+.. autoclass:: WNNetwork
+    :no-inherited-members:
diff --git a/docs/requirements.txt b/docs/requirements.txt
new file mode 100644
index 0000000..0c1422c
--- /dev/null
+++ b/docs/requirements.txt
@@ -0,0 +1,3 @@
+sphinx
+numpydoc
+sphinx_rtd_theme
\ No newline at end of file
diff --git a/mimikit/__init__.py b/mimikit/__init__.py
index 696d661..94673af 100644
--- a/mimikit/__init__.py
+++ b/mimikit/__init__.py
@@ -1,6 +1,26 @@
-__version__ = '0.1.10'
+__version__ = '0.2.0'
+
+from . import audios
+from . import connectors
+from . import data
+from . import extract
+from . import midi
+from . import models
+from . import modules
+from . import networks
+
+from .audios import *
+from .connectors import *
+from .data import *
+from .extract import *
+from .midi import *
+from .models.parts import *
+from .models import *
+from .modules import *
+from .networks import *
+
+from .file_walker import *
+from .get_trainer import *
+from .utils import *
+
 
-from .kit import get_trainer
-from .connectors.neptune import NeptuneConnector
-from .data import Database
-from .utils import show, audio, signal
diff --git a/mimikit/_make_notebook.py b/mimikit/_make_notebook.py
new file mode 100644
index 0000000..c8e7a39
--- /dev/null
+++ b/mimikit/_make_notebook.py
@@ -0,0 +1,3 @@
+import nbformat as nbf
+
+__all__ = []
diff --git a/mimikit/abstract/__init__.py b/mimikit/abstract/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/mimikit/abstract/features.py b/mimikit/abstract/features.py
new file mode 100644
index 0000000..361fca4
--- /dev/null
+++ b/mimikit/abstract/features.py
@@ -0,0 +1,45 @@
+import dataclasses as dtc
+
+import numpy as np
+from ..data import Feature
+from ..extract import from_recurrence_matrix
+
+
+@dtc.dataclass
+class SegmentLabels(Feature):
+    __ext__ = 'none'
+
+    input_key: str = 'x'
+    L: int = 6
+    k: int = None
+    sym: bool = True
+    bandwidth: float = 1.
+    thresh: float = 0.2
+    min_dur: int = 4
+
+    @property
+    def params(self):
+        return {'input_key': self.input_key}
+
+    @property
+    def encoders(self):
+        return {np.ndarray: lambda X: \
+            from_recurrence_matrix(X.T, L=self.L, k=self.k, sym=self.sym,
+                                   bandwidth=self.bandwidth, thresh=self.thresh,
+                                   min_dur=self.min_dur)}
+
+    def post_create(self, db, schema_key):
+        return getattr(db, schema_key).regions.to_labels()
+
+
+@dtc.dataclass
+class FilesLabels(Feature):
+    __ext__ = 'none'
+    input_key: str = 'x'
+
+    @property
+    def encoders(self):
+        return {np.ndarray: lambda x: np.ones((x.shape[0]))}
+
+    def post_create(self, db, schema_key):
+        return getattr(db, schema_key).files.to_labels()
\ No newline at end of file
diff --git a/mimikit/audios/__init__.py b/mimikit/audios/__init__.py
new file mode 100644
index 0000000..c37f784
--- /dev/null
+++ b/mimikit/audios/__init__.py
@@ -0,0 +1,5 @@
+from .features import *
+from .fmodules import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
diff --git a/mimikit/audios/features.py b/mimikit/audios/features.py
new file mode 100644
index 0000000..0406e34
--- /dev/null
+++ b/mimikit/audios/features.py
@@ -0,0 +1,144 @@
+import torch
+import numpy as np
+import librosa
+import dataclasses as dtc
+import IPython.display as ipd
+import soundfile as sf
+
+from ..data import Feature
+from . import fmodules as T
+
+__all__ = [
+    'AudioSignal',
+    'MuLawSignal',
+    'Spectrogram'
+]
+
+
+@dtc.dataclass
+class AudioSignal(Feature):
+    """
+    audio signal managers
+    """
+    __ext__ = 'audio'
+
+    sr: int = 22050
+    normalize: bool = True
+    emphasis: float = 0.
+
+    @property
+    def dim(self):
+        return 1
+
+    @property
+    def encoders(self):
+        def func(inputs):
+            if self.emphasis:
+                inputs = T.Emphasis(self.emphasis)(inputs)
+            if self.normalize:
+                inputs = T.Normalize()(inputs)
+            return inputs
+
+        return {
+            np.ndarray: func,
+            torch.Tensor: func
+        }
+
+    @property
+    def decoders(self):
+        def func(inputs):
+            if self.emphasis:
+                inputs = T.Deemphasis(self.emphasis)(inputs)
+            if self.normalize:
+                inputs = T.Normalize()(inputs)
+            return inputs
+
+        return {
+            np.ndarray: func,
+            torch.Tensor: func
+        }
+
+    def load(self, path):
+        y = T.FileToSignal(self.sr)(path)
+        y = self.encode(y)
+        return y
+
+    def display(self, inputs, **waveplot_kwargs):
+        waveplot_kwargs.setdefault('sr', self.sr)
+        return librosa.display.waveplot(inputs, **waveplot_kwargs)
+
+    def play(self, inputs):
+        return ipd.display(ipd.Audio(inputs, rate=self.sr))
+
+    def write(self, filename, inputs):
+        return sf.write(filename, inputs, self.sr, 'PCM_24')
+
+
+@dtc.dataclass(unsafe_hash=True)
+class MuLawSignal(AudioSignal):
+    q_levels: int = 256
+
+    @property
+    def dim(self):
+        return self.q_levels
+
+    @property
+    def encoders(self):
+        return {
+            np.ndarray: T.MuLawCompress(q_levels=self.q_levels),
+            torch.Tensor: T.MuLawCompress(q_levels=self.q_levels)
+        }
+
+    @property
+    def decoders(self):
+        return {
+            np.ndarray: T.MuLawExpand(self.q_levels),
+            torch.Tensor: T.MuLawExpand(self.q_levels)
+        }
+
+
+@dtc.dataclass(unsafe_hash=True)
+class Spectrogram(AudioSignal):
+    n_fft: int = 2048
+    hop_length: int = 512
+    magspec: bool = False
+
+    @property
+    def dim(self):
+        return self.n_fft // 2 + 1
+
+    @property
+    def encoders(self):
+        if self.magspec:
+            return {
+                np.ndarray: T.MagSpec(self.n_fft, self.hop_length),
+                torch.Tensor: T.MagSpec(self.n_fft, self.hop_length)
+            }
+        return {
+            np.ndarray: T.STFT(self.n_fft, self.hop_length),
+            torch.Tensor: T.STFT(self.n_fft, self.hop_length)
+        }
+
+    @property
+    def decoders(self):
+        if self.magspec:
+            return {
+                np.ndarray: T.GLA(self.n_fft, self.hop_length, n_iter=32),
+                torch.Tensor: T.GLA(self.n_fft, self.hop_length, n_iter=32)
+            }
+        return {
+            np.ndarray: T.ISTFT(self.n_fft, self.hop_length),
+            torch.Tensor: T.ISTFT(self.n_fft, self.hop_length)
+        }
+
+    def display(self, inputs, **waveplot_kwargs):
+        y = self.decode(inputs)
+        return super(Spectrogram, self).display(y, **waveplot_kwargs)
+
+    def play(self, inputs):
+        y = self.decode(inputs)
+        return super(Spectrogram, self).play(y)
+
+    def write(self, filename, inputs):
+        y = self.decode(inputs)
+        return super(Spectrogram, self).write(filename, y)
diff --git a/mimikit/audios/fmodules.py b/mimikit/audios/fmodules.py
new file mode 100644
index 0000000..82b7b43
--- /dev/null
+++ b/mimikit/audios/fmodules.py
@@ -0,0 +1,233 @@
+import librosa
+import torch
+import torchaudio.functional as F
+import torchaudio.transforms as T
+import numpy as np
+from scipy.signal import lfilter
+import dataclasses as dtc
+
+__all__ = [
+    'FileToSignal',
+    'Normalize',
+    'Emphasis',
+    'Deemphasis',
+    'MuLawCompress',
+    'MuLawExpand',
+    'STFT',
+    'ISTFT',
+    'MagSpec',
+    'GLA'
+]
+
+N_FFT = 2048
+HOP_LENGTH = 512
+SR = 22050
+Q_LEVELS = 256
+
+
+class FModule:
+    """
+    base class for implementing message passing in callable objects
+    """
+
+    @property
+    def functions(self) -> dict:
+        raise NotImplementedError
+
+    def __call__(self, inputs):
+        return self.functions[type(inputs)](inputs)
+
+
+@dtc.dataclass
+class FileToSignal(FModule):
+    sr: int = SR
+
+    @property
+    def functions(self):
+        return {
+            str: lambda path: librosa.load(path, sr=self.sr)[0]
+        }
+
+
+@dtc.dataclass
+class Normalize(FModule):
+    p: int = 1
+    dim: int = -1
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            return librosa.util.normalize(inputs, norm=self.p, axis=self.dim)
+
+        def torch_func(inputs):
+            return torch.nn.functional.normalize(inputs, p=self.p, dim=self.dim)
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class Emphasis(FModule):
+    emphasis: float = 0.
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            return lfilter([1, -self.emphasis], [1], inputs)
+
+        def torch_func(inputs):
+            return F.lfilter(inputs,
+                             torch.tensor([1, 0]).to(inputs),  # a0, a1
+                             torch.tensor([1, -self.emphasis]).to(inputs))  # b0, b1
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class Deemphasis(FModule):
+    emphasis: float = 0.
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            return lfilter([1 - self.emphasis], [1, -self.emphasis], inputs)
+
+        def torch_func(inputs):
+            return F.lfilter(inputs,
+                             torch.tensor([1, - self.emphasis]).to(inputs),  # a0, a1
+                             torch.tensor([1 - self.emphasis, 0]).to(inputs))  # b0, b1
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class MuLawCompress(FModule):
+    q_levels: int = Q_LEVELS
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            qx = librosa.mu_compress(inputs, self.q_levels - 1, quantize=True)
+            qx = qx + self.q_levels // 2
+            return qx
+
+        def torch_func(inputs):
+            mod = T.MuLawEncoding(self.q_levels)
+            return mod(inputs)
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class MuLawExpand(FModule):
+    q_levels: int = Q_LEVELS
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            return librosa.mu_expand(inputs - self.q_levels // 2, self.q_levels - 1, quantize=True)
+
+        def torch_func(inputs):
+            mod = T.MuLawDecoding(self.q_levels)
+            return mod(inputs)
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class STFT(FModule):
+    n_fft: int = N_FFT
+    hop_length: int = HOP_LENGTH
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            # returned shape is (time x freq)
+            return librosa.stft(inputs, n_fft=self.n_fft, hop_length=self.hop_length).T
+
+        def torch_func(inputs):
+            mod = T.Spectrogram(self.n_fft, hop_length=self.hop_length, power=1.,
+                                wkwargs=dict(device=inputs.device))
+            # returned shape is (..., time x freq)
+            return mod(inputs).transpose(-1, -2).contiguous()
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class ISTFT(FModule):
+    n_fft: int = N_FFT
+    hop_length: int = HOP_LENGTH
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            # inputs is of shape (time x freq)
+            return librosa.istft(inputs.T, n_fft=self.n_fft, hop_length=self.hop_length, )
+
+        def torch_func(inputs):
+            # inputs is of shape (time x freq)
+            y = torch.istft(inputs.transpose(-1, -2).contiguous(),
+                            n_fft=self.n_fft, hop_length=self.hop_length,
+                            window=torch.hann_window(self.n_fft, device=inputs.device))
+            return y
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
+
+
+@dtc.dataclass
+class MagSpec(STFT):
+
+    @property
+    def functions(self):
+        sup_f = super(MagSpec, self).functions
+        # dict comprehension would result in a single function for
+        # all types, so we declare the dict manually...
+        return {
+            np.ndarray: lambda x: abs(sup_f[np.ndarray](x)),
+            torch.Tensor: lambda x: abs(sup_f[torch.Tensor](x))
+        }
+
+
+@dtc.dataclass
+class GLA(FModule):
+    n_fft: int = N_FFT
+    hop_length: int = HOP_LENGTH
+    n_iter: int = 32
+
+    @property
+    def functions(self):
+        def np_func(inputs):
+            # inputs is of shape (time x freq)
+            return librosa.griffinlim(inputs.T, hop_length=self.hop_length, n_iter=self.n_iter)
+
+        def torch_func(inputs):
+            gla = T.GriffinLim(n_fft=self.n_fft, hop_length=self.hop_length, power=1.,
+                               wkwargs=dict(device=inputs.device))
+            # inputs is of shape (time x freq)
+            return gla(inputs.transpose(-1, -2).contiguous())
+
+        return {
+            np.ndarray: np_func,
+            torch.Tensor: torch_func
+        }
diff --git a/mimikit/connectors/__init__.py b/mimikit/connectors/__init__.py
index e69de29..93817a0 100644
--- a/mimikit/connectors/__init__.py
+++ b/mimikit/connectors/__init__.py
@@ -0,0 +1,4 @@
+from .neptune import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/connectors/neptune.py b/mimikit/connectors/neptune.py
index cb5020c..9d95742 100644
--- a/mimikit/connectors/neptune.py
+++ b/mimikit/connectors/neptune.py
@@ -17,6 +17,10 @@ class pytorch_lightning:
 
 from mimikit.data import Database
 
+__all__ = [
+    'NeptuneConnector'
+]
+
 
 class NeptuneConnector:
     NEPTUNE_TOKEN_KEY = "NEPTUNE_API_TOKEN"
@@ -90,7 +94,7 @@ def __init__(self,
         self.setup = setup
         self._session = None
         # initialize the token right-away
-        token = self.api_token
+        _ = self.api_token
 
     def path(self, setup_key: str, split: bool = False):
         """
@@ -248,7 +252,7 @@ def upload_database(self, setup_key: str, db: Database):
             params = {"name": db.h5_file,
                       "feature_name": feature,
                       "shape": feat_prox.shape,
-                      "files": len(db.metadata)}
+                      }
             params.update(feat_prox.attrs)
             exp = self.create_experiment(setup_key, params=params)
         else:
diff --git a/mimikit/data/__init__.py b/mimikit/data/__init__.py
index 2e4e4ef..eed24f6 100644
--- a/mimikit/data/__init__.py
+++ b/mimikit/data/__init__.py
@@ -1,6 +1,9 @@
-from .api import Database, FeatureProxy
-from .data_object import DataObject
-from .factory import make_root_db, AudioFileWalker
-from .regions import Regions
-from .transforms import file_to_fft, default_extract_func, N_FFT, HOP_LENGTH, SR
-from .freqnet_db import freqnet_db
+from .create import *
+from .database import *
+from .datamodule import *
+from .feature import *
+from .regions import *
+from .tbptt_sampler import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/data/api.py b/mimikit/data/api.py
deleted file mode 100644
index 9037ef9..0000000
--- a/mimikit/data/api.py
+++ /dev/null
@@ -1,285 +0,0 @@
-import h5py
-import numpy as np
-import pandas as pd
-import os
-from datetime import datetime
-
-from torch.utils.data.dataset import Subset
-
-from .regions import Regions
-from ..data.data_object import DataObject
-
-
-class FeatureProxy(object):
-    """
-    Interface to the numerical data (array) stored in a .h5 file created with ``mimikit.data.make_root_db``
-
-    ``FeatureProxy`` objects are like ``numpy`` arrays (on disk) and can be used as map-style ``Dataset``
-    because they implement ``__len__`` and ``__getitem__``.
-
-    Parameters
-    ----------
-    h5_file : str
-        name of the file
-    ds_name : str
-        name of the ``h5py.Dataset`` containing the data to be served
-
-    Attributes
-    ----------
-    N : int
-        the length of the first dimension of the underlying array
-    shape : tuple of int
-        the shape of the underlying array
-    attrs : dict
-        dictionary of additional information about the data as returned
-        by the ``extract_func`` passed to ``make_root_db``.
-        Typically, this is where you want to store the parameters of your extracting function, e.g.
-        sample rate, hop length etc.
-    """
-
-    def __init__(self, h5_file: str, ds_name: str):
-
-        self.h5_file = h5_file
-        self.name = ds_name
-        with h5py.File(h5_file, "r") as f:
-            ds = f[self.name]
-            self.N = ds.shape[0]
-            self.shape = ds.shape
-            self.attrs = {k: v for k, v in ds.attrs.items()}
-
-    def __len__(self):
-        return self.N
-
-    def __getitem__(self, item):
-        with h5py.File(self.h5_file, "r") as f:
-            rv = f[self.name][item]
-        return rv
-
-    def get(self, regions):
-        """
-        get the data (numpy array) corresponding to the rows of `regions`
-
-        Parameters
-        ----------
-        regions : Regions
-            the files you want to get as array
-
-        Returns
-        -------
-        data : np.ndarray
-            the concatenated data for all the requested files
-
-        Examples
-        --------
-        >>>from mimikit import Database
-        >>>db = Database("my-db.h5")
-        # only get files 0, 3 & 7
-        >>>db.fft.get(db.regions.iloc[[0, 3, 7]])
-
-        """
-        t_axis = self.attrs.get("time_axis", 0)
-        slices = regions.slices(t_axis)
-        return np.concatenate(tuple(self[slice_i] for slice_i in slices), axis=t_axis)
-
-    def add(self, array, filename=None):
-        """
-        EXPERIMENTAL! append `array` to the feature and fill the `"name"` column of `db.regions` with 'filename'
-        Parameters
-        ----------
-        array : np.ndarray
-            the array to append at the end of the feature
-        filename : str, optional
-            a name for the array being added
-
-        Returns
-        -------
-        new : FeatureProxy
-            the updated object
-        """
-        new = _add_data(self.h5_file, self.name, array, filename)
-        return new
-
-    def subset(self, indices):
-        """
-        transform self into a torch ``Subset`` (``Dataset``) containing only ``indices`` from the original data
-
-        Parameters
-        ----------
-        indices : ``Regions`` or any object that ``Subset`` accepts as indices
-            if ``indices`` is of type ``Regions``, the returned ``Subset`` will only contain the files (rows) present
-            in ``indices``.
-
-        Returns
-        -------
-        subset : torch.utils.data.Subset
-            the obtained subset
-        """
-        if isinstance(indices, Regions):
-            indices = indices.all_indices
-        return Subset(DataObject(self), indices)
-
-    def regions(self, indices):
-        """
-        transform self into a torch ``Subset`` containing the files whose indices have been passed
-        as ``ints`` in ``indices``
-
-        Parameters
-        ----------
-        indices : list of ints
-            correspond to the rows of db.regions you want to keep
-
-        Returns
-        -------
-        subset : torch.utils.data.Subset
-            the obtained subset
-
-        Examples
-        --------
-        >>>>db = Database("test.h5")
-        >>>>sub = db.fft.regions([1, 2, 3, 9])
-        """
-        regions = Database(self.h5_file).regions.iloc[indices]
-        return self.subset(regions)
-
-    def __repr__(self):
-        return "<FeatureProxy: '%s/%s'>" % (self.h5_file, self.name)
-
-
-class Database(object):
-    """
-    interface to .h5 databases created by mimikit
-
-    Parameters
-    ----------
-    h5_file : str
-        path to the .h5 file containing the data
-
-    Attributes
-    ----------
-    regions : Regions
-        pandas DataFrame where each row contains information about one stored file.
-        see ``Regions`` for more information
-
-    <feature_proxy> : FeatureProxy
-        each feature created by the extracting function passed to ``make_root_db``
-        is automatically added as attribute. If the extracting function returned a feature
-        by the name ``"fft"``, the attribute ``fft`` of type ``FeatureProxy`` will be automatically
-        added when the file is loaded and you will be able to access it through ``db.fft``.
-    """
-    def __init__(self, h5_file: str):
-        self.h5_file = h5_file
-        self.regions = Regions(self._get_dataframe("/regions"))
-        with h5py.File(h5_file, "r") as f:
-            self.attrs = {k: v for k, v in f.attrs.items()}
-            self.features = f.attrs["features"]
-            # add found features as self.feature_name = FeatureProxy(self.h5_file, feature_name)
-            self._register_features(self.features)
-
-    def visit(self, func=print):
-        """
-        wrapper for ``h5py.File.visititems()``
-
-        Parameters
-        ----------
-        func : function
-            a function to be applied recursively
-
-        Returns
-        -------
-        None
-        """
-        with h5py.File(self.h5_file, "r") as f:
-            f.visititems(func)
-
-    def _get_dataframe(self, key):
-        try:
-            return pd.read_hdf(self.h5_file, key=key)
-        except KeyError:
-            return pd.DataFrame()
-
-    def save_dataframe(self, key, df):
-        """
-        stores a ``pd.DataFrame`` object under ``key``
-
-        Parameters
-        ----------
-        key : str
-            the key under which ``df`` will be stored
-        df : pd.DataFrame
-            the ``DataFrame`` to be stored
-
-        Returns
-        -------
-        df : pd.DataFrame
-            the ``DataFrame`` as it has been stored
-        """
-        with h5py.File(self.h5_file, "r+") as f:
-            if key in f:
-                f.pop(key)
-        df.to_hdf(self.h5_file, key=key, mode="r+")
-        return self._get_dataframe(key)
-
-    def _register_features(self, names):
-        for name in names:
-            setattr(self, name, FeatureProxy(self.h5_file, name))
-        return None
-
-    def _register_dataframes(self, names):
-        for name in names:
-            setattr(self, name, self._get_dataframe(name))
-        return None
-
-    def __repr__(self):
-        return "<Database: '%s'>" % os.path.split(self.h5_file)[-1]
-
-
-def add_feature(h5_file, feature_name, array):
-    # TODO !
-    pass
-
-
-def _add_regions(h5_file, start, stop, duration, filename=None):
-    """
-    adds a row of regions to ``db.regions``
-    """
-    meta = pd.read_hdf(h5_file, "regions")
-    new = Regions.from_start_stop([start], [stop], [duration])
-    filename = datetime.now() if filename is None else filename
-    new["name"] = filename
-    new_meta = pd.concat((meta, new), axis=0, ignore_index=True)
-    new_meta.to_hdf(h5_file, "regions")
-    return new_meta
-
-
-def _add_data(h5_file, ds_name, array, filename):
-    """
-    EXPERIMENTAL!! adds data at the end of a ``Dataset`` in a .h5 file.
-
-    Parameters
-    ----------
-    h5_file : str
-        path to the file
-    ds_name : str
-        name of the ``h5py.Dataset`` to which ``array`` will be added
-    array : np.ndarray
-        the array to add
-    filename : str
-        a name for the array in ``db.regions``
-
-    Returns
-    -------
-    arr : np.ndarray
-        the array as it has been stored
-    """
-    N = array.shape[0]
-    with h5py.File(h5_file, "r+") as f:
-        if f[ds_name].shape[1:] != array.shape[1:]:
-            raise ValueError(
-                ("expected all but the first dimension of `array` to match %s. " % str(f[ds_name].shape[1:])) +
-                ("Got %s" % str(array.shape[1:])))
-        M = f[ds_name].shape[0]
-        f[ds_name].resize(M + N, axis=0)
-        f[ds_name][-N:] = array
-        rv = f[ds_name][-N:]
-    _add_regions(h5_file, M, N + M, N, filename)
-    return rv
\ No newline at end of file
diff --git a/mimikit/data/factory.py b/mimikit/data/create.py
similarity index 53%
rename from mimikit/data/factory.py
rename to mimikit/data/create.py
index 856a167..34e8d63 100644
--- a/mimikit/data/factory.py
+++ b/mimikit/data/create.py
@@ -1,85 +1,17 @@
 import h5py
-import numpy as np
 import pandas as pd
 from multiprocessing import cpu_count, Pool
-from typing import Iterable
 import os
 import warnings
 
-from .api import Database
 from .regions import Regions
-from .transforms import default_extract_func
-
-warnings.filterwarnings("ignore", message="PySoundFile failed.")
 
+__all__ = [
+    'make_root_db',
+    'write_feature'
+]
 
-class AudioFileWalker:
-
-    AUDIO_EXTENSIONS = ("wav", "aif", "aiff", "mp3", "m4a", "mp4")
-
-    def __init__(self, roots=None, files=None):
-        """
-        recursively find audio files from `roots` and/or collect audio files passed in `files`
-
-        Parameters
-        ----------
-        roots : str or list of str
-            a single path (string, os.Path) or an Iterable of paths from which to collect audio files recursively
-        files : str or list of str
-            a single path (string, os.Path) or an Iterable of paths
-
-        Examples
-        --------
-        >>> files = list(AudioFileWalker(roots="./my-directory", files=["sound.mp3"]))
-
-        Notes
-        ------
-        any file whose extension isn't in AudioFileWalker.AUDIO_EXTENSIONS will be ignored,
-        regardless whether it was found recursively or passed through the `files` argument.
-        """
-        generators = []
-
-        if roots is not None and isinstance(roots, Iterable):
-            if isinstance(roots, str):
-                if not os.path.exists(roots):
-                    raise FileNotFoundError("%s does not exist." % roots)
-                generators += [AudioFileWalker.walk_root(roots)]
-            else:
-                for r in roots:
-                    if not os.path.exists(r):
-                        raise FileNotFoundError("%s does not exist." % r)
-                generators += [AudioFileWalker.walk_root(root) for root in roots]
-
-        if files is not None and isinstance(files, Iterable):
-            if isinstance(files, str):
-                if not os.path.exists(files):
-                    raise FileNotFoundError("%s does not exist." % files)
-                generators += [(f for f in [files] if AudioFileWalker.is_audio_file(files))]
-            else:
-                for f in files:
-                    if not os.path.exists(f):
-                        raise FileNotFoundError("%s does not exist." % f)
-                generators += [(f for f in files if AudioFileWalker.is_audio_file(f))]
-
-        self.generators = generators
-
-    def __iter__(self):
-        for generator in self.generators:
-            for file in generator:
-                yield file
-
-    @staticmethod
-    def walk_root(root):
-        for directory, _, files in os.walk(root):
-            for audio_file in filter(AudioFileWalker.is_audio_file, files):
-                yield os.path.join(directory, audio_file)
-
-    @staticmethod
-    def is_audio_file(filename):
-        # filter out hidden files
-        if os.path.split(filename.strip("/"))[-1].startswith("."):
-            return False
-        return os.path.splitext(filename)[-1].strip(".") in AudioFileWalker.AUDIO_EXTENSIONS
+warnings.filterwarnings("ignore", message="PySoundFile failed.")
 
 
 def _sizeof_fmt(num, suffix='b'):
@@ -92,7 +24,21 @@ def _sizeof_fmt(num, suffix='b'):
 
 # Core function
 
-def file_to_db(abs_path, extract_func=default_extract_func, output_path=None, mode="w"):
+def write_feature(h5_file, feature_name, attrs, data, regions=None, files=None):
+    with h5py.File(h5_file, "r+") as f:
+        if feature_name in f:
+            f.pop(feature_name)
+        ds = f.create_dataset(name=feature_name, shape=data.shape, data=data)
+        ds.attrs.update(attrs)
+        f.attrs["features"] = [*f.attrs["features"], feature_name]
+    if regions is not None:
+        pd.DataFrame(regions).to_hdf(h5_file, feature_name + "_regions", "r+")
+    if files is not None:
+        pd.DataFrame(files).to_hdf(h5_file, feature_name + "_files", "r+")
+    return None
+
+
+def file_to_h5(abs_path, extract_func=None, output_path=None, mode="w"):
     """
     apply ``extract_func`` to ``abs_path`` and write the result in a .h5 file.
 
@@ -101,7 +47,7 @@ def file_to_db(abs_path, extract_func=default_extract_func, output_path=None, mo
     abs_path : str
         path to the file to be extracted
     extract_func : function
-        the function to use for the extraction - should take exactly one argument
+        the function to use for the extraction - should take exactly one argument: the path to the file to be extracted.
     output_path : str or None
         name of the created .h5 file if not None, else the name of the created .h5 file
         will be of the form : "<abs_path_without_extension>.h5"
@@ -120,19 +66,25 @@ def file_to_db(abs_path, extract_func=default_extract_func, output_path=None, mo
         if output_path[-3:] != ".h5":
             output_path += ".h5"
     rv = extract_func(abs_path)
-    if "regions" not in rv:
-        raise ValueError("Expected `extract_func` to return a ('regions', Regions) item. Found none")
     info = {}
-    f = h5py.File(output_path, mode)
-    for name, (attrs, data) in rv.items():
-        if issubclass(type(data), np.ndarray):
-            ds = f.create_dataset(name=name, shape=data.shape, data=data)
-            ds.attrs.update(attrs)
-            info[name] = {"dtype": ds.dtype, "shape": ds.shape}
-        elif issubclass(type(data), pd.DataFrame):
+    if os.path.exists(output_path) and mode == 'w':
+        os.remove(output_path)
+    if not os.path.exists(os.path.dirname(output_path)):
+        os.makedirs(os.path.dirname(output_path), exist_ok=True)
+    try:
+        f = h5py.File(output_path, mode)
+    except OSError as e:
+        print(output_path)
+        raise e
+    for name, (attrs, data, regions) in rv.items():
+        ds = f.create_dataset(name=name, shape=data.shape, data=data)
+        ds.attrs.update(attrs)
+        info[name] = {"dtype": ds.dtype, "shape": ds.shape}
+        if regions is not None:
             f.close()
-            pd.DataFrame(data).to_hdf(output_path, name, "r+")
+            pd.DataFrame(regions).to_hdf(output_path, name + "_regions", "r+")
             f = h5py.File(output_path, "r+")
+    f.attrs["features"] = list(rv.keys())
     f.flush()
     f.close()
     return output_path, info
@@ -141,7 +93,8 @@ def file_to_db(abs_path, extract_func=default_extract_func, output_path=None, mo
 # Multiprocessing routine
 
 def _make_db_for_each_file(file_walker,
-                           extract_func=default_extract_func,
+                           extract_func=None,
+                           destination="",
                            n_cores=cpu_count()):
     """
     apply ``extract_func`` to the files found by ``file_walker``
@@ -163,13 +116,14 @@ def _make_db_for_each_file(file_walker,
         each tuple in the list is of the form ``("<created_file>.h5", dict(feature_name=dict(dtype=..., shape=...), ...))``
     """
     # add ".tmp_" prefix to the output_paths
-    args = [(file, extract_func, os.path.join(os.path.split(file)[0], ".tmp_" + os.path.split(file)[1]))
-            for file in file_walker]
-    if len(args) > n_cores:
-        with Pool(n_cores) as p:
-            tmp_dbs_infos = p.starmap(file_to_db, args)
+    # args = [(file, extract_func, )
+    args = [(file, extract_func, os.path.join(destination, file.strip('.').strip('/')))
+            for n, file in enumerate(file_walker)]
+    if len(args) > 1:
+        with Pool(min(n_cores, len(args))) as p:
+            tmp_dbs_infos = p.starmap(file_to_h5, args)
     else:
-        tmp_dbs_infos = [file_to_db(*arg) for arg in args]
+        tmp_dbs_infos = [file_to_h5(*arg) for arg in args]
     return tmp_dbs_infos
 
 
@@ -197,13 +151,14 @@ def _aggregate_db_infos(infos):
         dtype = dtype.pop()
         shapes = [db[1][f]["shape"] for db in infos]
         dims = shapes[0][1:]
-        assert all(shp[1:] == dims for shp in
-                   shapes[1:]), "all features should have the same dimensions but for the first axis"
-        # collect the regions for the files
+        if not all(shp[1:] == dims for shp in shapes[1:]):
+            raise ValueError("all features should have the same dimensions but for the first axis." + \
+                             f" feature {str(f)} returned different shapes")
+        # collect the regions for the files (this is different from the possible segmentation regions!)
         regions = Regions.from_duration([s[0] for s in shapes])
         ds_shape = (regions.last_stop, *dims)
         regions.index = paths
-        ds_definitions[f] = {"shape": ds_shape, "dtype": dtype, "meta_regions": regions}
+        ds_definitions[f] = {"shape": ds_shape, "dtype": dtype, "files_regions": regions}
     return ds_definitions
 
 
@@ -237,24 +192,37 @@ def _aggregate_dbs(target, tmp_dbs_infos, mode="w"):
     for feature, info in features_infos.items():
         # collect the arguments for copying each file into its regions in the master datasets/features
         args += [(source, feature, indices) for source, indices in
-                 zip(info["meta_regions"].index, info["meta_regions"].slices(time_axis=0))]
+                 zip(info["files_regions"].index, info["files_regions"].slices(time_axis=0))]
+        # "<feature>_files" will store the slices corresponding to the files for this feature
+        files_regions = info["files_regions"]
+        files_regions.index = files_regions.index.str.replace(".tmp_", "").str.replace(".h5", "")
+        files_regions = files_regions.reset_index(drop=False)
+        files_regions = files_regions.rename(columns={"index": "name"})
+        pd.DataFrame(files_regions).to_hdf(target, feature + "_files", mode="r+")
 
     # copy the data
-    intra_regions = None
+    intra_regions = {}
     for source, key, indices in args:
         with h5py.File(source, "r") as src:
             data = src[key][()]
             attrs = {k: v for k, v in src[key].attrs.items()}
-        # concat the regions :
-        regions = pd.read_hdf(source, key="regions", mode="r")
-        regions.loc[:, ("start", "stop")] += indices[0].start
-        # remove ".tmp_" from the source's name
-        regions.loc[:, "name"] = "".join(source.split(".tmp_"))
-        intra_regions = pd.concat(([intra_regions] if intra_regions is not None else []) + [regions])
+            if key + "_regions" in src:
+                if key not in intra_regions:
+                    intra_regions[key] = []
+                # concat the regions :
+                regions = pd.read_hdf(source, key=key + "_regions", mode="r")
+                regions.loc[:, ("start", "stop")] += indices[0].start
+                regions = regions.reset_index(drop=False)
+                # remove ".tmp_" from the source's name
+                regions.loc[:, "name"] = "".join(source.split(".tmp_"))
+                intra_regions[key] += [regions]
         with h5py.File(target, "r+") as trgt:
             trgt[key][indices] = data
             trgt[key].attrs.update(attrs)
-    pd.DataFrame(intra_regions).to_hdf(target, "regions", mode="r+")
+    # "<feature>_regions" will store the segmenting slices for this feature
+    for key, regions in intra_regions.items():
+        if regions is not None:
+            pd.concat(regions).reset_index(drop=True).to_hdf(target, key + "_regions", mode="r+")
 
     # remove the temp_dbs
     for src, _ in tmp_dbs_infos:
@@ -262,8 +230,12 @@ def _aggregate_dbs(target, tmp_dbs_infos, mode="w"):
             os.remove(src)
 
 
-def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_func,
-                 n_cores=cpu_count()):
+def make_root_db(db_name,
+                 files_iterable=tuple(),
+                 extract_func=None,
+                 tmp_destination="/tmp/",
+                 n_cores=cpu_count()
+                 ):
     """
     extract and aggregate several files into a single .h5 Database
 
@@ -271,15 +243,12 @@ def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_f
     ----------
     db_name : str
         the name of the db to be created
-    roots : str or list of str, optional
-        directories from which to search recursively for audio files.
-        default is "./"
-    files : str or list of str, optional
-        single file(s) to include in the db.
-        default is `None`
+    files_iterable : iterable of str, optional
+        the files to be extracted
     extract_func : function, optional
         the function to use for the extraction - should take exactly one argument.
         the default transforms the file to the stft with n_fft=2048 and hop_length=512.
+    tmp_destination : str, optional
     n_cores : int, optional
         the number of cores to use to parallelize the extraction process.
         default is the number of available cores on the system.
@@ -289,7 +258,5 @@ def make_root_db(db_name, roots='./', files=None, extract_func=default_extract_f
     db : Database
         the created db
     """
-    walker = AudioFileWalker(roots, files)
-    tmp_dbs_infos = _make_db_for_each_file(walker, extract_func, n_cores)
+    tmp_dbs_infos = _make_db_for_each_file(files_iterable, extract_func, tmp_destination, n_cores)
     _aggregate_dbs(db_name, tmp_dbs_infos, "w")
-    return Database(db_name)
diff --git a/mimikit/data/data_object.py b/mimikit/data/data_object.py
deleted file mode 100644
index ae94dad..0000000
--- a/mimikit/data/data_object.py
+++ /dev/null
@@ -1,311 +0,0 @@
-import torch
-import numpy as np
-from torch.utils.data import Dataset as TorchDataset, Subset, random_split, DataLoader
-from torch.utils.data._utils.collate import default_convert
-from functools import update_wrapper
-
-
-def map_if_multi(attr):
-    """
-    decorator for applying a method recursively to ``getattr(self, attr)`` if ``getattr(self, attr)`` is a tuple.
-
-    Parameters
-    ----------
-    attr : str
-        name of the ``attribute`` to which the method will apply and that can possibly be a ``tuple``.
-
-    Returns
-    -------
-    decorator : function
-        the wrapping function for decorating the method
-    """
-    def decorator(fn):
-        def wrapper(self, *arg, **kwargs):
-            if isinstance(getattr(self, attr), tuple):
-                return tuple(getattr(x, fn.__name__)(*arg, **kwargs)[0] for x in getattr(self, attr))
-            else:
-                # for consistency, we always return a tuple
-                # noinspection PyRedundantParentheses
-                return (fn(self, *arg, **kwargs), )
-
-        return update_wrapper(wrapper, fn,
-                              # for some reason, __mod__ causes errors when using functools.wraps
-                              ("__name__", '__qualname__', '__doc__', '__annotations__'),
-                              ("__dict__",))
-    return decorator
-
-
-class DataObject(TorchDataset):
-    """
-    the goal of the ``DataObject`` class is to provide a single universal object for making smart ``Datasets``.
-
-    In its constructor, ``DataObject`` accepts ``numpy`` arrays, ``python iterables``, ``torch.Tensors`` and ``Datasets``,
-    ``FeatureProxies`` from ``mimikit``'s ``Databases`` and generally, any object that implements ``__len__`` and ``__getitem__``
-    or ``__iter__`` - i.e. objects that behaves like 'map-style' and 'iter-style' ``Datasets``.
-
-    Furthermore, ``DataObject`` handles nested ``tuples`` of such "data objects" as if they were single containers with
-    multiple fields, thus allowing you to easily define the shape and content of your batches at the very top-level.
-    For example, if you have raw images and labels, you just need
-    ```
-    dataloader = Dataloader(DataObject((images, labels)), ....)
-    ```
-    in order to later do
-    ```
-    images, labels = batch
-    ```
-    in your ``training_step()``.
-
-    ``DataObject`` instances also implement some of the most common ``numpy`` and ``torch`` properties like ``shape``
-    or ``dtype`` and some handy methods like ``to_tensor()``, ``to(device)`` or ``split([...])`` thus removing a lot of
-    boilerplate imports and code-lines that deal with low-level data management while maintaining access to the parts
-    of the low-level interfaces that are always needed when defining models and trainings.
-    """
-
-    @property
-    def data(self):
-        """
-        give access to the underlying data to wrappers and co
-        """
-        return self._object
-
-    @property
-    def n_features(self):
-        return self._n_features
-
-    @property
-    def dims(self):
-        return self._shape
-
-    @property
-    def shape(self):
-        return self._shape
-
-    @property
-    def dtype(self):
-        return self._dtype
-
-    @property
-    def device(self):
-        return self._device
-
-    @property
-    def style(self):
-        return self._style
-
-    @property
-    def attrs(self):
-        """
-        @return: `attrs` : dict of user-specified infos he wants to keep close (e.g. {"sr": 22050, "time-axis": 0})
-        """
-        return self._attrs
-
-    def __init__(self,
-                 data_object=None,
-                 wrappers=None,
-                 augmentations=None,
-                 attrs=None):
-        """
-        Here we initialize the instance.
-        """
-        if data_object is None:
-            raise ValueError("`data_object` can not be None.")
-        self._object = data_object
-        self._wrapper = wrappers
-        self._augmentation = augmentations
-        self._attrs = attrs
-        self._n_features = 1 if not isinstance(data_object, tuple) else len(data_object)
-        self._is_multi = self._n_features > 1
-
-        if self._is_multi:
-            # we cast to Dataset recursively
-            self._object = tuple(x if isinstance(x, DataObject) else DataObject(x) for x in self._object)
-
-        self._style = self.get_style()
-        if any(style is None for style in self._style):
-            raise ValueError("Expected data_object to either implement __getitem__ and __len__, or __iter__."
-                             " object of class %s implements none of those." % str(type(data_object)))
-        if self._is_multi:
-            if not all(self._style[0] == style for style in self._style[1:]):
-                raise TypeError("Expected all data_objects to be of the same style. Got %s" % str(self._style))
-            lengths = tuple(len(obj) for obj in self._object)
-            if self._style[0] is "map" and not all(lengths[0] == n for n in lengths[1:]):
-                raise ValueError("Expected all 'map-style' data_objects to be of same lengths. Got %s" % \
-                                 str(lengths))
-        self._dtype = self.get_dtype()
-        if any(dtype is None for dtype in self._dtype):
-            raise TypeError("Expected elements of `data_object` to be of torch, numpy or built-in numerical dtype."
-                            " Got following dtypes %s" % str(self._dtype))
-        self._shape = self.get_shape()
-        self._device = self.get_device()
-
-    def __call__(self, *args, **kwargs):
-        """
-        Method used to wrap around an other Dataset. Only used in the wrapper classes.
-        """
-        pass
-
-    def __len__(self):
-        return len(self._object[0]) if self._is_multi else len(self._object)
-
-    @map_if_multi("_object")
-    def __getitem__(self, item):
-        if isinstance(item, tuple):
-            return tuple(self._object[i] for i in item)
-        return self._object[item]
-
-    def __iter__(self):
-        if isinstance(self._object, tuple):
-            return zip(*[iter(obj) for obj in self._object])
-        return iter(self._object)
-
-    @staticmethod
-    def has_len(data_object):
-        return getattr(data_object, "__len__", False)
-
-    @staticmethod
-    def _get_style(data_object):
-        if getattr(data_object, "__getitem__", False) and getattr(data_object, "__len__", False):
-            return "map"
-        elif getattr(data_object, "__iter__", False):
-            return "iter"
-        return None
-
-    @map_if_multi("_object")
-    def get_style(self):
-        return DataObject._get_style(self._object)
-
-    @staticmethod
-    def _get_elem(data_object):
-        if DataObject._get_style(data_object) == "map":
-            return data_object[0]
-        elif DataObject._get_style(data_object) == "iter":
-            return next(iter(data_object))
-        return None
-
-    @map_if_multi("_object")
-    def get_elem(self):
-        return DataObject._get_elem(self._object)
-
-    @staticmethod
-    def _get_shape(data_object):
-        shape = getattr(data_object, "shape", None)
-        if shape is None:
-            # we try to fet lengths recursively and add `None` when we get elements with no length
-            # thus, in the case of a hierarchy of generators, returning at least the number of 'axes'
-            N = len(data_object) if DataObject.has_len(data_object) else None
-            shape = (N, )
-            elem = DataObject._get_elem(data_object)
-            while elem is not None and DataObject._get_elem(elem) != elem:
-                shape = (*shape, len(elem) if DataObject.has_len(elem) else None)
-                elem = DataObject._get_elem(elem)
-            # last None in the shape is bottom-level element
-            shape = shape[:-1]
-        if isinstance(shape, torch.Size):
-            shape = tuple(shape)
-        return shape
-
-    @map_if_multi("_object")
-    def get_shape(self):
-        return DataObject._get_shape(self._object)
-
-    @staticmethod
-    def _get_dtype(data_object):
-        elem = DataObject._get_elem(data_object)
-        dtype = getattr(data_object, "dtype", None)
-        if dtype is None:
-            dtype = getattr(elem, "dtype", None) or (type(elem) if type(elem) in (int, float, complex) else None)
-            while dtype is None and elem is not None and DataObject._get_elem(elem) != elem:
-                elem = DataObject._get_elem(elem)
-                dtype = getattr(elem, "dtype", None) or (type(elem) if type(elem) in (int, float, complex) else None)
-        if isinstance(dtype, torch.dtype):
-            pass
-        elif dtype in (float, int, complex):
-            dtype = getattr(torch, np.dtype(dtype).name, None)
-        elif getattr(dtype, "name", False):  # it's a numpy array!
-            dtype = getattr(torch, dtype.name, dtype)
-        return dtype
-
-    @map_if_multi("_object")
-    def get_dtype(self):
-        return DataObject._get_dtype(self._object)
-
-    @staticmethod
-    def _get_device(data_object):
-        elem = DataObject._get_elem(data_object)
-        device = getattr(data_object, "device", None)
-        if device is None:
-            device = getattr(elem, "device", "cpu")
-        if isinstance(device, torch.device):
-            device = device.type
-        return device
-
-    @map_if_multi("_object")
-    def get_device(self):
-        return DataObject._get_device(self._object)
-
-    @map_if_multi("_object")
-    def to_tensor(self):
-        if issubclass(type(self._object), torch.Tensor):
-            return
-        elif "iter" in self.style:
-            raise TypeError("Cannot convert 'iter' style objects to tensor.")
-        maybe_tensor = default_convert(self._object[:])
-        if issubclass(type(maybe_tensor), torch.Tensor):
-            self._object = maybe_tensor
-            return
-        try:
-            self._object = torch.tensor(self._object)
-        except Exception as e:
-            raise e
-        return
-
-    @map_if_multi("_object")
-    def select(self, indices, inplace=False):
-        if not self.has_len(self._object):
-            raise TypeError("Cannot select indices of objects that don't implement `__len__`")
-        if inplace:
-            self._object = self._object[np.sort(indices)]
-        else:
-            self._object = Subset(self._object, indices)
-
-    @map_if_multi("_object")
-    def to(self, device):
-        """move any underlying tensor to some device"""
-        if isinstance(self._object, torch.Tensor):
-            self._object = self._object.to(device)
-
-    def split(self, splits):
-        """
-        @param splits: Sequence of floats or ints possibly containing None. The sequence of elements
-        corresponds to the proportion (floats), the number of examples (ints) or the absence of
-        train-set, validation-set, test-set, other sets... in that order.
-        @return: the *-sets
-        """
-        nones = []
-        if any(x is None for x in splits):
-            if splits[0] is None:
-                raise ValueError("the train-set's split cannot be None")
-            nones = [i for i, x in zip(range(len(splits)), splits) if x is None]
-            splits = [x for x in splits if x is not None]
-        if all(type(x) is float for x in splits):
-            splits = [x / sum(splits) for x in splits]
-            N = len(self)
-            # leave the last one out for now because of rounding
-            as_ints = [int(N * x) for x in splits[:-1]]
-            # check that the last is not zero
-            if N - sum(as_ints) == 0:
-                raise ValueError("the last split rounded to zero element. Please provide a greater float or consider "
-                                 "passing ints.")
-            as_ints += [N - sum(as_ints)]
-            splits = as_ints
-        sets = list(random_split(self, splits))
-        if any(nones):
-            sets = [None if i in nones else sets.pop(0) for i in range(len(sets + nones))]
-        return tuple(sets)
-
-    def load(self, **kwargs):
-        """pack self into a batch producer (Dataloader)"""
-        return DataLoader(self, **kwargs)
-
-    def __repr__(self):
-        return "<DataObject shape:%s, dtype:%s>" % (self.shape, self.dtype)
\ No newline at end of file
diff --git a/mimikit/data/database.py b/mimikit/data/database.py
new file mode 100644
index 0000000..d263d05
--- /dev/null
+++ b/mimikit/data/database.py
@@ -0,0 +1,414 @@
+from functools import partial
+import h5py
+import torch
+from torch.utils.data._utils.collate import default_convert
+from torch.utils.data import random_split, Dataset
+import numpy as np
+import pandas as pd
+import os
+from argparse import Namespace
+import pytorch_lightning as pl
+
+from .create import make_root_db, write_feature
+from .feature import Feature
+from .regions import Regions
+from ..file_walker import FileWalker, EXTENSIONS
+
+__all__ = [
+    "Database"
+]
+
+
+class FeatureProxy(object):
+    """
+    Interface to the numerical data (array) stored in a .h5 file created with ``mimikit.data.make_root_db``
+
+    ``FeatureProxy`` objects are like ``numpy`` arrays (on disk) and can be used as map-style ``Dataset``
+    because they implement ``__len__`` and ``__getitem__``.
+
+    Parameters
+    ----------
+    h5_file : str
+        name of the file
+    ds_name : str
+        name of the ``h5py.Dataset`` containing the data to be served
+
+    Attributes
+    ----------
+    N : int
+        the length of the first dimension of the underlying array
+    shape : tuple of int
+        the shape of the underlying array
+    attrs : dict
+        dictionary of additional information about the data as returned
+        by the ``extract_func`` passed to ``make_root_db``.
+        Typically, this is where you want to store the parameters of your extracting function, e.g.
+        sample rate, hop length etc.
+    files : Regions
+        pandas DataFrame where each row contains name, start, stop & duration of one stored file.
+        see ``Regions`` for more information
+    regions : Regions or None
+        pandas DataFrame containing index, start, stop & duration for segments ('sub-regions' of the files)
+        This attribute isn't None only if the ``extract()`` method of the parent ``Database`` class returned such
+        an object for this feature.
+    """
+
+    def __init__(self, h5_file: str, ds_name: str, keep_open=False):
+
+        self.h5_file = h5_file
+        self.name = ds_name
+        with h5py.File(h5_file, "r") as f:
+            ds = f[self.name]
+            self.N = ds.shape[0]
+            self.shape = ds.shape
+            self.dtype = ds.dtype
+            self.attrs = {k: v for k, v in ds.attrs.items()}
+            has_files = self.name + "_files" in f
+            has_regions = self.name + "_regions" in f
+        self.files = Regions(pd.read_hdf(h5_file, self.name + "_files", mode="r")) if has_files else None
+        self.regions = Regions(pd.read_hdf(h5_file, self.name + "_regions", mode="r")) if has_regions else None
+        self._f = h5py.File(h5_file, "r+") if keep_open else None
+
+    def __len__(self):
+        return self.N
+
+    def __getitem__(self, item):
+        if self._f is not None:
+            return self._f[self.name][item]
+        with h5py.File(self.h5_file, "r") as f:
+            rv = f[self.name][item]
+        return rv
+
+    def __setitem__(self, item, value):
+        if self._f is not None:
+            self._f[self.name][item] = value
+        with h5py.File(self.h5_file, "r+") as f:
+            f[self.name][item] = value
+        return
+
+    def get_regions(self, regions):
+        """
+        get the data (numpy array) corresponding to the rows of `regions`
+
+        Parameters
+        ----------
+        regions : Regions
+            the files you want to get as array
+
+        Returns
+        -------
+        data : np.ndarray
+            the concatenated data for all the requested files
+
+        Examples
+        --------
+
+            .. code-block::
+
+                from mimikit import Database
+                X = Database("my-db.h5").fft
+                # only get files 0, 3 & 7
+                X.get_regions(X.regions.iloc[[0, 3, 7]])
+        """
+        slices = regions.slices(0)
+        return np.concatenate(tuple(self[slice_i] for slice_i in slices), axis=0)
+
+    def close(self):
+        if self._f is not None:
+            self._f.close()
+        return self
+
+    def __repr__(self):
+        return "<FeatureProxy: '%s/%s'>" % (self.h5_file, self.name)
+
+
+class Database(Dataset):
+    """
+    interface to .h5 databases created by mimikit
+
+    Parameters
+    ----------
+    h5_file : str
+        path to the .h5 file containing the data
+    keep_open : bool, optional
+        whether to keep the h5 file open or close it after each query.
+        Default is ``False``.
+
+    Attributes
+    ----------
+    <feature_proxy> : FeatureProxy
+        each feature created by the extracting function passed to ``make_root_db``
+        is automatically added as attribute. If the extracting function returned a feature
+        by the name ``"fft"``, the attribute ``fft`` of type ``FeatureProxy`` will be automatically
+        added when the file is loaded and you will be able to access it through ``db.fft``.
+    """
+
+    schema = {}
+
+    def __init__(self, h5_file: str, keep_open=False):
+        self.h5_file = h5_file
+        self.keep_open = keep_open
+        with h5py.File(h5_file, "r") as f:
+            self.attrs = {k: v for k, v in f.attrs.items()}
+            self.features = f.attrs.get("features", ["fft"])
+        # add found features as self.feature_name = FeatureProxy(self.h5_file, feature_name)
+        self.params = Namespace()
+        self._register_features(self.features, keep_open)
+
+    def copy(self):
+        return type(self)(self.h5_file, self.keep_open)
+
+    @staticmethod
+    def extract(path, **kwargs):
+        raise NotImplementedError
+
+    @classmethod
+    def make(cls, db_name, files_ext='audio', items=tuple(), **kwargs):
+        walker = FileWalker(files_ext, items)
+        make_root_db(db_name, walker, partial(cls.extract, **kwargs))
+        return cls(db_name)
+
+    @classmethod
+    def _load(cls, path, schema={}):
+        """
+        default extract_func for Database.build. Roughly equivalent to :
+            ``{feat_name: feat.load(path) for feat_name, feat in features_dict.items()}``
+
+        """
+        out = {}
+        for f_name, f in schema.items():
+            # check that f has `load` and that `path` is of the right type
+            if getattr(type(f), 'load', Feature.load) != Feature.load and \
+                    os.path.splitext(path)[-1].strip('.') in EXTENSIONS[f.__ext__]:
+                obj = f.load(path)
+            # feature is derived from the output of an other
+            elif hasattr(f, 'input_key') and f.input_key in out:
+                # we us __call__ instead of load, the former being eq. to a transform
+                obj = f(out[f.input_key][1])
+            else:
+                obj = None
+            # obj is (data, regions)
+            if isinstance(obj, tuple) and tuple(map(type, obj)) == (np.ndarray, Regions):
+                out[f_name] = getattr(f, 'params', {}), *obj
+            elif isinstance(obj, Regions):
+                out[f_name] = getattr(f, 'params', {}), np.array([]), obj
+            elif isinstance(obj, np.ndarray):
+                out[f_name] = getattr(f, 'params', {}), obj, None
+        return out
+
+    @classmethod
+    def create(cls, db_name, sources=tuple(), schema={}):
+        """
+        creates a db from the schema provided in `features_dict` and the files or root directories found in `items`
+
+        Parameters
+        ----------
+        db_name : str
+            the name of the file to be created
+        sources : str or iterable of str
+            the sources to be passed to `FileWalker`
+        schema : dict
+            keys (`str`) are the names of the features, values are `Feature` objects
+
+        Returns
+        -------
+        db : Database
+            an instance of the created db
+
+        """
+        # get the set of file extensions from the features and instantiate a walker
+        exts = {f.__ext__ for f in schema.values() if getattr(f, '__ext__', False)}
+        walker = FileWalker(exts, sources)
+        # run the extraction job
+        make_root_db(db_name, walker, partial(cls._load, schema=schema))
+        # add post-build features
+        db = cls(db_name)
+        for f_name, f in schema.items():
+            # let features the chance to update them selves confronted to their whole set
+            if getattr(type(f), "post_create", Feature.post_create) != Feature.post_create:
+                rv = f.post_create(db, f_name)
+                if isinstance(rv, np.ndarray):
+                    rv = (rv, getattr(db, f_name).regions)
+                elif isinstance(rv, Regions):
+                    rv = (getattr(db, f_name)[:], rv)
+                elif tuple(*map(type, rv)) == (np.ndarray, Regions):
+                    pass
+                write_feature(db_name, f_name, getattr(f, 'params', {}),
+                              *rv, files=getattr(db, f_name).files)
+        db = cls(db_name)
+        db.schema = schema
+        return db
+
+    @classmethod
+    def make_temp(cls, roots=None, files=None, **kwargs):
+        return cls.make("/tmp/%s-tmp.h5" % cls.__name__, roots, files, **kwargs)
+
+    def _visit(self, func=print):
+        with h5py.File(self.h5_file, "r") as f:
+            f.visititems(func)
+
+    def _register_features(self, names, keep_open):
+        for name in names:
+            setattr(self, name, FeatureProxy(self.h5_file, name, keep_open))
+            setattr(self.params, name, getattr(self, name).attrs)
+        return self
+
+    def _register_dataframes(self, names):
+        for name in names:
+            try:
+                df = pd.read_hdf(self.h5_file, key=name)
+            except KeyError:
+                df = pd.DataFrame()
+            setattr(self, name, df)
+        return self
+
+    def __repr__(self):
+        return "<Database: '%s'>" % os.path.split(self.h5_file)[-1]
+
+    def bind(self, dataset_cls):
+        """
+        extends and if necessary overwrites the class of this instance with `dataset_cls`
+
+        if `dataset_cls` contains a `__init__` method, it must take 2 arguments (self, model=None)
+        and will be renamed to `prepare_dataset`.
+
+        for example usage of this method, see `mimikit.data.DataModule.prepare_data()`
+
+        Parameters
+        ----------
+        dataset_cls: type
+            the interface to be added to the class of self
+
+        Returns
+        -------
+        self
+            the extended instance
+        """
+        def _bind(func, as_name=None):
+            """
+            Bind the function *func* to the class of *self*
+            """
+            if as_name is None:
+                as_name = func.__name__
+            setattr(type(self), as_name, func)
+
+        for k, v in dataset_cls.__dict__.items():
+            if k in ('__init__', 'prepare_dataset'):
+                assert v.__code__.co_argcount == 2, \
+                    "__init__ method of dataset_cls should take 2 arguments: self, model=None"
+                # rename init
+                _bind(v, 'prepare_dataset')
+            elif k in ('__len__', '__getitem__') or k not in type(self).__dict__:
+                if '__call__' in dir(v):
+                    _bind(v, k)
+                else:
+                    # class attributes are likewise attached to type(self)
+                    setattr(type(self), k, v)
+
+        return self
+
+    def prepare_dataset(self, model: pl.LightningModule):
+        """
+        placeholder for implementing what need to be done before serving data
+
+        Parameters
+        ----------
+        model : pl.LightningModule
+            The model that will consume this dataset.
+
+        Returns
+        -------
+        None
+        """
+        raise NotImplementedError
+
+    def __len__(self):
+        raise NotImplementedError
+
+    def __getitem__(self, item):
+        raise NotImplementedError
+
+    def split(self, splits):
+        """
+        performs random splits on self
+
+        Parameters
+        ----------
+        splits: Sequence of floats or ints possibly containing None.
+            The sequence of elements corresponds to the proportion (floats), the number of examples (ints) or the absence of
+            train-set, validation-set, test-set, other sets... in that order.
+
+        Returns
+        -------
+        splits: tuple
+            the *-sets
+        """
+        nones = []
+        if any(x is None for x in splits):
+            if splits[0] is None:
+                raise ValueError("the train-set's split cannot be None")
+            nones = [i for i, x in zip(range(len(splits)), splits) if x is None]
+            splits = [x for x in splits if x is not None]
+        if all(type(x) is float for x in splits):
+            splits = [x / sum(splits) for x in splits]
+            N = len(self)
+            # leave the last one out for now because of rounding
+            as_ints = [int(N * x) for x in splits[:-1]]
+            # check that the last is not zero
+            if N - sum(as_ints) == 0:
+                raise ValueError("the last split rounded to zero element. Please provide a greater float or consider "
+                                 "passing ints.")
+            as_ints += [N - sum(as_ints)]
+            splits = as_ints
+        sets = list(random_split(self, splits))
+        if any(nones):
+            sets = [None if i in nones else sets.pop(0) for i in range(len(sets + nones))]
+        return tuple(sets)
+
+    def to_tensor(self):
+        for feat in self.features:
+            as_tensor = self._to_tensor(getattr(self, feat))
+            setattr(self, feat, as_tensor)
+        return self
+
+    def to(self, device):
+        for feat in self.features:
+            self._to(getattr(self, feat), device)
+        return self
+
+    @staticmethod
+    def _to_tensor(obj):
+        if type(obj) is torch.Tensor:
+            return obj
+        # converting obj[:] makes sure we get the data out of any db.feature object
+        maybe_tensor = default_convert(obj[:])
+        if type(maybe_tensor) is torch.Tensor:
+            return maybe_tensor
+        try:
+            obj = torch.tensor(obj)
+        except Exception as e:
+            raise e
+        return obj
+
+    @staticmethod
+    def _to(obj, device):
+        """move any underlying tensor to some device"""
+        if getattr(obj, "to", False):
+            return obj.to(device)
+        raise TypeError("object %s has no `to()` attribute" % str(obj))
+
+
+def join(target_name, feature_names, databases):
+
+    with h5py.File(target_name, "w") as f:
+        for feat_name in feature_names:
+            lengths = [getattr(db, feat_name).shape[0] for db in databases]
+            dim = set([getattr(db, feat_name).shape[1] for db in databases]).pop()
+            layout = h5py.VirtualLayout(shape=(sum(lengths), dim))
+            offset = 0
+            for i, n in enumerate(lengths):
+                vsource = h5py.VirtualSource(databases[i].h5_file, feat_name, shape=(n, dim))
+                layout[offset:offset + n] = vsource
+                offset += n
+            ds = f.create_virtual_dataset(feat_name, layout)
+            ds.attrs = getattr(databases[0], feat_name).attrs
diff --git a/mimikit/data/datamodule.py b/mimikit/data/datamodule.py
new file mode 100644
index 0000000..43525eb
--- /dev/null
+++ b/mimikit/data/datamodule.py
@@ -0,0 +1,340 @@
+import dataclasses as dtc
+import os
+import pytorch_lightning as pl
+import torch
+import numpy as np
+from numpy.lib.stride_tricks import as_strided as np_as_strided
+from torch.utils.data import Dataset, DataLoader
+from typing import Iterable, Optional, Callable
+import re
+from random import randint
+from torch._six import container_abcs, string_classes
+
+from . import Database
+
+__all__ = [
+    'AsSlice',
+    'AsFramedSlice',
+    'Input',
+    'Target',
+    'process_batch',
+    'DefaultDataset',
+    'DataModule'
+]
+
+
+@dtc.dataclass
+class Getter:
+    """
+    base class for implementing data getter
+
+    Parameters
+    ----------
+
+    Attributes
+    ----------
+    n : int or None
+        the length of the underlying data
+    """
+    n: Optional[int] = dtc.field(default=None, init=False)
+
+    def __call__(self, feat_data, item):
+        """
+        apply this instance's logic to get data from ``feat_data`` for a given ``item``
+
+        Parameters
+        ----------
+        feat_data: [np.ndarray, torch.Tensor, mimikit.FeatureProxy]
+
+        item: int
+            the index emitted from a sampler
+
+        Returns
+        -------
+        data: Any
+            the examples corresponding to this item
+        """
+        return feat_data[item]
+
+    def __len__(self):
+        return self.n
+
+
+@dtc.dataclass()
+class AsSlice(Getter):
+    """
+    maps an ``item`` to a slice of data
+
+    Parameters
+    ----------
+    shift : int
+        the slice will start at the index `item + shift`
+    length : int
+        the length of the slice
+    stride : int
+        sub-sampling factor. Every `stride` datapoints `item` increases of `1`
+
+    Examples
+    --------
+
+    .. testcode::
+
+       import mimikit as mmk
+
+       slicer = mmk.AsSlice(shift=2, length=3)
+       data, item = list(range(10)), 2
+
+       # now use it like a function :
+       sliced = slicer(data, item)
+
+       print(sliced)
+
+    will output:
+
+    .. testoutput::
+
+       [4, 5, 6]
+    """
+    shift: int = 0
+    length: int = 1
+    stride: int = 1
+
+    def __call__(self, feat_data, item):
+        i = item * self.stride
+        return feat_data[slice(i + self.shift, i + self.shift + self.length)]
+
+    def __len__(self):
+        return (self.n - (self.shift + self.length) + 1) // self.stride
+
+
+@dtc.dataclass
+class AsFramedSlice(AsSlice):
+    frame_size: int = 1
+    as_strided: bool = False
+
+    def __call__(self, feat_data, item):
+        sliced = super(AsFramedSlice, self).__call__(feat_data, item)
+        if self.as_strided:
+            if isinstance(sliced, np.ndarray):
+                itemsize = sliced.dtype.itemsize
+                as_strided = lambda arr: np_as_strided(arr,
+                                                       shape=(self.length, self.frame_size),
+                                                       strides=(itemsize, itemsize))
+            else:
+                as_strided = lambda tensor: torch.as_strided(tensor,
+                                                             size=(self.length, self.frame_size),
+                                                             stride=(1, 1))
+
+            with torch.no_grad():
+                return as_strided(sliced)
+        else:
+            return sliced.reshape(-1, self.frame_size)
+
+
+@dtc.dataclass
+class Input:
+    db_key: str = ''
+    getter: Getter = Getter()
+    transform: Callable = lambda x: x
+
+    def __len__(self):
+        return len(self.getter)
+
+
+class Target(Input):
+    """exactly equivalent to Input, just makes code simpler to read."""
+    pass
+
+
+np_str_obj_array_pattern = re.compile(r'[SaUO]')
+
+
+def process_batch(batch, test=lambda x: False, func=lambda x: x):
+    """
+    recursively apply func to the elements of data if test(element) is True.
+    This is used in DefaultDataset to process elements (Input or Target) packed in tuples, list, dict etc...
+    """
+    elem_type = type(batch)
+    if test(batch):
+        return func(batch)
+    elif isinstance(batch, container_abcs.Mapping):
+        return {key: process_batch(batch[key], test, func) for key in batch}
+    elif isinstance(batch, tuple) and hasattr(batch, '_fields'):  # namedtuple
+        return elem_type(*(process_batch(d, test, func) for d in batch))
+    elif isinstance(batch, container_abcs.Sequence) and not isinstance(batch, string_classes):
+        return [process_batch(d, test, func) for d in batch]
+    else:
+        return batch
+
+
+def _is_batchitem(obj):
+    return isinstance(obj, (Input, Target))
+
+
+class DefaultDataset(Dataset):
+
+    def prepare_dataset(self, batch=tuple()):
+        super(Dataset, self).__init__()
+
+        # pass the lengths of the db features to the getters
+        def cache_lengths(feat):
+            if feat.getter.n is None:
+                setattr(feat.getter, 'n', len(getattr(self, feat.db_key)))
+            return feat
+
+        self.batch = process_batch(batch, _is_batchitem, cache_lengths)
+
+        # get the minimum length of all batchitems
+        self.N = float('inf')
+
+        def set_n_to_min(feat):
+            self.N = min(len(feat), self.N)
+            return feat
+
+        process_batch(self.batch, _is_batchitem, set_n_to_min)
+
+        return self
+
+    def __getitem__(self, item):
+        def get_data(feat):
+            return feat.transform(feat.getter(getattr(self, feat.db_key), item))
+
+        return process_batch(self.batch, _is_batchitem, get_data)
+
+    def __len__(self):
+        return self.N
+
+
+def _implements_mapstyle(obj):
+    if isinstance(obj, type):
+        getitem, flen = getattr(obj, '__getitem__', False), getattr(obj, '__len__', False)
+    else:
+        getitem, flen = getattr(obj.__class__, '__getitem__', False), getattr(obj.__class__, '__len__', False)
+
+    return getitem not in (Database.__getitem__, Dataset.__getitem__) and flen is not Database.__len__
+
+
+@dtc.dataclass
+class DataModule(pl.LightningDataModule):
+    """
+    encapsulate the logic for creating, initializing and serving databases & datasets
+    """
+
+    model: pl.LightningModule = None
+    db: [Database, Dataset, str] = None
+    keep_open: bool = True
+    sources: [str, Iterable[str]] = ''
+    schema: dict = dtc.field(default_factory=dict)
+    dataset_cls: type = None
+    in_mem_data: bool = True
+    splits: tuple = tuple()
+    loader_kwargs: dict = dtc.field(default_factory=dict)
+
+    def __post_init__(self):
+        super(DataModule, self).__init__()
+        self.has_split = False
+        self.full_ds = None
+        self.datasets = {}
+
+    def prepare_data(self, *args, **kwargs):
+        """
+        creates a db if necessary and possible
+        """
+        # get a db object (create it if necessary)
+        if isinstance(self.db, (str, os.PathLike)):
+            if not os.path.exists(self.db):
+                # user must provide sources and schema to build db
+                if not self.sources and self.schema:
+                    raise ValueError("You need to provide sources and a schema in order to create a db")
+                Database.create(self.db, self.sources, self.schema)
+
+    def _instantiate_db(self):
+        if isinstance(self.db, (str, os.PathLike)):
+            self.db = Database(self.db, keep_open=self.keep_open)
+        elif isinstance(self.db, (Database, Dataset)):
+            pass
+        else:
+            raise TypeError("Expected `db` to be of type str, Database or Dataset. Got " + str(type(self.db)))
+
+    def _init_ds(self, ds, stage):
+        # upgrade db to a Dataset if it isn't one already
+        if _implements_mapstyle(ds):
+            pass
+        # user provided Dataset class
+        elif self.dataset_cls is not None:
+            ds = ds.bind(self.dataset_cls)
+        # model has a batch signature
+        elif getattr(self.model, 'batch_signature', False):
+            ds = ds.bind(DefaultDataset)
+        else:
+            raise RuntimeError("couldn't instantiate a Dataset with the provided arguments")
+        if hasattr(ds, 'prepare_dataset') and hasattr(self.model, 'batch_signature'):
+            ds = ds.prepare_dataset(self.model.batch_signature(stage))
+        return ds
+
+    def _split(self):
+        if not self.splits:
+            self.datasets['full'] = self.full_ds
+            self.datasets['fit'] = self.full_ds
+        else:
+            sets = self.full_ds.split(self.splits)
+            # store the sets as attr
+            for ds, stage in zip(sets, ["fit", "val", "test"]):
+                if ds is not None:
+                    self.datasets[stage] = ds
+        setattr(self, 'has_split', True)
+        return self
+
+    def _move_to_mem(self, ds):
+        if isinstance(ds, Database):
+            if self.in_mem_data and torch.cuda.is_available():
+                ds.to_tensor()
+                ds.to("cuda")
+        return ds
+
+    def setup(self, stage=None):
+        if not self.has_prepared_data:
+            self.prepare_data()
+        if not isinstance(self.db, (Database, Dataset)):
+            self._instantiate_db()
+        if self.full_ds is None:
+            self.full_ds = self._init_ds(self.db, 'fit')
+            self.full_ds = self._move_to_mem(self.full_ds)
+        if not self.has_split:
+            self._split()
+        # update loader kwargs
+        if hasattr(self.model, 'loader_kwargs'):
+            self.loader_kwargs = self.model.loader_kwargs(stage, self)
+
+    def full_dataloader(self):
+        self.setup()
+        return DataLoader(self.full_ds, **self.loader_kwargs)
+
+    def train_dataloader(self):
+        if not self.has_setup_fit:
+            self.setup("fit")
+        return DataLoader(self.datasets['fit'], **self.loader_kwargs)
+
+    def val_dataloader(self):
+        if 'val' not in self.datasets:
+            return None
+        if not self.has_setup_fit:
+            self.setup("val")
+        kwargs = self.loader_kwargs.copy()
+        return DataLoader(self.datasets['val'], **kwargs)
+
+    def test_dataloader(self):
+        if 'test' not in self.datasets:
+            return None
+        if not self.has_setup_test:
+            self.setup("test")
+        kwargs = self.loader_kwargs.copy()
+        return DataLoader(self.datasets['test'], **kwargs)
+
+    def get_prompts(self, indices=tuple()):
+        ds = self.full_ds
+        if hasattr(ds, 'prepare_dataset') and hasattr(self.model, 'batch_signature'):
+            ds = ds.prepare_dataset(self.model.batch_signature('test'))
+        N = len(ds)
+        data = [ds[randint(0, N) if ix is None else ix] for ix in indices]
+        return torch.utils.data.dataloader.default_collate(data)
diff --git a/mimikit/data/feature.py b/mimikit/data/feature.py
new file mode 100644
index 0000000..c178ac1
--- /dev/null
+++ b/mimikit/data/feature.py
@@ -0,0 +1,139 @@
+import dataclasses as dtc
+
+__all__ = [
+    'Feature'
+]
+
+
+@dtc.dataclass
+class Feature:
+    """
+    class for centralizing methods pertaining to a particular type, class or format of data.
+    """
+    __ext__ = None
+
+    @property
+    def dim(self):
+        """
+        Returns
+        -------
+        dim: int, tuple
+            information about the shape / dimensionality of the feature
+        """
+        raise NotImplementedError
+
+    @property
+    def params(self):
+        """
+        Returns
+        -------
+        params: dict
+            the parameters of this Feature
+        """
+        return dtc.asdict(self)
+
+    @property
+    def encoders(self):
+        """
+        Returns
+        -------
+        encoders: dict
+            dict or mapping where keys are types and values callables taking only one argument
+        """
+        class AnyIdentity:
+            def __getitem__(self, item):
+                return lambda x: x
+        return AnyIdentity()
+
+    @property
+    def decoders(self):
+        """
+        Returns
+        -------
+        decoders: dict
+            dict or mapping where keys are types and values callables taking only one argument
+        """
+        class AnyIdentity:
+            def __getitem__(self, item):
+                return lambda x: x
+
+        return AnyIdentity()
+
+    def __call__(self, inputs, encode=True):
+        """
+        only apply the layer of encoder/decoder owned by this class. Use it for augmentation, transformation.
+
+        Parameters
+        ----------
+        inputs
+
+        encode: bool, opt
+            apply self.encoders if True, self.decoders otherwise
+
+        Returns
+        -------
+
+        """
+        funcs = self.encoders if encode else self.decoders
+        return funcs[type(inputs)](inputs)
+
+    def encode(self, inputs):
+        """
+        apply all the layers of encoders inherited by this class to `inputs`.
+
+        Parameters
+        ----------
+        inputs
+
+        Returns
+        -------
+
+        """
+        if hasattr(super(), 'encoders'):
+            inputs = super().encoders[type(inputs)](inputs)
+        return self.encoders[type(inputs)](inputs)
+
+    def decode(self, inputs):
+        """
+        apply all the layers of decoders inherited by this class to `inputs`.
+
+        Parameters
+        ----------
+        inputs
+
+        Returns
+        -------
+
+        """
+        inputs = self.decoders[type(inputs)](inputs)
+        if hasattr(super(), 'decoders'):
+            inputs = super().decoders[type(inputs)](inputs)
+        return inputs
+
+    def load(self, path):
+        """
+        reads the file at `path` and returns it as numpy array
+
+        Parameters
+        ----------
+        path
+
+        Returns
+        -------
+
+        """
+        # since this is optional, we don't raise any exception
+        pass
+
+    def post_create(self, db, schema_key):
+        # since this is optional, we don't raise any exception
+        pass
+
+    def play(self, inputs):
+        raise NotImplementedError
+
+    def display(self, inputs):
+        raise NotImplementedError
+
+    def write(self, filename, inputs):
+        raise NotImplementedError
diff --git a/mimikit/data/freqnet_db.py b/mimikit/data/freqnet_db.py
deleted file mode 100644
index c062778..0000000
--- a/mimikit/data/freqnet_db.py
+++ /dev/null
@@ -1,108 +0,0 @@
-import argparse
-from functools import partial
-from multiprocessing import cpu_count
-
-from . import Database, make_root_db, file_to_fft, AudioFileWalker
-from ..connectors.neptune import NeptuneConnector
-
-
-parser = argparse.ArgumentParser(prog="freqnet-db",
-                                 description="transform audio file to FFTs with specified parameters and put "
-                                             "them in .h5 Database that FreqNets can consume")
-
-parser.add_argument("target", type=str,
-                    help="the name (path) of the db you want to create")
-parser.add_argument("--roots", "-r",
-                    type=str, nargs="*",
-                    default=None,
-                    help="list of paths from which to search for audio files to include in the db")
-parser.add_argument("--files", "-f",
-                    type=str, nargs="*",
-                    default=None,
-                    help="list of audio files to include in the db")
-parser.add_argument("--n-fft", '-d',
-                    type=int, default=2048,
-                    help="the fft size used to transform the files (default=2048)")
-parser.add_argument("--hop-length", "-o",
-                    type=int, default=512,
-                    help="the hop length used to transform the files (default=512)")
-parser.add_argument("--sample-rate", "-s",
-                    type=int, default=22050,
-                    help="the sample rate used to transform the files (default=22050)")
-parser.add_argument("--neptune-path", '-p',
-                    type=str, default=None,
-                    help="path to the neptune.ai project or experiment you wish to upload the db to"
-                         " (requires that you stored your neptune api token in the environment of this script)")
-
-
-def freqnet_db(target,
-               roots=None,
-               files=None,
-               n_fft=2048,
-               hop_length=512,
-               sample_rate=22050,
-               neptune_path=None):
-    """
-    transform found audio files to STFTs with specified parameters and put
-    them in a .h5 ``Database``
-
-    Parameters
-    ----------
-    target : str
-        the name (or path) of the .h5 db you want to create
-    roots : str or list of str, optional
-        list of paths from which to search for audio files to include in the db
-    files : str or list of str, optional
-        list of audio files to include in the db
-
-        .. note::
-            if no ``roots`` and no ``files`` are provided, the function will search
-            for audio files in the current working directory.
-
-    n_fft : int, optional
-        the fft size used to transform the files (default is 2048)
-    hop_length : int, optional
-        the hop length used to transform the files (default is 512)
-    sample_rate : int, optional
-        the sample rate used to transform the files (default is 22050)
-    neptune_path : str or None, optional
-        name of the neptune.ai project you wish to upload the db to.
-        `neptune_path` is expected to be of the form `"username/project"`.
-        Requires that you stored your neptune api token in the environment.
-
-    Returns
-    -------
-    db : Database
-        the created db
-    """
-    namespace = argparse.Namespace(target=target, roots=roots, files=files, n_fft=n_fft,
-                                   hop_length=hop_length, sample_rate=sample_rate,
-                                   neptune_path=neptune_path)
-    main(namespace)
-    return Database(target)
-
-
-def main(namespace=None):
-    if namespace is None:
-        args = parser.parse_args()
-    else:
-        args = namespace
-    transform = partial(file_to_fft,
-                        n_fft=args.n_fft,
-                        hop_length=args.hop_length,
-                        sr=args.sample_rate)
-    if args.roots is None and args.files is None:
-        args.roots = "./"
-    walker = AudioFileWalker(roots=args.roots, files=args.files)
-    print("Found following audio files :", "\n", *["\t" + file + "\n" for file in list(walker)])
-    print("Making the database...")
-    make_root_db(args.target, roots=args.roots, files=args.files, extract_func=transform, n_cores=cpu_count()//2)
-
-    if args.neptune_path is not None:
-        path = args.neptune_path.split("/")
-        user, rest = path[0], "/".join(path[1:]) if path[-1] is not None else path[1]
-        connector = NeptuneConnector(user=user, setup={"db": rest})
-        db = Database(args.target)
-        print("uploading database to neptune...")
-        connector.upload_database("db", db)
-    print("done!")
diff --git a/mimikit/data/regions.py b/mimikit/data/regions.py
index 7149a2e..f0f378f 100644
--- a/mimikit/data/regions.py
+++ b/mimikit/data/regions.py
@@ -1,10 +1,14 @@
 import pandas as pd
 import numpy as np
 
+__all__ = [
+    'Regions'
+]
+
 
 def ssts(item, axis=0):
     arr = np.atleast_2d(item[['start', 'stop']].values)
-    slices = map(lambda a: (*[slice(0, None)] * axis, slice(a[0], a[1]), ), arr)
+    slices = map(lambda a: (*[slice(0, None)] * axis, slice(a[0], a[1]),), arr)
     return slices
 
 
@@ -120,6 +124,10 @@ def from_frame_definition(total_duration, frame_length, stride=1, butlasts=0):
         stops = starts + durations
         return Regions.from_start_stop(starts, stops, durations)
 
+    def to_labels(self):
+        return np.hstack([np.ones((tp.duration,), dtype=np.int) * tp.Index
+                          for tp in self.itertuples()])
+
 
 @pd.api.extensions.register_dataframe_accessor("soft_q")
 class SoftQueryAccessor:
diff --git a/mimikit/data/tbptt_sampler.py b/mimikit/data/tbptt_sampler.py
new file mode 100644
index 0000000..d20e7ee
--- /dev/null
+++ b/mimikit/data/tbptt_sampler.py
@@ -0,0 +1,42 @@
+import math
+
+import torch
+from torch.utils.data import Sampler, RandomSampler, BatchSampler
+
+
+__all__ = [
+    'TBPTTSampler'
+]
+
+
+class TBPTTSampler(Sampler):
+    """
+    yields batches of indices for performing Truncated Back Propagation Through Time
+    """
+
+    def __init__(self,
+                 n_samples,
+                 batch_size=64,  # nbr of "tracks" per batch
+                 chunk_length=8 * 16000,  # total length of a track
+                 seq_len=512  # nbr of samples per backward pass
+                 ):
+        super().__init__(None)
+        self.n_samples = n_samples
+        self.chunk_length = chunk_length
+        self.seq_len = seq_len
+        self.n_chunks = self.n_samples // self.chunk_length
+        self.remainder = self.n_samples % self.chunk_length
+        self.n_per_chunk = self.chunk_length // self.seq_len
+        self.batch_size = min(batch_size, self.n_chunks)
+
+    def __iter__(self):
+        smp = RandomSampler(torch.arange(self.n_chunks))
+        for top in BatchSampler(smp, self.batch_size, False):  # don't drop last!
+            offsets = torch.randint(0, self.remainder, (self.batch_size,))
+            top = tuple(o + (t * self.chunk_length) for t, o in zip(top, offsets))
+            for start in range(self.n_per_chunk):
+                # start indices of the batch
+                yield tuple(t + (start * self.seq_len) for t in top)
+
+    def __len__(self):
+        return int(max(1, math.floor(self.n_chunks / self.batch_size)) * self.n_per_chunk)
\ No newline at end of file
diff --git a/mimikit/data/transforms.py b/mimikit/data/transforms.py
deleted file mode 100644
index 5f66ee0..0000000
--- a/mimikit/data/transforms.py
+++ /dev/null
@@ -1,39 +0,0 @@
-import librosa
-from .regions import Regions
-
-N_FFT = 2048
-HOP_LENGTH = 512
-SR = 22050
-MU = 255
-
-
-def stft(file, n_fft=N_FFT, hop_length=HOP_LENGTH, sr=SR):
-    y, sr = librosa.load(file, sr=sr)
-    fft = librosa.stft(y, n_fft=n_fft, hop_length=hop_length)
-    # returns the feature and its attributes
-    return fft, dict(n_fft=n_fft, hop_length=hop_length, sr=sr)
-
-
-def file_to_fft(abs_path, n_fft=N_FFT, hop_length=HOP_LENGTH, sr=SR):
-    fft, params = stft(abs_path, n_fft, hop_length, sr)
-    fft = abs(fft)
-    regions = Regions.from_duration([fft.shape[1]])
-    params.update(dict(time_axis=0))
-    return dict(fft=(params, fft.T), regions=({}, regions))
-
-
-def mu_law_compress(file, mu=MU, sr=SR):
-    y, sr = librosa.load(file, sr=sr)
-    y = librosa.util.normalize(y)
-    qx = librosa.mu_compress(y, mu, quantize=True)
-    qx = qx + (MU + 1) // 2
-    return qx, dict(mu=mu, sr=sr)
-
-
-def file_to_qx(abs_path, mu=MU, sr=SR):
-    qx, params = mu_law_compress(abs_path, mu, sr)
-    regions = Regions.from_duration([qx.shape[0]])
-    return dict(qx=(params, qx.reshape(-1, 1)), regions=({}, regions))
-
-
-default_extract_func = file_to_fft
diff --git a/mimikit/extract/__init__.py b/mimikit/extract/__init__.py
index e69de29..e80f116 100644
--- a/mimikit/extract/__init__.py
+++ b/mimikit/extract/__init__.py
@@ -0,0 +1,4 @@
+from .segment import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/extract/segment.py b/mimikit/extract/segment.py
index 4890b8b..e7a9911 100644
--- a/mimikit/extract/segment.py
+++ b/mimikit/extract/segment.py
@@ -4,6 +4,18 @@
 from scipy.ndimage import convolve
 from ..data.regions import Regions
 
+__all__ = [
+    'from_recurrence_matrix'
+]
+
+
+def checker(N):
+    block = np.zeros((N * 2 + 1, N * 2 + 1), dtype=np.int32)
+    for k in range(-N, N + 1):
+        for l in range(-N, N + 1):
+            block[k + N, l + N] = np.sign(k) * np.sign(l)
+    return block / abs(block).sum()
+
 
 def from_recurrence_matrix(X,
                            L=6,
@@ -12,12 +24,6 @@ def from_recurrence_matrix(X,
                            bandwidth=1.,
                            thresh=0.2,
                            min_dur=4):
-    def checker(N):
-        block = np.zeros((N * 2 + 1, N * 2 + 1), dtype=np.int32)
-        for k in range(-N, N + 1):
-            for l in range(-N, N + 1):
-                block[k + N, l + N] = np.sign(k) * np.sign(l)
-        return block / abs(block).sum()
 
     R = recurrence_matrix(
         X, metric="cosine", mode="affinity",
@@ -32,4 +38,4 @@ def checker(N):
     mx = mx * (np.diff(mx, append=R.shape[0]) >= min_dur)
     mx = mx[mx > 0]
     stops = np.r_[mx, R.shape[0]]
-    return Regions.from_stop(stops)
\ No newline at end of file
+    return Regions.from_stop(stops)
diff --git a/mimikit/file_walker.py b/mimikit/file_walker.py
new file mode 100644
index 0000000..7967380
--- /dev/null
+++ b/mimikit/file_walker.py
@@ -0,0 +1,81 @@
+import os
+from typing import Iterable
+
+__all__ = [
+    "EXTENSIONS",
+    "FileWalker"
+]
+
+AUDIO_EXTENSIONS = {"wav", "aif", "aiff", "mp3", "m4a", "mp4"}
+IMAGE_EXTENSIONS = {'png', 'jpeg'}
+MIDI_EXTENSIONS = {"mid"}
+
+EXTENSIONS = dict(audio=AUDIO_EXTENSIONS,
+                  img=IMAGE_EXTENSIONS,
+                  midi=MIDI_EXTENSIONS,
+                  none={})
+
+
+class FileWalker:
+
+    def __init__(self, files_ext, sources=None):
+        """
+        recursively find files from `items` whose extensions match the ones implied in `files_ext`
+
+        Parameters
+        ----------
+        files_ext : str or list of str
+            type(s) of files' extensions to be matched. Must be either 'audio' or 'midi'.
+        sources : str or iterable of str
+            a single path (string, os.Path) or an iterable of paths. Each item can either be the root of
+            a directory which will be searched recursively or a single file.
+
+        Examples
+        --------
+        >>> files = list(FileWalker(files_ext='midi', sources=["my-root-dir", 'piece.mid']))
+
+        """
+        if isinstance(files_ext, str):
+            files_ext = [files_ext]
+        if not all(ext in EXTENSIONS.keys() for ext in files_ext):
+            raise ValueError("Expected all files_ext to be one of 'audio' or 'midi'.")
+        self._matching_ext = {ext for file_type in files_ext for ext in EXTENSIONS[file_type]}
+
+        generators = []
+
+        if sources is not None and isinstance(sources, Iterable):
+            if isinstance(sources, str):
+                if not os.path.exists(sources):
+                    raise FileNotFoundError("%s does not exist." % sources)
+                if os.path.isdir(sources):
+                    generators += [self.walk_root(sources)]
+                else:
+                    if self.is_matching_file(sources):
+                        generators += [[sources]]
+            else:
+                for item in sources:
+                    if not os.path.exists(item):
+                        raise FileNotFoundError("%s does not exist." % item)
+                    if os.path.isdir(item):
+                        generators += [self.walk_root(item)]
+                    else:
+                        if self.is_matching_file(item):
+                            generators += [[item]]
+
+        self.generators = generators
+
+    def __iter__(self):
+        for generator in self.generators:
+            for file in generator:
+                yield file
+
+    def walk_root(self, root):
+        for directory, _, files in os.walk(root):
+            for file in filter(self.is_matching_file, files):
+                yield os.path.join(directory, file)
+
+    def is_matching_file(self, filename):
+        # filter out hidden files
+        if os.path.split(filename.strip("/"))[-1].startswith("."):
+            return False
+        return os.path.splitext(filename)[-1].strip(".") in self._matching_ext
diff --git a/mimikit/freqnet/__init__.py b/mimikit/freqnet/__init__.py
deleted file mode 100644
index f152609..0000000
--- a/mimikit/freqnet/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from .base import FreqNetModel, FreqOptim, FreqData
-from .freaksnet import FreaksNet
-from .freq_layer import FreqLayer
-from .freqnet import FreqNet
-from .frexpert_mixture import FrexpertMixture
-from .layer_wise_loss import LayerWiseLossFreqNet
-from .hierarchical_kernel import HKFreqNet
diff --git a/mimikit/freqnet/base.py b/mimikit/freqnet/base.py
deleted file mode 100644
index aef1155..0000000
--- a/mimikit/freqnet/base.py
+++ /dev/null
@@ -1,255 +0,0 @@
-import torch
-from pytorch_lightning import LightningModule, LightningDataModule
-from torch.utils.data import DataLoader
-import librosa
-from abc import ABC
-
-from ..data import DataObject, FeatureProxy, HOP_LENGTH
-from ..kit import ShiftedSeqsPair, MMKHooks, LoggingHooks, tqdm
-
-
-class ManyOneCycleLR(torch.optim.lr_scheduler.OneCycleLR):
-    def step(self, epoch=None):
-        try:
-            super(ManyOneCycleLR, self).step(epoch=epoch)
-        except ValueError:
-            self.last_epoch = 0
-            super(ManyOneCycleLR, self).step(epoch=epoch)
-
-
-class FreqOptim:
-    """
-    simple class to modularize the optimization setup used in a ``FreqNetModel``.
-
-    Here the method ``configure_optimizers()`` returns an ``Adam`` optimizer and a slightly modified
-    ``OneCycleLR`` scheduler (the only modification being that it won't raise a ``ValueError`` if you use for more
-    steps than what it expects and will instead restarts its cycle)
-    """
-    def __init__(self,
-                 model,
-                 max_lr=1e-3,
-                 betas=(.9, .9),
-                 div_factor=3.,
-                 final_div_factor=1.,
-                 pct_start=.25,
-                 cycle_momentum=False):
-        self.model = model
-        self.max_lr = max_lr
-        self.betas = betas
-        self.div_factor = div_factor
-        self.final_div_factor = final_div_factor
-        self.pct_start = pct_start
-        self.cycle_momentum = cycle_momentum
-        # those are set in setup when we know what the trainer and datamodule got
-        self.steps_per_epoch = None
-        self.max_epochs = None
-
-    def configure_optimizers(self):
-        self.opt = torch.optim.Adam(self.model.parameters(), lr=self.max_lr, betas=self.betas)
-        self.sched = ManyOneCycleLR(self.opt,
-                                    steps_per_epoch=self.steps_per_epoch,
-                                    epochs=self.max_epochs,
-                                    max_lr=self.max_lr, div_factor=self.div_factor,
-                                    final_div_factor=self.final_div_factor,
-                                    pct_start=self.pct_start,
-                                    cycle_momentum=self.cycle_momentum
-                                    )
-        return [self.opt], [{"scheduler": self.sched, "interval": "step", "frequency": 1}]
-
-
-class FreqData(LightningDataModule):
-    """
-    boilerplate subclass of ``pytorch_lightning.LightningDataModule`` to handle the data of a ``FreqNetModel``.
-
-    the data is passed through ``data_object``, wrapped in a ``ShiftedSeqsPair`` in ``prepare_data()`` and
-    in ``setup("fit")`` it is moved to the gpu if ``in_mem_data`` is ``True`` and split into train, val and test sets
-    according to ``splits``
-    """
-    def __init__(self,
-                 model,
-                 data_object=None,
-                 input_seq_length=64,
-                 batch_size=64,
-                 in_mem_data=True,
-                 splits=None,
-                 **loader_kwargs,
-                 ):
-        super(FreqData, self).__init__()
-        self.model = model
-        self.ds = DataObject(data_object)
-        self.input_seq_length = input_seq_length
-        self.batch_size = batch_size
-        self.in_mem_data = in_mem_data
-        self.splits = splits
-        loader_kwargs.setdefault("drop_last", False)
-        loader_kwargs.setdefault("shuffle", True)
-        self.loader_kwargs = loader_kwargs
-        self.train_ds, self.val_ds, self.test_ds = None, None, None
-
-    def prepare_data(self, *args, **kwargs):
-        if not (getattr(self.model, "targets_shifts_and_lengths", False)):
-            raise TypeError("Expected `model` to implement `targets_shifts_and_lengths(input_length)`"
-                            " in order to compute the right slices for the batches")
-        targets_def = self.model.targets_shifts_and_lengths(self.input_seq_length)
-        wrapper = ShiftedSeqsPair(self.input_seq_length, targets_def)
-        self.ds = wrapper(self.ds)
-
-    def setup(self, stage=None):
-        if stage == "fit":
-            if self.in_mem_data and torch.cuda.is_available():
-                self.ds.to_tensor()
-                self.ds.to("cuda")
-            if self.splits is None:
-                self.splits = (1.,)
-            sets = self.ds.split(self.splits)
-            for ds, attr in zip(sets, ["train_ds", "val_ds", "test_ds"]):
-                setattr(self, attr, ds)
-
-    def train_dataloader(self):
-        if not self.has_prepared_data:
-            self.prepare_data()
-        if not self.has_setup_fit:
-            self.setup("fit")
-        return DataLoader(self.train_ds, batch_size=self.batch_size, **self.loader_kwargs)
-
-    def val_dataloader(self, shuffle=False):
-        has_val = len(self.splits) >= 2 and self.splits[1] is not None
-        if not has_val:
-            return None
-        if not self.has_prepared_data:
-            self.prepare_data()
-        if not self.has_setup_fit:
-            self.setup("fit")
-        kwargs = self.loader_kwargs.copy()
-        kwargs["shuffle"] = shuffle
-        return DataLoader(self.val_ds, batch_size=self.batch_size, **kwargs)
-
-    def test_dataloader(self, shuffle=False):
-        has_test = len(self.splits) >= 3 and self.splits[2] is not None
-        if not has_test:
-            return None
-        if not self.has_prepared_data:
-            self.prepare_data()
-        if not self.has_setup_fit:
-            self.setup("test")
-        kwargs = self.loader_kwargs.copy()
-        kwargs["shuffle"] = shuffle
-        return DataLoader(self.test_ds, batch_size=self.batch_size, **kwargs)
-
-
-class FreqNetModel(MMKHooks,
-                   LoggingHooks,
-                   LightningModule,
-                   ABC):
-    """
-    base class for all ``FreqNets`` that handles optim, data and a few handy methods like ``generate()``
-    """
-    @property
-    def data(self):
-        """short-hand to quickly access the data object passed to the constructor"""
-        return self.datamodule.ds.data
-
-    def __init__(self,
-                 data_object=None,
-                 input_seq_length=64,
-                 batch_size=64,
-                 in_mem_data=True,
-                 splits=[.8, .2],
-                 max_lr=1e-3,
-                 betas=(.9, .9),
-                 div_factor=3.,
-                 final_div_factor=1.,
-                 pct_start=.25,
-                 cycle_momentum=False,
-                 **loaders_kwargs):
-        super(FreqNetModel, self).__init__()
-        # dimensionality of inputs is automatically available
-        if data_object is not None and not isinstance(data_object, str):
-            self.input_dim = data_object.shape[-1]
-            self.datamodule = FreqData(self, data_object, input_seq_length, batch_size,
-                                       in_mem_data, splits, **loaders_kwargs)
-        else:
-            raise ValueError("Please pass a valid data_object to instantiate a FreqNetModel."
-                             " If you are loading from a checkpoint, you can do so by modifying your method call to :\n"
-                             "FreqNetModel.load_from_checkpoint(path_to_ckpt, data_object=my_data_object)")
-        self.optim = FreqOptim(self, max_lr, betas, div_factor, final_div_factor, pct_start,
-                               cycle_momentum)
-        # calling this updates self.hparams from any subclass : call it when subclassing!
-        self.save_hyperparameters()
-
-    def training_step(self, batch, batch_idx):
-        batch, target = batch
-        output = self.forward(batch)
-        recon = self.loss_fn(output, target)
-        return {"loss": recon}
-
-    def validation_step(self, batch, batch_idx):
-        batch, target = batch
-        output = self.forward(batch)
-        recon = self.loss_fn(output, target)
-        return {"val_loss": recon}
-
-    def setup(self, stage: str):
-        if stage == "fit":
-            self.logger.log_hyperparams(self.hparams)
-            self.optim.max_epochs = self.trainer.max_epochs
-            self.optim.steps_per_epoch = len(self.datamodule.train_dataloader())
-
-    def configure_optimizers(self):
-        return self.optim.configure_optimizers()
-
-    def _set_hparams(self, hp):
-        """
-        Small handy hook to
-            - copy params of a db in hparams
-            - make sure we don't store arrays or tensors in hparams by replacing them
-            by their string
-        """
-        # replace inputs passed to the hparams with their string
-        if "data_object" in hp:
-            if isinstance(hp["data_object"], FeatureProxy):
-                # copy the attrs of the db to hparams
-                for k, v in hp["data_object"].attrs.items():
-                    hp[k] = v
-            string = repr(hp["data_object"])
-            hp["data_object"] = string[:88] + ("..." if len(string) > 88 else "")
-        super(FreqNetModel, self)._set_hparams(hp)
-
-    # Convenience Methods for generative audio models (see also, `LoggingHooks.log_audio`):
-
-    def random_train_example(self):
-        return next(iter(self.datamodule.train_dataloader()))[0][0:1]
-
-    def random_val_example(self):
-        return next(iter(self.datamodule.val_dataloader(shuffle=True)))[0][0:1]
-
-    def generation_slices(self):
-        raise NotImplementedError("subclasses of `FreqNetModel` have to implement `generation_slices`")
-
-    def targets_shifts_and_lengths(self, input_length):
-        raise NotImplementedError("subclasses of `FreqNetModel` have to implement `targets_shifts_and_lengths`")
-
-    def generate(self, prompt, n_steps, hop_length=HOP_LENGTH, time_domain=True):
-        was_training = self.training
-        self.eval()
-        initial_device = self.device
-        self.to("cuda" if torch.cuda.is_available() else "cpu")
-        if not isinstance(prompt, torch.Tensor):
-            prompt = torch.from_numpy(prompt)
-        if len(prompt.shape) < 3:
-            prompt = prompt.unsqueeze(0)
-        input_slice, output_slice = self.generation_slices()
-        generated = prompt.to(self.device)
-        for _ in tqdm(range(n_steps), desc="Generate", dynamic_ncols=True, leave=False, unit="step"):
-            with torch.no_grad():
-                out = self(generated[:, input_slice])
-                generated = torch.cat((generated, out[:, output_slice]), dim=1)
-        if time_domain:
-            generated = generated.transpose(1, 2).squeeze()
-            generated = librosa.griffinlim(generated.cpu().numpy(), hop_length=hop_length, n_iter=64)
-            generated = torch.from_numpy(generated)
-        else:  # for consistency with time_domain :
-            generated = generated.cpu()
-        self.to(initial_device)
-        self.train() if was_training else None
-        return generated.unsqueeze(0)
diff --git a/mimikit/freqnet/freaksnet.py b/mimikit/freqnet/freaksnet.py
deleted file mode 100644
index 08adfa5..0000000
--- a/mimikit/freqnet/freaksnet.py
+++ /dev/null
@@ -1,59 +0,0 @@
-import torch.nn as nn
-
-from .freq_layer import FreqLayer
-from .freqnet import FreqNet
-from .modules import mean_L1_prop, GatedLinearInput, AbsLinearOutput
-
-
-class FreaksNet(FreqNet):
-    """
-    Simple adaptation of FreqNet to have varying parameters for each layer.
-    to each block of layer corresponds a tuple of values for the params listed in FeqNet.LAYER_PARAMS
-    """
-
-    def __init__(self,
-                 loss_fn=mean_L1_prop,
-                 input_dim=1025,
-                 model_dim=512,
-                 groups=((1, 1, 1),),
-                 n_layers=(3,),
-                 kernel_size=((2, 2, 2),),
-                 strict=((False, True, False),),
-                 accum_outputs=((-1, 1, -1),),
-                 concat_outputs=((1, -1, 1),),
-                 pad_input=((0, 0, 0),),
-                 learn_padding=((False, False, False),),
-                 with_skip_conv=((False, True, False,), ),
-                 with_residual_conv=((True, True, False), ),
-                 **data_optim_kwargs):
-        super(FreaksNet, self).__init__(**data_optim_kwargs)
-        self._loss_fn = loss_fn
-        self.input_dim = input_dim
-        self.model_dim = model_dim
-        self.groups = groups
-        self.n_layers = n_layers
-        self.strict = strict
-        self.kernel_size = kernel_size
-        self.accum_outputs = accum_outputs
-        self.concat_outputs = concat_outputs
-        self.pad_input = pad_input
-        self.learn_padding = learn_padding
-        self.with_skip_conv = with_skip_conv
-        self.with_residual_conv = with_residual_conv
-
-        # Input Encoder
-        self.inpt = GatedLinearInput(self.input_dim, self.model_dim)
-
-        # Auto-regressive Part
-        layer_kwargs = {attr: getattr(self, attr) for attr in self.LAYER_KWARGS}
-        # for simplicity we keep all the layers in a flat list
-        self.layers = nn.ModuleList([
-            FreqLayer(layer_index=i, input_dim=model_dim, layer_dim=model_dim,
-                      **{k: v[j][i] for k, v in layer_kwargs.items()})
-            for j, n_layers in enumerate(self.n_layers) for i in range(n_layers)
-        ])
-
-        # Output Decoder
-        self.outpt = AbsLinearOutput(self.model_dim, self.input_dim)
-
-        self.save_hyperparameters()
diff --git a/mimikit/freqnet/freq_layer.py b/mimikit/freqnet/freq_layer.py
deleted file mode 100644
index 34309f9..0000000
--- a/mimikit/freqnet/freq_layer.py
+++ /dev/null
@@ -1,152 +0,0 @@
-import torch
-import torch.nn as nn
-import numpy as np
-import math
-
-from .modules import GatedConv, LearnablePad1d
-
-
-def accum(x, y, shift=1):
-    if shift == 0:
-        if x.size(-1) == y.size(-1):
-            return x + y
-        else:
-            raise ValueError("zero shift with size: %i and %i" % (x.size(-1), y.size(-1)))
-    side = np.sign(shift)
-    aligned_x = slice(*((shift, None)[::side]))
-    n_aligned = x.size(-1) - abs(shift)
-    aligned_y = slice(*((None, n_aligned * side)[::side]))
-    compl_y = slice(*((n_aligned * side, None)[::side]))
-    aligned = x[:, :, aligned_x] + y[:, :, aligned_y]
-    rest = y[:, :, compl_y]
-    return torch.cat((aligned, rest)[::side], dim=-1)
-
-
-def concat(x, y, shift=1):
-    if shift == 0:
-        if x.size(-1) != y.size(-1):
-            return concat(x, y, - x.size(-1) - y.size(-1))
-        return y
-    side = np.sign(shift)
-    compl_x = slice(*((None, shift)[::side]))
-    return torch.cat((x[:, :, compl_x], y)[::side], dim=-1)
-
-
-class FreqLayer(nn.Module):
-    stride = 1
-    bias = True
-    sides = {1: 1, -1: -1, 0: 0, "left": 1, "right": -1}
-
-    def __init__(self,
-                 layer_index,
-                 input_dim=512,
-                 layer_dim=512,
-                 groups=1,
-                 strict=False,
-                 kernel_size=2,
-                 accum_outputs=0,
-                 concat_outputs=0,
-                 pad_input=0,
-                 learn_padding=False,
-                 with_skip_conv=False,
-                 with_residual_conv=False,
-                 ):
-        super(FreqLayer, self).__init__()
-        self.layer_index = layer_index
-        self.input_dim = input_dim
-        self.layer_dim = layer_dim
-        self.groups = groups
-        self.strict = strict
-        self.kernel_size = kernel_size
-        if pad_input != 0 and concat_outputs != 0:
-            raise ValueError("pad_input and concat_outputs can not be both non zero")
-        # accum and concat have opposite signs :
-        # accum left is a shift from the end (negative), whereas concat left is a shift from the beginning (positive)
-        self.accum_outputs = - self.sides.get(accum_outputs, 0)
-        self.concat_outputs = self.sides.get(concat_outputs, 0)
-        self.pad_input = self.sides.get(pad_input, 0)
-        self.learn_padding = learn_padding
-        self.with_skip_conv = with_skip_conv
-        self.with_residual_conv = with_residual_conv
-
-        self.pad = LearnablePad1d(self.input_dim, self.padding, self.learn_padding)
-
-        convs_kwargs = dict(dilation=self.dilation, stride=self.stride,
-                            groups=self.groups, bias=self.bias)
-        self.gate = GatedConv(self.input_dim, self.layer_dim, self.kernel_size, **convs_kwargs)
-
-        self.skips = nn.Conv1d(self.layer_dim, self.layer_dim, kernel_size=1, **convs_kwargs) \
-            if self.with_skip_conv else None
-        self.residuals = nn.Conv1d(self.layer_dim, self.input_dim, kernel_size=1, **convs_kwargs) \
-            if self.with_residual_conv else None
-
-    def forward(self, x, skip=None):
-        input = self.pad(x)
-
-        y = self.gate(input)
-        h = self.skips(y) if self.with_skip_conv else None
-        y = self.residuals(y) if self.with_residual_conv else y
-
-        accum_outputs, concat_outputs = self.accum_outputs, self.concat_outputs
-        shift = self.rel_shift()
-
-        if accum_outputs:
-            y = accum(x, y, shift * accum_outputs)
-
-        if concat_outputs:
-            y = concat(x, y, shift * concat_outputs)
-
-        if self.with_skip_conv:
-            if skip is None:
-                skip = torch.zeros_like(h).to(h)
-            if accum_outputs:
-                # we do it on the same side as the outputs
-                skip = accum(skip, h, shift * accum_outputs)
-            else:
-                # we accum on the right anyway, otherwise skips don't make any sense,
-                # since they aren't inputs to anything...
-                skip = accum(skip, h, shift * 1)
-            if concat_outputs:
-                skip = concat(skip, h, shift * concat_outputs)
-        else:
-            # still need to output a tensor
-            skip = y
-        return y, skip
-
-    @property
-    def dilation(self):
-        return self.kernel_size ** self.layer_index
-
-    @property
-    def padding(self):
-        """
-        signed amount of padding necessary to output as many time-steps as there were in the inputs.
-        The sign of padding corresponds to the desired side (1: left, -1: right)
-        """
-        return self.pad_input * (self.kernel_size - 1) * self.dilation
-
-    def receptive_field(self):
-        """
-        amount of inputs necessary for 1 output (this is independent of the shift!)
-        """
-        return self.kernel_size ** (self.layer_index + 1)
-
-    def shift(self):
-        """total shift at this layer wrt. to the beginning of its block"""
-        return self.receptive_field() + self.layer_index * int(self.strict)
-
-    def rel_shift(self):
-        """relative shift at this layer wrt. to the previous layer"""
-        return (int(self.strict) + (self.kernel_size - 1) * self.dilation) if self.pad_input == 0 else int(self.strict)
-
-    def output_length(self, input_length):
-        if abs(self.concat_outputs):
-            # concats appends past steps which grow by 1 at each layer when strict=True
-            return input_length + int(self.strict)
-        if abs(self.pad_input):
-            # no matter what, padding input gives the same output shape
-            return input_length
-        # output is gonna be less than input
-        numerator = input_length - self.dilation * (self.kernel_size - 1) - 1
-        denominator = self.stride
-        return math.floor(1 + numerator / denominator)
diff --git a/mimikit/freqnet/freqnet.py b/mimikit/freqnet/freqnet.py
deleted file mode 100644
index 78f39d4..0000000
--- a/mimikit/freqnet/freqnet.py
+++ /dev/null
@@ -1,126 +0,0 @@
-import torch.nn as nn
-import numpy as np
-
-from .modules import GatedLinearInput, AbsLinearOutput, mean_L1_prop
-from .freq_layer import FreqLayer
-from .base import FreqNetModel
-
-
-class FreqNet(FreqNetModel):
-    LAYER_KWARGS = ["groups", "strict", "accum_outputs", "concat_outputs", "kernel_size",
-                    "pad_input", "learn_padding", "with_skip_conv", "with_residual_conv"]
-
-    def __init__(self,
-                 loss_fn=mean_L1_prop,
-                 model_dim=512,
-                 groups=1,
-                 n_layers=(2,),
-                 kernel_size=2,
-                 strict=False,
-                 accum_outputs=0,
-                 concat_outputs=0,
-                 pad_input=0,
-                 learn_padding=False,
-                 with_skip_conv=True,
-                 with_residual_conv=True,
-                 **data_optim_kwargs):
-        super(FreqNet, self).__init__(**data_optim_kwargs)
-        self._loss_fn = loss_fn
-        self.model_dim = model_dim
-        self.groups = groups
-        self.n_layers = n_layers
-        self.kernel_size = kernel_size
-        self.strict = strict
-        self.accum_outputs = accum_outputs
-        self.concat_outputs = concat_outputs
-        self.pad_input = pad_input
-        self.learn_padding = learn_padding
-        self.with_skip_conv = with_skip_conv
-        self.with_residual_conv = with_residual_conv
-
-        # Input Encoder
-        self.inpt = GatedLinearInput(self.input_dim, self.model_dim)
-
-        # Auto-regressive Part
-        layer_kwargs = {attr: getattr(self, attr) for attr in self.LAYER_KWARGS}
-        # for simplicity we keep all the layers in a flat list
-        self.layers = nn.ModuleList([
-            FreqLayer(layer_index=i, input_dim=model_dim, layer_dim=model_dim, **layer_kwargs)
-            for n_layers in self.n_layers for i in range(n_layers)
-        ])
-
-        # Output Decoder
-        self.outpt = AbsLinearOutput(self.model_dim, self.input_dim)
-
-        self.save_hyperparameters()
-
-    def forward(self, x):
-        """
-        """
-        x = self.inpt(x)
-        skips = None
-        for layer in self.layers:
-            x, skips = layer(x, skips)
-        x = self.outpt(skips)
-        return x
-
-    def loss_fn(self, predictions, targets):
-        return self._loss_fn(predictions, targets)
-
-    def all_rel_shifts(self):
-        """sequence of shifts from one layer to the next"""
-        return tuple(layer.rel_shift() for layer in self.layers)
-
-    def shift(self):
-        """total shift of the network"""
-        if not self.strict and (self.pad_input == 1 or self.concat_outputs == 1):
-            return 1
-        elif self.pad_input == -1 or self.concat_outputs == -1:
-            return self.receptive_field()
-        elif self.strict and self.concat_outputs == 1:
-            return 0
-        elif self.strict and self.pad_input == 1:
-            return len(self.layers)
-        else:
-            return sum(self.all_rel_shifts()) + int(not self.strict)
-
-    def all_shifts(self):
-        """the accumulated shift at each layer"""
-        return tuple(np.cumsum(self.all_rel_shifts()) + int(not self.strict))
-
-    def receptive_field(self):
-        block_rf = []
-        for i, layer in enumerate(self.layers[:-1]):
-            if self.layers[i + 1].layer_index == 0:
-                block_rf += [layer.receptive_field() - 1]
-        block_rf += [self.layers[-1].receptive_field()]
-        return sum(block_rf)
-
-    def output_length(self, input_length):
-        return self.all_output_lengths(input_length)[-1]
-
-    def all_output_lengths(self, input_length):
-        out_length = input_length
-        lengths = []
-        for layer in self.layers:
-            out_length = layer.output_length(out_length)
-            lengths += [out_length]
-        return tuple(lengths)
-
-    def targets_shifts_and_lengths(self, input_length):
-        return [(self.shift(), self.output_length(input_length))]
-
-    def generation_slices(self):
-        # input is always the last receptive field
-        input_slice = slice(-self.receptive_field(), None)
-        if not self.strict and (self.pad_input == 1 or self.concat_outputs == 1):
-            # then there's only one future time-step : the last output
-            output_slice = slice(-1, None)
-        elif self.strict and (self.pad_input == 1 or self.concat_outputs == 1):
-            # then there are as many future-steps as they are layers and they all are
-            # at the end of the outputs
-            output_slice = slice(-len(self.layers), None)
-        else:
-            # for all other cases, the first output has the shift of the whole network
-            output_slice = slice(None, 1)
-        return input_slice, output_slice
diff --git a/mimikit/freqnet/frexpert_mixture.py b/mimikit/freqnet/frexpert_mixture.py
deleted file mode 100644
index 49c2fa6..0000000
--- a/mimikit/freqnet/frexpert_mixture.py
+++ /dev/null
@@ -1,47 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .modules.loss_functions import mean_L1_prop
-from .base import FreqNetModel
-from .freqnet import FreqNet
-
-
-class FrexpertMixture(FreqNetModel):
-    """
-    Use the Base FreqNet class to load a bunch of FreqNets for a mixture
-    """
-
-    def __init__(self, loss_fn=mean_L1_prop, **data_optim_kwargs):
-        super(FrexpertMixture, self).__init__(**data_optim_kwargs)
-        self.loss_fn = loss_fn
-        self.save_hyperparameters()
-        data_obj = data_optim_kwargs.get("data_object", None)
-        # we use nets instead of layers, but we still store them in the layers attributes!
-        self.nets = nn.ModuleList([
-            FreqNet(data_object=data_obj,
-                    model_dim=512,
-                    groups=1,
-                    n_layers=(3,),
-                    strict=False,
-                    accum_outputs=1),
-            FreqNet(data_object=data_obj,
-                    model_dim=1024,
-                    groups=2,
-                    n_layers=(3,),
-                    strict=False,
-                    accum_outputs=-1)
-        ])
-        # a simple MLP to output the weight of each net given the input
-        self.proba_f = nn.Sequential(nn.Linear(1025, 1025), nn.Tanh(),
-                                     nn.Linear(1025, 2), nn.Softmax(dim=-1))
-
-    def forward(self, x):
-        ys = torch.stack([net(x) for net in self.nets], dim=2)
-        # we crop the begining of x to have matching shapes
-        pr_ys = self.proba_f(x[:, -ys.size(1):]).unsqueeze(-1)
-        # sum by the "net-dim"
-        y = (pr_ys * ys).sum(dim=2)
-        return y
-
-    def targets_shifts_and_lengths(self, input_length):
-        return self.nets[0].targets_shifts_and_lengths(input_length)
\ No newline at end of file
diff --git a/mimikit/freqnet/hierarchical_kernel.py b/mimikit/freqnet/hierarchical_kernel.py
deleted file mode 100644
index 5c24f6d..0000000
--- a/mimikit/freqnet/hierarchical_kernel.py
+++ /dev/null
@@ -1,84 +0,0 @@
-import torch
-import torch.nn as nn
-
-from .freqnet import FreqNet
-from .modules import mean_L1_prop
-
-
-class HKFreqNet(FreqNet):
-    """
-    Add a hierarchical quality to the conditioning sequence by combining FreqLayers in a
-    final Conv2d making depth in the network proportional to the distance to the predicted time-step.
-
-    N.B : Could be fun to try the other way around : the closer to now -> the deeper in the net.
-    """
-
-    def __init__(self,
-                 loss_fn=mean_L1_prop,
-                 model_dim=512,
-                 groups=1,
-                 n_layers=2,
-                 strict=False,
-                 kernel_size=2,
-                 accum_outputs=0,
-                 concat_outputs=0,
-                 pad_input=0,
-                 learn_padding=False,
-                 with_skip_conv=True,
-                 with_residual_conv=True,
-                 **data_optim_kwargs):
-        super(HKFreqNet, self).__init__(loss_fn, model_dim, groups, (n_layers,), kernel_size, strict,
-                                        accum_outputs, concat_outputs, pad_input, learn_padding,
-                                        with_skip_conv, with_residual_conv,
-                                        **data_optim_kwargs)
-        # the 4th dimension in the final conv2d is 2 encoder outputs + n_conv_layers
-        self.conv_f = nn.Conv2d(model_dim, model_dim, (1, n_layers + 2), groups=groups)
-        self.conv_g = nn.Conv2d(model_dim, model_dim, (1, n_layers + 2), groups=groups)
-
-    def forward(self, x):
-        """
-
-        """
-        x = self.inpt(x)
-
-        # the conv2d layer adds a power of 2 to the n standard layers before it
-        n_outs = self.output_length(x.size(-1))
-
-        # stack the 2 encoder outputs
-        hk_input = torch.stack((x[:, :, -n_outs - 1:-1], x[:, :, -n_outs:]), dim=-1)
-        skip = None
-
-        for layer in self.layers:
-            x, skip = layer(x, skip)
-            # we only keep n_outs outputs finishing at the -`layer.shift()` step
-            to_stack = slice(-n_outs - layer.shift(), -layer.shift())
-            hk_input = torch.cat((hk_input, skip[:, :, to_stack].unsqueeze(-1)), dim=-1)
-
-        # Do the Gated2dConv on the gathered pieces of layers
-        f = nn.Tanh()(self.conv_f(hk_input))
-        g = nn.Sigmoid()(self.conv_g(hk_input))
-
-        y = self.outpt((f * g).squeeze(-1))
-        return y
-
-    def receptive_field(self):
-        # rf of the standard layers
-        rf = super(HKFreqNet, self).receptive_field()
-        # rf WITH HK layer :
-        return rf * 2
-
-    def all_rel_shifts(self):
-        """sequence of shifts from one layer to the next"""
-        stack = super(HKFreqNet, self).all_rel_shifts()
-        # add the last HK layer
-        return (*stack, 2 * stack[-1])
-
-    def all_output_lengths(self, input_length):
-        out_length = input_length
-        lengths = []
-        for layer in self.layers:
-            out_length = layer.output_length(out_length)
-            lengths += [out_length]
-        # add the last HK layer
-        lengths += [out_length - (2 ** (len(self.layers)))]
-        return tuple(lengths)
diff --git a/mimikit/freqnet/layer_wise_loss.py b/mimikit/freqnet/layer_wise_loss.py
deleted file mode 100644
index 3d6d785..0000000
--- a/mimikit/freqnet/layer_wise_loss.py
+++ /dev/null
@@ -1,54 +0,0 @@
-import torch
-
-from .freqnet import FreqNet
-
-
-class LayerWiseLossFreqNet(FreqNet):
-
-    def forward(self, x):
-        """
-        """
-        x = self.inpt(x)
-        # we collect all the layers outputs
-        skips = None
-        outputs = []
-        for layer in self.layers:
-            x, skips = layer(x, skips)
-            # decode each layer's output
-            outputs += [self.outpt(skips)]
-        if self.training:
-            # pass them all to the loss function
-            return tuple(outputs)
-        # else redirect to a special method for inference
-        return self.infer_layer_wise(outputs)
-    
-    def training_step(self, batch, batch_idx):
-        batch = batch[0], batch[1:]
-        return super(LayerWiseLossFreqNet, self).training_step(batch, batch_idx)
-
-    def infer_layer_wise(self, outputs):
-        """
-        Place to implement the inference method when trained layer-wise.
-        this default method returns the last output if the network isn't strict, otherwise, it collects the future
-        time-steps where they are first predicted in the network (Xt+1 @ layer_1, Xt+2 @ layer_2...)
-        @param outputs: list containing all the layers outputs
-        @return: the prediction of the network
-        """
-        if not self.strict:
-            return outputs[-1][-1]
-        # we collect one future time step pro layer
-        rf = self.receptive_field()
-        future_steps = []
-        for layer_out, shift in zip(outputs, self.all_shifts()):
-            # we keep everything from the last layer, hence the condition on the slice
-            step = slice(rf - shift, rf - shift +1) if rf - shift > 0 else slice(None, None)
-            future_steps += [layer_out[:, step]]
-        return torch.cat(future_steps, dim=1)
-
-    def loss_fn(self, predictions, targets):
-        denominator = len(self.layers)
-        return sum(self._loss_fn(pred, trg) / denominator for pred, trg in zip(predictions, targets))
-
-    def targets_shifts_and_lengths(self, input_length):
-        shifts, lengths = self.all_shifts(), self.all_output_lengths(input_length)
-        return list(zip(shifts, lengths))
\ No newline at end of file
diff --git a/mimikit/freqnet/modules/__init__.py b/mimikit/freqnet/modules/__init__.py
deleted file mode 100644
index e39d94e..0000000
--- a/mimikit/freqnet/modules/__init__.py
+++ /dev/null
@@ -1,4 +0,0 @@
-from .io import *
-from .gated_units import *
-from .loss_functions import *
-from .padding import *
\ No newline at end of file
diff --git a/mimikit/freqnet/modules/gated_units.py b/mimikit/freqnet/modules/gated_units.py
deleted file mode 100644
index 09793e0..0000000
--- a/mimikit/freqnet/modules/gated_units.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import torch.nn as nn
-
-
-class GatedLinear(nn.Module):
-    def __init__(self, in_dim, out_dim, **kwargs):
-        super(GatedLinear, self).__init__()
-        self.fcg = nn.Linear(in_dim, out_dim, **kwargs)
-        self.fcf = nn.Linear(in_dim, out_dim, **kwargs)
-
-    def forward(self, x):
-        f = nn.Tanh()(self.fcf(x))
-        g = nn.Sigmoid()(self.fcg(x))
-        return f * g
-
-
-class GatedConv(nn.Module):
-    def __init__(self, c_in, c_out, kernel_size=2, transpose=False, **kwargs):
-        super(GatedConv, self).__init__()
-        mod = nn.Conv1d if not transpose else nn.ConvTranspose1d
-        self.conv_f = mod(c_in, c_out, kernel_size, **kwargs)
-        self.conv_g = mod(c_in, c_out, kernel_size, **kwargs)
-        self.tanh = nn.Tanh()
-        self.sig = nn.Sigmoid()
-
-    def forward(self, x):
-        f = self.tanh(self.conv_f(x))
-        g = self.sig(self.conv_g(x))
-        return f * g
diff --git a/mimikit/freqnet/modules/io.py b/mimikit/freqnet/modules/io.py
deleted file mode 100644
index 040823d..0000000
--- a/mimikit/freqnet/modules/io.py
+++ /dev/null
@@ -1,28 +0,0 @@
-import torch.nn as nn
-from .gated_units import GatedLinear
-
-
-class PermuteTF(nn.Module):
-    def forward(self, x):
-        return x.transpose(1, 2).contiguous()
-
-
-class GatedLinearInput(GatedLinear):
-    def __init__(self, in_dim, out_dim, permute_tf=True):
-        super(GatedLinearInput, self).__init__(in_dim, out_dim)
-        self.post = PermuteTF() if permute_tf else nn.Identity()
-
-    def forward(self, x):
-        x = super(GatedLinearInput, self).forward(x)
-        return self.post(x)
-
-
-class AbsLinearOutput(nn.Module):
-    def __init__(self, in_dim, out_dim, permute_tf=True, **kwargs):
-        super(AbsLinearOutput, self).__init__()
-        self.fc = nn.Sequential(
-            PermuteTF() if permute_tf else nn.Identity(),
-            nn.Linear(in_dim, out_dim, **kwargs))
-
-    def forward(self, x):
-        return self.fc(x).abs()
diff --git a/mimikit/freqnet/modules/loss_functions.py b/mimikit/freqnet/modules/loss_functions.py
deleted file mode 100644
index a9e3182..0000000
--- a/mimikit/freqnet/modules/loss_functions.py
+++ /dev/null
@@ -1,6 +0,0 @@
-from torch import nn as nn
-
-
-def mean_L1_prop(output, target):
-    L = nn.L1Loss(reduction="none")(output, target).sum(dim=(0, -1), keepdim=True)
-    return 100 * (L / target.abs().sum(dim=(0, -1), keepdim=True)).mean()
\ No newline at end of file
diff --git a/mimikit/freqnet/modules/padding.py b/mimikit/freqnet/modules/padding.py
deleted file mode 100644
index e247b97..0000000
--- a/mimikit/freqnet/modules/padding.py
+++ /dev/null
@@ -1,21 +0,0 @@
-import torch
-from torch import nn as nn
-
-
-class LearnablePad1d(nn.Module):
-    def __init__(self, input_dim, amount, learn_padding=False):
-        super(LearnablePad1d, self).__init__()
-        self.pad = None
-        self.side = torch.sign(torch.tensor(amount)).item()
-        amount = abs(amount)
-        if amount > 0:
-            if learn_padding:
-                self.pad = nn.Parameter(torch.zeros(input_dim, amount), requires_grad=True)
-            else:
-                self.pad = torch.zeros(input_dim, amount)
-
-    def forward(self, x):
-        if self.pad is None:
-            return x
-        self.pad = self.pad.to(x)
-        return torch.cat((torch.stack([self.pad] * x.size(0)), x)[::self.side], dim=-1)
\ No newline at end of file
diff --git a/mimikit/freqnet/wavenet.py b/mimikit/freqnet/wavenet.py
deleted file mode 100644
index 9ede450..0000000
--- a/mimikit/freqnet/wavenet.py
+++ /dev/null
@@ -1,88 +0,0 @@
-import torch.nn as nn
-import torch.nn.functional as F
-from functools import partial
-from multiprocessing import cpu_count
-
-from . import FreqNet
-from .modules.io import PermuteTF
-from ..data import Database, make_root_db
-from ..connectors.neptune import NeptuneConnector
-from ..data.transforms import file_to_qx
-
-
-def wavenet_db(target,
-               roots=None,
-               files=None,
-               mu=255,
-               sample_rate=22050,
-               neptune_path=None):
-    transform = partial(file_to_qx,
-                        mu=mu,
-                        sr=sample_rate)
-    if roots is None and files is None:
-        roots = "./"
-    make_root_db(target, roots=roots, files=files, extract_func=transform, n_cores=cpu_count() // 2)
-
-    if neptune_path is not None:
-        path = neptune_path.split("/")
-        user, rest = path[0], "/".join(path[1:]) if path[-1] is not None else path[1]
-        connector = NeptuneConnector(user=user, setup={"db": rest})
-        db = Database(target)
-        print("uploading database to neptune...")
-        connector.upload_database("db", db)
-    print("done!")
-    return Database(target)
-
-
-def wavenet_loss_fn(output, target):
-    return F.nll_loss(output.view(-1, output.size(-1)), target.view(-1))
-
-
-class WaveNet(FreqNet):
-
-    def __init__(self,
-                 input_dim=256,
-                 model_dim=256,
-                 groups=1,
-                 n_layers=(1,),
-                 with_skip_conv=True,
-                 with_residual_conv=True,
-                 **data_optim_kwargs
-                 ):
-        super(WaveNet, self).__init__(
-            loss_fn=wavenet_loss_fn,
-            model_dim=model_dim,
-            groups=groups,
-            n_layers=n_layers,
-            strict=False,
-            accum_outputs="right",
-            concat_outputs=None,
-            pad_input="left",
-            learn_padding=False,
-            with_skip_conv=with_skip_conv,
-            with_residual_conv=with_residual_conv,
-            **data_optim_kwargs)
-
-        self.input_dim = input_dim
-
-        self.inpt = nn.Sequential(
-            nn.Embedding(self.input_dim, self.model_dim), PermuteTF())
-
-        self.outpt = nn.Sequential(
-            nn.ReLU(inplace=True),
-            nn.Conv1d(self.model_dim, self.model_dim, 1),
-            nn.ReLU(inplace=True),
-            nn.Conv1d(self.model_dim, self.input_dim, 1),
-            PermuteTF(),
-            nn.LogSoftmax(dim=1)
-        )
-
-        self.save_hyperparameters()
-
-    def forward(self, x):
-        x = self.inpt(x.squeeze())
-        skips = None
-        for layer in self.layers:
-            x, skips = layer(x, skips)
-        x = self.outpt(skips)
-        return x
diff --git a/mimikit/kit/get_trainer.py b/mimikit/get_trainer.py
similarity index 98%
rename from mimikit/kit/get_trainer.py
rename to mimikit/get_trainer.py
index d49b5a8..4080c0e 100644
--- a/mimikit/kit/get_trainer.py
+++ b/mimikit/get_trainer.py
@@ -2,12 +2,15 @@
 import torch
 from pytorch_lightning.loggers import NeptuneLogger
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
-from .loggers import MMKDefaultLogger
-from .checkpoint import MMKCheckpoint
-from .callbacks import EpochProgressBarCallback
+
+from .models.parts import MMKDefaultLogger, MMKCheckpoint, EpochProgressBarCallback
 from pytorch_lightning.trainer import Trainer
 import warnings
 
+__all__ = [
+    'get_trainer'
+]
+
 
 def get_trainer(root_dir=None,
                 version=None,
diff --git a/mimikit/kit/__init__.py b/mimikit/kit/__init__.py
deleted file mode 100644
index 01a34a8..0000000
--- a/mimikit/kit/__init__.py
+++ /dev/null
@@ -1,7 +0,0 @@
-from .callbacks import *
-from .checkpoint import *
-from .loggers import *
-from .hooks import *
-from .get_trainer import get_trainer
-from .ds_wrappers import *
-
diff --git a/mimikit/kit/ds_wrappers.py b/mimikit/kit/ds_wrappers.py
deleted file mode 100644
index 1fa2f73..0000000
--- a/mimikit/kit/ds_wrappers.py
+++ /dev/null
@@ -1,62 +0,0 @@
-from copy import copy
-
-from ..data.data_object import DataObject
-
-
-class DSWrapper(DataObject):
-
-    def __init__(self):
-        pass
-
-    def upgrade(self, dataset):
-        """
-        dynamically extends an object with own methods and attributes.
-        @param dataset: object to be extended
-        @return: extended object will be an instance of object.__class__ and of self.__class__
-        """
-        new = copy(dataset)
-        bases = (self.__class__, new.__class__)
-        name = self.__class__.__name__ + "Dataset"
-        new.__dict__.update(self.__dict__)
-        new.__class__ = type(name, bases, new.__dict__)
-        return new
-
-    def __call__(self, dataset: DataObject):
-        return self.upgrade(dataset)
-
-
-class InputEqualTarget(DSWrapper):
-    def __getitem__(self, item):
-        return tuple(self.data[i] for i in [item, item])
-
-    def __len__(self):
-        return len(self.data)
-
-
-class ShiftedSeqsPair(DSWrapper):
-    def __init__(self, input_length, targets, stride=1):
-        """
-        @param input_length:
-        @param targets: [(shift, length)... ]
-        @param stride:
-        """
-        self.shifts, self.lengths = list(zip(*[(0, input_length), *targets]))
-        self.stride = stride
-        self.N = None
-
-    def __call__(self, dataset: DataObject):
-        # grab the number of time-steps BEFORE we upgrade
-        self.N = len(dataset.data)
-        return super(ShiftedSeqsPair, self).__call__(dataset)
-
-    def __len__(self):
-        ln = (self.N - max(shift + ln for shift, ln in zip(self.shifts, self.lengths)) + 1) // self.stride
-        return ln
-
-    def __getitem__(self, item):
-        i = item * self.stride
-        items_slices = [slice(i + shift, i + shift + length)
-                         for shift, length in zip(self.shifts, self.lengths)]
-        return tuple(self.data[idx] for idx in items_slices)
-
-
diff --git a/mimikit/midi/__init__.py b/mimikit/midi/__init__.py
new file mode 100644
index 0000000..882f1f0
--- /dev/null
+++ b/mimikit/midi/__init__.py
@@ -0,0 +1,4 @@
+from .features import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
diff --git a/mimikit/midi/features.py b/mimikit/midi/features.py
new file mode 100644
index 0000000..be9574b
--- /dev/null
+++ b/mimikit/midi/features.py
@@ -0,0 +1,162 @@
+import torch
+import numpy as np
+import muspy
+import dataclasses as dtc
+from typing import Optional
+from collections import OrderedDict
+import IPython.display as ipd
+
+__all__ = [
+    'MidiMusic',
+    'MidiEvent'
+]
+
+
+class MMKNote(muspy.Note):
+    """
+    muspy.Note with `channel` and `program` attributes
+    """
+    _attributes = OrderedDict(
+        [
+            ("time", int),
+            ("pitch", int),
+            ("duration", int),
+            ("velocity", int),
+            ("pitch_str", str),
+            ("channel", int),
+            ("program", int)
+        ]
+    )
+    _optional_attributes = ["velocity", "pitch_str", "channel", "program"]
+
+    def __init__(
+            self,
+            time: int,
+            pitch: int,
+            duration: int,
+            velocity: Optional[int] = None,
+            pitch_str: Optional[str] = None,
+            channel: Optional[int] = None,
+            program: Optional[int] = None
+    ):
+        super(MMKNote, self).__init__(time, pitch, duration, velocity, pitch_str)
+        self.channel = channel
+        self.program = program
+
+
+class MMKMusic(muspy.Music):
+
+    def __init__(self,
+                 metadata=None,
+                 resolution=None,
+                 tempos=None,
+                 key_signatures=None,
+                 time_signatures=None,
+                 downbeats=None,
+                 lyrics=None,
+                 annotations=None,
+                 tracks=None
+                 ):
+        super(MMKMusic, self).__init__(metadata, resolution, tempos, key_signatures,
+                                       time_signatures, downbeats, lyrics,
+                                       annotations, tracks)
+        # add channel and program attributes to the notes
+        for t in self.tracks:
+            t.notes = [MMKNote(**n.to_ordered_dict(), channel=t.name, program=t.program)
+                       for n in t.notes]
+
+
+@dtc.dataclass
+class MidiMusic:
+    __ext__ = 'midi'
+
+    sr: int = 44100
+    backend: str = 'mido'
+    duplicate_note_mode: str = 'fifo'
+
+    @property
+    def params(self):
+        return dtc.asdict(self)
+
+    def load(self, path):
+        music = muspy.read_midi(path, self.backend, duplicate_note_mode=self.duplicate_note_mode)
+        return MMKMusic.from_dict(music.to_ordered_dict())
+
+    def display(self, inputs, **kwargs):
+        muspy.show_pianoroll(inputs, **kwargs)
+
+    def synthesize(self, inputs):
+        return inputs.synthesize(rate=self.sr).sum(axis=1).reshape(-1)
+
+    def play(self, inputs):
+        y = self.synthesize(inputs)
+        ipd.display(ipd.Audio(y, rate=self.sr))
+
+
+@dtc.dataclass
+class MidiEvent(MidiMusic):
+
+    resolution: int = muspy.DEFAULT_RESOLUTION
+    program: int = 0
+    is_drum: bool = False
+    use_single_note_off_event: bool = False
+    use_end_of_sequence_event: bool = False
+    encode_velocity: bool = False
+    force_velocity_event: bool = True
+    max_time_shift: int = 100
+    velocity_bins: int = 32
+    default_velocity: int = 64
+
+    @property
+    def encoders(self):
+        return {
+            muspy.Music:
+                lambda music: muspy.to_event_representation(music,
+                                                            self.use_single_note_off_event,
+                                                            self.use_end_of_sequence_event,
+                                                            self.encode_velocity,
+                                                            self.force_velocity_event,
+                                                            self.max_time_shift,
+                                                            self.velocity_bins)
+        }
+
+    @property
+    def decoders(self):
+        return {
+            np.ndarray:
+                lambda arr: muspy.from_event_representation(arr,
+                                                            self.resolution,
+                                                            self.program,
+                                                            self.is_drum,
+                                                            self.use_single_note_off_event,
+                                                            self.use_end_of_sequence_event,
+                                                            self.max_time_shift,
+                                                            self.velocity_bins,
+                                                            self.default_velocity,
+                                                            self.duplicate_note_mode
+                                                            ),
+            torch.Tensor:
+                lambda tensor: self.decode(tensor.detach().cpu().numpy())
+        }
+
+    def encode(self, inputs):
+        return self.encoders[type(inputs)](inputs)
+
+    def decode(self, inputs):
+        return self.decoders[type(inputs)](inputs)
+
+    def load(self, path):
+        music = super(MidiEvent, self).load(path)
+        return self.encode(music)
+
+    def display(self, inputs, **kwargs):
+        inputs = self.decode(inputs)
+        super(MidiEvent, self).display(inputs)
+
+    def synthesize(self, inputs):
+        inputs = self.decode(inputs)
+        return super(MidiEvent, self).synthesize(inputs)
+
+    def play(self, inputs):
+        inputs = self.decode(inputs)
+        return super(MidiEvent, self).play(inputs)
diff --git a/mimikit/models/__init__.py b/mimikit/models/__init__.py
new file mode 100644
index 0000000..441285c
--- /dev/null
+++ b/mimikit/models/__init__.py
@@ -0,0 +1,10 @@
+from .freqnet import *
+from .model import *
+from .parts import *
+from .sample_rnn import *
+from .seq2seqlstm import *
+from .wavenet import *
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
+
+
diff --git a/mimikit/models/freqnet.py b/mimikit/models/freqnet.py
new file mode 100644
index 0000000..743457f
--- /dev/null
+++ b/mimikit/models/freqnet.py
@@ -0,0 +1,160 @@
+import torch
+import dataclasses as dtc
+
+from ..abstract.features import Feature, SegmentLabels, FilesLabels
+from ..audios import Spectrogram
+from ..data import Input, AsSlice, Target
+from .parts import SuperAdam, SequenceModel, mean_L1_prop, IData
+from .model import model
+from ..networks import FreqNetNetwork, WNNetwork
+
+__all__ = [
+    'FreqNetData',
+    'FreqNet',
+    'main'
+]
+
+
+@dtc.dataclass
+class FreqNetData(IData):
+
+    feature: Feature = None
+    batch_size: int = 64
+    batch_seq_length: int = 64
+
+    @classmethod
+    def schema(cls, sr=22050, emphasis=0., n_fft=2048, hop_length=512,
+               segment_labels=False, files_labels=False):
+
+        schema = {"fft": Spectrogram(sr=sr, emphasis=emphasis,
+                                     n_fft=n_fft, hop_length=hop_length,
+                                     magspec=True)}
+        if segment_labels:
+            schema.update({
+                'loc': SegmentLabels(input_key='fft')
+            })
+        if files_labels:
+            schema.update({
+                'glob': FilesLabels(input_key='fft')
+            })
+        return schema
+
+    @classmethod
+    def dependant_hp(cls, db):
+        hp = {}
+        if 'loc' in db.features:
+            hp.update(dict(n_cin_classes=len(db.loc.regions)))
+        if 'glob' in db.features:
+            hp.update(dict(n_gin_classes=len(db.glob.files)))
+        return dict(
+            feature=db.schema['fft'], input_dim=db.schema['fft'].dim, **hp
+        )
+
+    def batch_signature(self, stage='fit'):
+        inpt = Input('fft', AsSlice(shift=0, length=self.batch_seq_length))
+        trgt = Target('fft', AsSlice(shift=self.shift,
+                                     length=self.output_shape((-1, self.batch_seq_length, -1))[1]))
+        # where are we conditioned?
+        loc, glob = self.n_cin_classes is not None, self.n_gin_classes is not None
+        if loc:
+            pass
+        if glob:
+            pass
+        if stage in ('full', 'fit', 'train', 'val'):
+            return inpt, trgt
+        # test, predict, generate...
+        return inpt
+
+    def loader_kwargs(self, stage, datamodule):
+        return dict(
+            batch_size=self.batch_size,
+            drop_last=False,
+            shuffle=True
+        )
+
+
+@model
+@dtc.dataclass
+class FreqNet(
+    # data configuration :
+    FreqNetData,
+    # optimizer :
+    SuperAdam,
+    # training step, generate routine & interface :
+    SequenceModel,
+    # overrides wavenet's inputs/outputs modules :
+    FreqNetNetwork,
+    # we inherit the networks __init__ and all generate_* from wavenet
+    WNNetwork
+):
+
+    @staticmethod
+    def loss_fn(output, target):
+        return {"loss": mean_L1_prop(output, target)}
+
+    def setup(self, stage: str):
+        SequenceModel.setup(self, stage)
+        SuperAdam.setup(self, stage)
+
+    def encode_inputs(self, inputs: torch.Tensor):
+        return self.feature.encode(inputs)
+
+    def decode_outputs(self, outputs: torch.Tensor):
+        return self.feature.decode(outputs)
+
+
+def main(
+        sources='./data',
+        sr=22050, n_fft=2048, hop_length=512,
+        segments_labels=False, files_labels=False):
+    import mimikit as mmk
+    import os
+
+    schema = mmk.FreqNet.schema(sr, emphasis=0., n_fft=n_fft, hop_length=hop_length,
+                                segment_labels=False, files_labels=False)
+
+    db_path = '/tmp/freqnet_db.h5'
+    db = mmk.Database.create(db_path, sources, schema)
+
+    net = mmk.FreqNet(
+        **mmk.FreqNet.dependant_hp(db),
+        cin_dim=None,
+        gin_dim=None,
+        n_layers=(4,),
+        gate_dim=1024,
+        groups=2,
+        batch_size=16,
+        batch_seq_length=32,
+        max_lr=1e-3,
+        div_factor=10,
+        betas=(.9, .92),
+
+    )
+    print(net.hparams)
+
+    dm = mmk.DataModule(net, db,
+                        splits=tuple())
+
+    cb = mmk.GenerateCallBack(3, indices=[None]*4,
+                              n_steps=1000,
+                              play_audios=False,
+                              plot_audios=False,
+                              log_audios=True,
+                              log_dir=os.path.abspath('outputs/freqnet'))
+
+    trainer = mmk.get_trainer(root_dir=None,
+                              max_epochs=100,
+                              callbacks=[cb],
+                              checkpoint_callback=False)
+
+    trainer.fit(net, datamodule=dm)
+
+    prp = dm.get_prompts([None] * 4)
+    print(prp.size())
+
+    out = net.generate(prp, 32)
+    print('Done!')
+
+
+if __name__ == '__main__':
+    main()
diff --git a/mimikit/models/model.py b/mimikit/models/model.py
new file mode 100644
index 0000000..dd1e1eb
--- /dev/null
+++ b/mimikit/models/model.py
@@ -0,0 +1,77 @@
+import dataclasses as dtc
+from inspect import signature
+
+import pytorch_lightning as pl
+
+
+def _filter_cls_kwargs(cls, kwargs):
+    valids = signature(cls.__init__).parameters.keys()
+    return {k: v for k, v in kwargs.items() if k in valids}
+
+
+def model(cls):
+    """
+
+    Parameters
+    ----------
+    cls
+
+    Returns
+    -------
+
+    """
+
+    bases = cls.__bases__
+    bases = [cls, *bases]
+
+    # get the __init__ signatures of the bases
+    sigs = [dict(signature(d.__init__).parameters) for d in bases
+            if d is not pl.LightningModule]
+    # collect positional arguments for validation
+    posits = {k: v for s in sigs for k, v in s.items() if v.default is v.empty}
+    if any(k for k in posits.keys() if k != "self"):
+        for b in bases:
+            posits = {k: v for k, v in dict(signature(b.__init__).parameters).items() if v.default is v.empty}
+            if any(k for k in posits.keys() if k != "self"):
+                raise ValueError("All dependencies should have kwargs only."
+                                 " Got following positional args: " + str([k for k in posits.keys() if k != "self"]) +\
+                                 "for the following parent class: " + str(b))
+
+    # collect keywords arguments
+    kws = {k: v for s in sigs for k, v in s.items() if v.default is not v.empty}
+
+    # build fields for make_dataclass
+    fields = [(k, type(v.default), dtc.field(default=v.default)) for k, v in kws.items()]
+
+    # method to initialize all the bases
+    def __post_init__(self):
+        super(pl.LightningModule, self).__init__()
+        for b in bases:
+            if hasattr(b, '__post_init__'):
+                # avoid recursion if b is a dataclass already
+                b.__post_init__(self)
+            else:
+                b.__init__(self, **_filter_cls_kwargs(b, self.__dict__))
+
+        # set hparams manually since lightning will not be able to infer them...
+        self._set_hparams({k: getattr(self, k) for k in self.__args__})
+
+    # namespace of the class
+    ns = dict(cls.__dict__)
+    ns.update({
+        "__post_init__": __post_init__,
+        '__args__': tuple(kws.keys()),
+    })
+    return dtc.make_dataclass(cls.__qualname__,
+                              fields,
+                              bases=(*(b for b in bases if b is not cls),
+                                     *((pl.LightningModule,) if pl.LightningModule not in bases else tuple())),
+                              namespace=ns,
+                              init=True, repr=False, eq=False, frozen=False, unsafe_hash=True
+                              )
+
+
+class Model(type):
+    def __new__(mcs, name, bases, namespace):
+        cls = super(Model, mcs).__new__(mcs, name, bases, namespace)
+        return model(cls)
\ No newline at end of file
diff --git a/mimikit/models/parts/__init__.py b/mimikit/models/parts/__init__.py
new file mode 100644
index 0000000..7fa4950
--- /dev/null
+++ b/mimikit/models/parts/__init__.py
@@ -0,0 +1,11 @@
+from .callbacks import *
+from .checkpoint import *
+from .idata import *
+from .hooks import *
+from .loggers import *
+from .loss_functions import *
+from .optim import *
+from .sequence_model import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/kit/callbacks.py b/mimikit/models/parts/callbacks.py
similarity index 56%
rename from mimikit/kit/callbacks.py
rename to mimikit/models/parts/callbacks.py
index 2795075..4e0b9d5 100644
--- a/mimikit/kit/callbacks.py
+++ b/mimikit/models/parts/callbacks.py
@@ -1,6 +1,13 @@
 from pytorch_lightning.callbacks.progress import Callback
 from IPython import get_ipython
 
+__all__ = [
+    'is_notebook',
+    'EpochProgressBarCallback',
+    'GradNormCallback',
+    'tqdm'
+]
+
 
 def is_notebook():
     shell = get_ipython().__class__.__name__
@@ -24,6 +31,20 @@ def on_train_start(self, trainer, pl_module):
         self.epoch_bar = tqdm(range(1, trainer.max_epochs), unit="epoch",
                               position=0, leave=False, dynamic_ncols=True)
 
-    def on_epoch_end(self, trainer, pl_module):
+    def on_train_epoch_end(self, trainer, *args):
+        if trainer.val_dataloaders is not None and any(trainer.val_dataloaders):
+            return
+        else:
+            self.epoch_bar.update()
+
+    def on_validation_epoch_end(self, *args):
         self.epoch_bar.update()
 
+
+class GradNormCallback(Callback):
+
+    def __init__(self):
+        self.gradnorms = []
+
+    def on_after_backward(self, trainer, pl_module: 'LightningModule') -> None:
+        self.gradnorms += [pl_module.grad_norm(1.)]
\ No newline at end of file
diff --git a/mimikit/kit/checkpoint.py b/mimikit/models/parts/checkpoint.py
similarity index 80%
rename from mimikit/kit/checkpoint.py
rename to mimikit/models/parts/checkpoint.py
index 96647d0..256df6a 100644
--- a/mimikit/kit/checkpoint.py
+++ b/mimikit/models/parts/checkpoint.py
@@ -3,6 +3,10 @@
 from typing import Iterable
 import os
 
+__all__ = [
+    'MMKCheckpoint'
+]
+
 
 class MMKCheckpoint(ModelCheckpoint):
     def __init__(self,
@@ -13,22 +17,20 @@ def __init__(self,
         self.epochs = epochs
 
         super(MMKCheckpoint, self).__init__(
-            filepath=None,
             monitor=None,
             verbose=False,
             save_last=None,
             save_top_k=None,
-            save_weights_only=False, # we just save optimizers separately
-            mode="auto",
+            save_weights_only=False,  # we just save optimizers separately
+            mode="min",
             period=1,
-            prefix="",
             dirpath=os.path.join(dirpath, "states"),
             filename=filename
         )
 
     def should_save(self, epoch, step):
         if type(self.epochs) is int:
-            return (epoch % self.epochs) == 0
+            return epoch > 0 and (epoch % self.epochs) == 0
         elif type(self.epochs) is float:  # Doesn't work for now...
             return (step % int(self.epochs)) == 0
         elif isinstance(self.epochs, Iterable):
@@ -36,9 +38,9 @@ def should_save(self, epoch, step):
         return False
 
     def save_checkpoint(self, trainer, pl_module):
-        epoch, global_step = trainer.current_epoch, trainer.global_step
+        epoch, global_step = trainer.current_epoch + 1, trainer.global_step
         if trainer.state == TrainerState.INTERRUPTED or \
-                epoch == (trainer.max_epochs - 1) or self.should_save(epoch, global_step):
+                epoch == trainer.max_epochs or self.should_save(epoch, global_step):
             # hacks the code from github
             self._add_backward_monitor_support(trainer)
             self._validate_monitor_key(trainer)
@@ -50,7 +52,7 @@ def save_checkpoint(self, trainer, pl_module):
 
             self._save_model(filepath, trainer, pl_module)
 
-    def on_save_checkpoint(self, trainer, model):
+    def on_save_checkpoint(self, trainer, model, checkpoint):
         """returns the state to be saved in the checkpoint"""
         return {"epochs": self.epochs}
 
diff --git a/mimikit/kit/hooks.py b/mimikit/models/parts/hooks.py
similarity index 85%
rename from mimikit/kit/hooks.py
rename to mimikit/models/parts/hooks.py
index 0ff7d1a..d34f1aa 100644
--- a/mimikit/kit/hooks.py
+++ b/mimikit/models/parts/hooks.py
@@ -1,5 +1,6 @@
 from pytorch_lightning.utilities.cloud_io import load as pl_load
 from pytorch_lightning.callbacks.model_checkpoint import ModelCheckpoint
+from pytorch_lightning import LightningModule
 from time import time, gmtime
 from typing import Optional, Union, Callable, Dict, IO
 import torch
@@ -15,11 +16,15 @@
     warnings.warn("It seems neptune and/or neptunecontrib are not installed. "
                   "You won't be able to log checkpoints and data to neptune.")
 
-from .. import __version__ as version
-from ..data.transforms import SR
+from ... import __version__ as version
 
+__all__ = [
+    'LoggingHooks',
+    'MMKHooks'
+]
 
-class LoggingHooks:
+
+class LoggingHooks(LightningModule):
 
     @property
     def neptune_experiment(self):
@@ -51,7 +56,7 @@ def training_step_end(self, out):
     def validation_step_end(self, out):
         return self.log_output(out)
 
-    def on_train_epoch_end(self, *args):
+    def _flush_ep_metrics(self):
         to_print = "Epoch %i " % self.current_epoch
         to_log = {}
         for k, v in self._ep_metrics.items():
@@ -61,6 +66,16 @@ def on_train_epoch_end(self, *args):
         if getattr(self, "logger", None) is not None:
             self.logger.log_metrics(to_log)
 
+    def on_train_epoch_end(self, *args):
+        if any(self.trainer.val_dataloaders):
+            # wait until validation end
+            return
+        else:
+            self._flush_ep_metrics()
+
+    def on_validation_epoch_end(self, *args):
+        self._flush_ep_metrics()
+
     def on_fit_start(self):
         self._ep_time = time()
 
@@ -76,7 +91,7 @@ def on_fit_end(self):
             self.hparams.update({"training_time_sec": duration})
             self.logger.log_hyperparams(self.hparams)
 
-    def log_audio(self, filename, audio_tensor, sample_rate=SR, experiments=None):
+    def log_audio(self, filename, audio_tensor, sample_rate=22050, experiments=None):
         # TensorBoards write their own "Event" file which is quite cumbersome to extract afterwards...
         # NeptuneLoggers need an audio file written on disk and store it afterwards as html on the server...
         # Just to be sure and even if it is then in 3 places : we add the audio in root_dir/audios of the model!
@@ -125,7 +140,7 @@ def _check_version(other_v):
                       (", try to install the right version with `pip install mimikit==%s" % v))
 
 
-class MMKHooks:
+class MMKHooks(LightningModule):
     _loaded_checkpoint = None
 
     @classmethod
@@ -190,3 +205,25 @@ def on_load_checkpoint(self, checkpoint):
                 checkpoint["lr_schedulers"] = optim_state["lr_schedulers"]
 
         return checkpoint
+
+    @staticmethod
+    def load_checkpoint_hparams(checkpoint_path):
+        ckpt = pl_load(checkpoint_path)
+        return ckpt['hyper_parameters']
+
+    @staticmethod
+    def list_available_states(root_dir):
+        return [os.path.abspath(os.path.join(root_dir, "states", f))
+                for f in os.listdir(os.path.join(root_dir, "states")) if "epoch=" in f]
+
+    @staticmethod
+    def list_available_epochs(root_dir):
+        epochs = []
+        for f in os.listdir(os.path.join(root_dir, "states")):
+            if "epoch=" in f:
+                epochs += [int(f[6:][:-5])]
+        return epochs
+
+    @classmethod
+    def load_epoch_checkpoint(cls, root_dir, epoch, **kwargs):
+        return cls.load_from_checkpoint(os.path.join(root_dir, "states", "epoch=%i.ckpt" % epoch), **kwargs)
diff --git a/mimikit/models/parts/idata.py b/mimikit/models/parts/idata.py
new file mode 100644
index 0000000..0b5e738
--- /dev/null
+++ b/mimikit/models/parts/idata.py
@@ -0,0 +1,35 @@
+from pytorch_lightning import LightningModule
+from typing import Iterable
+import dataclasses as dtc
+from torch.utils.data import Dataset
+
+
+from ...data import Database, DataModule
+
+__all__ = [
+    'IData'
+]
+
+
+class IData(LightningModule):
+
+    def __init__(self):
+        LightningModule.__init__(self)
+
+    @property
+    def dm(self) -> DataModule:
+        return self.trainer.data_connector.datamodule
+
+    @classmethod
+    def schema(cls, *args, **kwargs) -> dict:
+        raise NotImplementedError
+
+    @classmethod
+    def dependant_hp(cls, db: Database) -> dict:
+        return {}
+
+    def batch_signature(self, stage):
+        raise NotImplementedError
+
+    def loader_kwargs(self, stage, datamodule) -> dict:
+        raise NotImplementedError
diff --git a/mimikit/kit/loggers.py b/mimikit/models/parts/loggers.py
similarity index 98%
rename from mimikit/kit/loggers.py
rename to mimikit/models/parts/loggers.py
index ca64e4e..24eec7b 100644
--- a/mimikit/kit/loggers.py
+++ b/mimikit/models/parts/loggers.py
@@ -5,6 +5,10 @@
 from typing import Optional
 import os
 
+__all__ = [
+    'MMKDefaultLogger'
+]
+
 
 class MMKExperiment(Experiment):
     """
diff --git a/mimikit/models/parts/loss_functions.py b/mimikit/models/parts/loss_functions.py
new file mode 100644
index 0000000..cf44e39
--- /dev/null
+++ b/mimikit/models/parts/loss_functions.py
@@ -0,0 +1,64 @@
+import torch
+import torch.nn as nn
+
+__all__ = [
+    'angular_distance',
+    'cosine_similarity',
+    'mean_L1_prop'
+]
+
+
+def mean_L1_prop(output, target):
+    L = nn.L1Loss(reduction="none")(output, target).sum(dim=(0, -1), keepdim=True)
+    return 100 * (L / target.abs().sum(dim=(0, -1), keepdim=True)).mean()
+
+
+def cosine_similarity(X, Y, eps=1e-10):
+    """
+    safely computes the cosine similarity between matrices X and Y.
+
+    Shapes:
+    -------
+    X : (*, N, D)
+    Y : (*, M, D)
+    D_xy : (*, N, M)
+
+    Notes:
+    ------
+    The need for this function arises from the fact that torch.nn.CosineSimilarity only computes the
+    diagonal of D_xy, as in cosine_sim(output, target)
+    """
+    if not isinstance(eps, torch.Tensor):
+        eps = torch.tensor(eps).to(X)
+
+    dot_prod = torch.matmul(X, Y.transpose(-2, -1))
+    norms = torch.norm(X, p=2, dim=-1).unsqueeze(-1) * torch.norm(Y, p=2, dim=-1).unsqueeze(-2)
+    cos_theta = dot_prod / torch.maximum(norms, eps)
+    return cos_theta
+
+
+def angular_distance(X, Y, eps=1e-20):
+    """
+    angular distance is a valid distance metric based on the cosine similarity
+    see https://en.wikipedia.org/wiki/Cosine_similarity#Angular_distance_and_similarity
+
+    Shapes:
+    -------
+    X : (*, N, D)
+    Y : (*, M, D)
+    D_xy : (*, N, M)
+    """
+    if not isinstance(eps, torch.Tensor):
+        eps = torch.tensor(eps).to(X)
+
+    def safe_acos(x):
+        # torch.acos returns nan near -1 and 1... see https://github.com/pytorch/pytorch/issues/8069
+        return torch.acos(torch.clamp(x, min=-1 + eps / 2, max=1 - eps / 2))
+
+    have_negatives = torch.any(X < 0) or torch.any(Y < 0)
+    cos_theta = cosine_similarity(X, Y, eps)
+
+    pi = torch.acos(torch.zeros(1)).item() * 2
+    D_xy = (1 + int(not have_negatives)) * safe_acos(cos_theta) / pi
+
+    return D_xy
\ No newline at end of file
diff --git a/mimikit/models/parts/optim.py b/mimikit/models/parts/optim.py
new file mode 100644
index 0000000..66979a5
--- /dev/null
+++ b/mimikit/models/parts/optim.py
@@ -0,0 +1,87 @@
+import torch
+from pytorch_lightning import LightningModule
+import dataclasses as dtc
+from typing import Tuple, Optional
+
+__all__ = [
+    'Adam',
+    'SuperAdam'
+]
+
+
+class Adam(LightningModule):
+    def __init__(self,
+                 max_lr=1e-3,
+                 betas=(.9, .9)
+                 ):
+        super(LightningModule, self).__init__()
+        self.max_lr = max_lr
+        self.betas = betas
+
+    def configure_optimizers(self):
+        self.opt = torch.optim.Adam(self.parameters(), lr=self.max_lr, betas=self.betas)
+        return self.opt
+
+
+class ManyOneCycleLR(torch.optim.lr_scheduler.OneCycleLR):
+    def step(self, epoch=None):
+        try:
+            super(ManyOneCycleLR, self).step(epoch=epoch)
+        except ValueError:
+            self.last_epoch = 0
+            self._step_count = 1
+            super(ManyOneCycleLR, self).step(epoch=epoch)
+
+
+@dtc.dataclass(repr=False, eq=False)
+class SuperAdam(LightningModule):
+
+    max_lr: float = 1e-3
+    betas: Tuple[float, float] = (.9, .9)
+    div_factor: float = 3.
+    final_div_factor: float = 1.
+    pct_start: float = 0.
+    cycle_momentum: bool = False
+    total_steps: Optional[int] = None
+
+    def __post_init__(self):
+        super(LightningModule, self).__init__()
+        # those are set in setup when we know what the trainer and datamodule got
+        self.steps_per_epoch = None
+        self.max_epochs = None
+        self.opt, self.sched = None, None
+
+    def configure_optimizers(self):
+        self.opt = torch.optim.Adam(self.parameters(), lr=self.max_lr, betas=self.betas)
+        steps = dict(total_steps=self.total_steps) if self.total_steps is not None \
+            else dict(steps_per_epoch=self.steps_per_epoch, epochs=self.max_epochs)
+        self.sched = ManyOneCycleLR(self.opt,
+                                    **steps,
+                                    max_lr=self.max_lr, div_factor=self.div_factor,
+                                    final_div_factor=self.final_div_factor,
+                                    pct_start=self.pct_start,
+                                    cycle_momentum=self.cycle_momentum
+                                    )
+        return [self.opt], [{"scheduler": self.sched, "interval": "step", "frequency": 1}]
+
+    def setup(self, stage: str):
+        if stage == "fit":
+            self.max_epochs = self.trainer.max_epochs
+            self.steps_per_epoch = len(self.train_dataloader())
+
+
+class RMS(LightningModule):
+    def __init__(self,
+                 lr=1e-2,
+                 alpha=0.99,
+                 eps=1e-8,
+                 weight_decay=0,
+                 momentum=0,
+                 centered=False
+                 ):
+        super(LightningModule, self).__init__()
+        self.args = dict(lr=lr, alpha=alpha, eps=eps, weight_decay=weight_decay, momentum=momentum, centered=centered)
+
+    def configure_optimizers(self):
+        self.opt = torch.optim.RMSprop(self.parameters(), **self.args)
+        return self.opt
\ No newline at end of file
diff --git a/mimikit/models/parts/sequence_model.py b/mimikit/models/parts/sequence_model.py
new file mode 100644
index 0000000..a19e425
--- /dev/null
+++ b/mimikit/models/parts/sequence_model.py
@@ -0,0 +1,121 @@
+import torch
+import pytorch_lightning as pl
+from abc import ABC
+import matplotlib.pyplot as plt
+import os
+
+from .hooks import MMKHooks, LoggingHooks
+from ...utils import audio
+
+__all__ = [
+    'SequenceModel',
+    'GenerateCallBack'
+]
+
+
+class SequenceModel(MMKHooks,
+                    LoggingHooks,
+                    pl.LightningModule,
+                    ABC):
+    loss_fn = None
+
+    def __init__(self):
+        super(pl.LightningModule, self).__init__()
+        MMKHooks.__init__(self)
+        LoggingHooks.__init__(self)
+
+    def training_step(self, batch, batch_idx):
+        batch, target = batch
+        output = self.forward(batch)
+        return self.loss_fn(output, target)
+
+    def validation_step(self, batch, batch_idx):
+        batch, target = batch
+        output = self.forward(batch)
+        L = self.loss_fn(output, target)
+        return {"val_loss" if k == "loss" else k: v for k, v in L.items()}
+
+    def setup(self, stage: str):
+        if stage == "fit" and getattr(self, "logger", None) is not None:
+            self.logger.log_hyperparams(self.hparams)
+
+    def encode_inputs(self, inputs: torch.Tensor):
+        raise NotImplementedError
+
+    def decode_outputs(self, outputs: torch.Tensor):
+        raise NotImplementedError
+
+    def before_generate(self, *args, **kwargs):
+        # prepare model
+        self._was_training = self.training
+        self._initial_device = self.device
+        self.eval()
+        self.to("cuda" if torch.cuda.is_available() else "cpu")
+        torch.set_grad_enabled(False)
+
+    def after_generate(self, *args, **kwargs):
+        # reset model
+        self.to(self._initial_device)
+        self.train() if self._was_training else None
+        torch.set_grad_enabled(True)
+
+    def generate(self, prompt, n_steps=16000,
+                 encode_prompt=False, decode_outputs=False,
+                 **kwargs):
+        # prepare model
+        self.before_generate()
+        if encode_prompt:
+            prompt = self.encode_inputs(prompt)
+        output = self.generate_(prompt, n_steps, **kwargs)
+
+        if decode_outputs:
+            output = self.decode_outputs(output)
+
+        self.after_generate()
+        return output
+
+
+class GenerateCallBack(pl.callbacks.Callback):
+
+    def __init__(self, every_n_epochs=10, indices=3, n_steps=1000,
+                 plot_audios=True, play_audios=True, log_audios=False,
+                 log_dir=None,
+                 **gen_kwargs):
+        self.every_n_epochs = every_n_epochs
+        self.indices = indices
+        self.n_steps = n_steps
+        self.kwargs = gen_kwargs
+        self.log_audios = log_audios
+        self.plot_audios = plot_audios
+        self.play_audios = play_audios
+        self.log_dir = log_dir
+
+    def on_train_epoch_start(self, trainer, model):
+        # TODO : avoid having to do that!
+        dm = trainer.datamodule
+        dm.full_ds = dm._init_ds(dm.full_ds, 'fit')
+
+    def on_epoch_end(self, trainer: pl.Trainer, model: SequenceModel):
+        if (trainer.current_epoch + 1) % self.every_n_epochs != 0:
+            return
+        dm = trainer.datamodule
+        prompt = dm.get_prompts(self.indices)
+        output = model.generate(prompt, self.n_steps, decode_outputs=True, **self.kwargs)
+        sr = model.feature.params.get('sr', 22050)
+        hop_length = model.feature.params.get('hop_length', 512)
+        for i in range(output.size(0)):
+            y = output[i].detach().cpu().numpy()
+            if self.plot_audios:
+                plt.figure(figsize=(20, 2))
+                plt.plot(y)
+                plt.show()
+            if self.play_audios:
+                audio(y, sr=sr, hop_length=hop_length)
+
+        if self.log_audios:
+            for i in range(output.size(0)):
+                filename = "epoch=%i - prompt=%i" % (trainer.current_epoch, i)
+                if self.log_dir is not None:
+                    os.makedirs(self.log_dir, exist_ok=True)
+                    filename = os.path.join(self.log_dir, filename)
+                model.log_audio(filename, output[i].unsqueeze(0), sample_rate=sr)
diff --git a/mimikit/models/sample_rnn.py b/mimikit/models/sample_rnn.py
new file mode 100644
index 0000000..9e18d1a
--- /dev/null
+++ b/mimikit/models/sample_rnn.py
@@ -0,0 +1,146 @@
+import inspect
+import click
+import torch
+import torch.nn as nn
+import dataclasses as dtc
+
+from ..audios.features import MuLawSignal
+from ..data import TBPTTSampler, Feature, AsSlice, AsFramedSlice, Input, Target
+from .parts import SuperAdam, SequenceModel
+from ..networks.sample_rnn import SampleRNNNetwork
+
+from . import model, IData
+
+__all__ = [
+    'SampleRNNData',
+    'SampleRNN',
+    'main'
+]
+
+
+@dtc.dataclass
+class SampleRNNData(IData):
+    feature: Feature = None
+    batch_size: int = 16
+    chunk_len: int = 16000 * 8
+    batch_seq_len: int = 512
+
+    @classmethod
+    def schema(cls, sr=22050, emphasis=0., q_levels=256):
+        return {'qx': MuLawSignal(sr=sr, emphasis=emphasis, q_levels=q_levels)}
+
+    def batch_signature(self, stage='fit'):
+        batch_seq_len, frame_sizes = tuple(getattr(self, key) for key in ["batch_seq_len", "frame_sizes"])
+        shifts = [frame_sizes[0] - size for size in frame_sizes]
+        inputs = []
+        for fs, shift in zip(frame_sizes[:-1], shifts[:-1]):
+            inputs.append(
+                Input('qx', AsFramedSlice(shift, batch_seq_len, frame_size=fs,
+                                          as_strided=False)))
+        inputs.append(
+            Input('qx', AsFramedSlice(shifts[-1], batch_seq_len, frame_size=frame_sizes[-1],
+                                      as_strided=True)))
+        targets = Target('qx', AsSlice(shift=frame_sizes[0], length=batch_seq_len))
+        if stage in ('fit', 'train', 'val'):
+            return inputs, targets
+        # test, predict, generate
+        return Input('qx', AsSlice(0, batch_seq_len))
+
+    @classmethod
+    def dependant_hp(cls, db):
+        return dict(feature=db.schema['qx'], q_levels=db.schema['qx'].q_levels)
+
+    def loader_kwargs(self, stage, datamodule):
+        ds = datamodule.datasets[stage]
+        N = len(ds)
+        batch_sampler = TBPTTSampler(N,
+                                     self.batch_size,
+                                     self.chunk_len,
+                                     self.batch_seq_len)
+        return dict(batch_sampler=batch_sampler)
+
+
+@model
+class SampleRNN(
+    # db schema, batch sig
+    SampleRNNData,
+    # training step, generate routine & interface
+    SequenceModel,
+    # optimizer & scheduler
+    SuperAdam,
+    # nn.Module, generate_ function
+    SampleRNNNetwork
+):
+
+    def setup(self, stage: str):
+        SequenceModel.setup(self, stage)
+        SuperAdam.setup(self, stage)
+
+    @staticmethod
+    def loss_fn(output, target):
+        criterion = nn.CrossEntropyLoss(reduction="mean")
+        return {"loss": criterion(output.view(-1, output.size(-1)), target.view(-1))}
+
+    def on_train_batch_start(self, batch, batch_idx, dataloader_idx):
+        if (batch_idx * self.batch_seq_len) % self.chunk_len == 0:
+            self.reset_h0()
+
+    def encode_inputs(self, inputs: torch.Tensor):
+        return self.feature.encode(inputs)
+
+    def decode_outputs(self, outputs: torch.Tensor):
+        return self.feature.decode(outputs)
+
+
+def make_click_command(f):
+    aspec = inspect.getfullargspec(f)
+    if aspec.defaults is not None:
+        for arg, default in zip(reversed(aspec.args), reversed(aspec.defaults)):
+            opt = click.option('--' + arg.replace("_", "-"), default=default)
+            f = opt(f)
+    return click.command()(f)
+
+
+@make_click_command
+def main(sources=['./data'], sr=16000, q_levels=256):
+    import mimikit as mmk
+    import os
+
+    schema = mmk.SampleRNN.schema(sr, 0., q_levels)
+
+    net = mmk.SampleRNN(
+        feature=schema['qx'],
+        q_levels=q_levels,
+        frame_sizes=(16, 4, 4),
+        n_rnn=2,
+
+        max_lr=7e-4,
+        betas=(.9, .91),
+        div_factor=5,
+    )
+
+    print(net.hparams)
+
+    dm = mmk.DataModule(net, "/tmp/srnn.h5",
+                        sources=sources, schema=schema,
+                        splits=tuple(),
+                        in_mem_data=True)
+
+    cb = mmk.GenerateCallBack(5, indices=[None] * 4,
+                              n_steps=16000*10,
+                              play_audios=False,
+                              plot_audios=False,
+                              log_audios=True,
+                              log_dir=os.path.abspath('outputs/sample-rnn'),
+                              temperature=torch.tensor([[.9], [.999], [1.1], [1.25]]).to('cuda'))
+
+    trainer = mmk.get_trainer(root_dir=None,
+                              max_epochs=50,
+                              callbacks=[cb],
+                              checkpoint_callback=False)
+
+    trainer.fit(net, datamodule=dm)
+
+
+if __name__ == '__main__':
+    main()
diff --git a/mimikit/models/seq2seqlstm.py b/mimikit/models/seq2seqlstm.py
new file mode 100644
index 0000000..06a5a8b
--- /dev/null
+++ b/mimikit/models/seq2seqlstm.py
@@ -0,0 +1,68 @@
+import numpy as np
+import pytorch_lightning as pl
+import torch
+import torch.nn as nn
+import dataclasses as dtc
+
+from ..audios import Spectrogram
+from ..data import Feature, AsSlice, Input, Target
+from .parts import SuperAdam, SequenceModel, IData
+from ..networks import Seq2SeqLSTM
+from .parts.loss_functions import mean_L1_prop
+
+__all__ = [
+    'Seq2SeqLSTMModel'
+]
+
+
+@dtc.dataclass
+class Seq2SeqData(IData):
+    feature: Feature = None
+    batch_size: int = 16
+
+    @classmethod
+    def schema(cls, sr=22050, emphasis=0., n_fft=2048, hop_length=512):
+        schema = {"fft": Spectrogram(sr=sr, emphasis=emphasis,
+                                     n_fft=n_fft, hop_length=hop_length,
+                                     magspec=True)}
+        return schema
+
+    @classmethod
+    def dependant_hp(cls, db):
+        return dict(
+            feature=db.schema['fft'], input_dim=db.schema['fft'].dim
+        )
+
+    def batch_signature(self, stage='fit'):
+        inpt = Input('fft', AsSlice(shift=0, length=self.shift))
+        trgt = Target('fft', AsSlice(shift=self.shift,
+                                     length=self.shift))
+        if stage in ('full', 'fit', 'train', 'val'):
+            return inpt, trgt
+        # test, predict, generate...
+        return inpt
+
+    def loader_kwargs(self, stage, datamodule):
+        return dict(
+            batch_size=self.batch_size,
+            drop_last=False,
+            shuffle=True
+        )
+
+
+class Seq2SeqLSTMModel(
+    Seq2SeqData,
+    SequenceModel,
+    SuperAdam,
+    Seq2SeqLSTM,
+):
+
+    @staticmethod
+    def loss_fn(output, target):
+        return {"loss": mean_L1_prop(output, target)}
+
+    def encode_inputs(self, inputs: torch.Tensor):
+        return self.feature.encode(inputs)
+
+    def decode_outputs(self, outputs: torch.Tensor):
+        return self.feature.decode(outputs)
diff --git a/mimikit/models/wavenet.py b/mimikit/models/wavenet.py
new file mode 100644
index 0000000..daf5595
--- /dev/null
+++ b/mimikit/models/wavenet.py
@@ -0,0 +1,94 @@
+import dataclasses as dtc
+import torch.nn as nn
+import torch
+
+from ..abstract.features import SegmentLabels, FilesLabels
+from ..audios import MuLawSignal
+from ..data import Feature, Input, AsSlice, Target
+from . import model
+from ..networks import WNNetwork
+from .parts import IData, SuperAdam, SequenceModel
+
+__all__ = [
+    'WaveNetData',
+    'WaveNet'
+]
+
+
+@dtc.dataclass
+class WaveNetData(IData):
+    feature: Feature = None
+    batch_size: int = 32
+    batch_seq_length: int = 64
+
+    @classmethod
+    def schema(cls, sr=22050, emphasis=0., q_levels=256,
+               segment_labels=False, files_labels=False):
+
+        schema = {"qx": MuLawSignal(sr=sr, emphasis=emphasis, q_levels=q_levels)}
+        if segment_labels:
+            schema.update({
+                'loc': SegmentLabels(input_key='qx')
+            })
+        if files_labels:
+            schema.update({
+                'glob': FilesLabels(input_key='qx')
+            })
+        return schema
+
+    @classmethod
+    def dependant_hp(cls, db):
+        hp = {}
+        if 'loc' in db.features:
+            hp.update(dict(n_cin_classes=len(db.loc.regions)))
+        if 'glob' in db.features:
+            hp.update(dict(n_gin_classes=len(db.glob.files)))
+        return dict(
+            feature=db.schema['qx'], q_levels=db.schema['qx'].q_levels, **hp
+        )
+
+    def batch_signature(self, stage='fit'):
+        inpt = Input('qx', AsSlice(shift=0, length=self.batch_seq_length))
+        trgt = Target('qx', AsSlice(shift=self.shift,
+                                    length=self.output_shape((-1, self.batch_seq_length, -1))[1]))
+        # where are we conditioned?
+        loc, glob = self.n_cin_classes is not None, self.n_gin_classes is not None
+        if loc:
+            pass
+        if glob:
+            pass
+        if stage in ('full', 'fit', 'train', 'val'):
+            return inpt, trgt
+        # test, predict, generate...
+        return inpt
+
+    def loader_kwargs(self, stage, datamodule):
+        return dict(
+            batch_size=self.batch_size,
+            drop_last=False,
+            shuffle=True
+        )
+
+
+@model
+class WaveNet(
+    WaveNetData,
+    SequenceModel,
+    SuperAdam,
+    WNNetwork,
+):
+
+    @staticmethod
+    def loss_fn(output, target):
+        criterion = nn.CrossEntropyLoss(reduction="mean")
+        return {"loss": criterion(output.view(-1, output.size(-1)), target.view(-1))}
+
+    def setup(self, stage: str):
+        SequenceModel.setup(self, stage)
+        SuperAdam.setup(self, stage)
+
+    def encode_inputs(self, inputs: torch.Tensor):
+        return self.feature.encode(inputs)
+
+    def decode_outputs(self, outputs: torch.Tensor):
+        return self.feature.decode(outputs)
diff --git a/mimikit/modules/__init__.py b/mimikit/modules/__init__.py
new file mode 100644
index 0000000..2332631
--- /dev/null
+++ b/mimikit/modules/__init__.py
@@ -0,0 +1,5 @@
+from .homs import *
+from .ops import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/modules/homs.py b/mimikit/modules/homs.py
new file mode 100644
index 0000000..50d0406
--- /dev/null
+++ b/mimikit/modules/homs.py
@@ -0,0 +1,218 @@
+from dataclasses import dataclass
+from functools import wraps
+
+import torch.nn as nn
+import torch
+from copy import deepcopy
+
+
+__all__ = [
+    'HOM',
+    'HOMSequential',
+    'SequentialAdd',
+    'SequentialMul',
+    'SequentialCollect',
+    'Paths',
+    'FPaths',
+    'AddPaths',
+    'MulPaths',
+    'block',
+    'Tiers',
+    'Skips',
+    'GatedUnit'
+]
+
+# *******************************************************************
+# ************************ HOM **************************************
+# *******************************************************************
+
+
+class NotAModule(nn.Module):
+    """
+    the nn.Module equivalent of None
+    """
+    pass
+
+
+class HOM(nn.Sequential, nn.ModuleList):
+
+    @staticmethod
+    def replace_none(module):
+        return module if module is not None else NotAModule()
+
+    @staticmethod
+    def output_or_none(mod, x):
+        if type(mod) is NotAModule or x is None:
+            return None
+        return mod(x)
+
+    @staticmethod
+    def flat_inputs(*inputs):
+        return inputs[0] if len(inputs) == 1 and isinstance(inputs[0], tuple) else inputs
+
+    def __init__(self, *modules):
+        super().__init__(*(self.replace_none(m) for m in modules))
+
+    def forward(self, *inputs):
+        raise NotImplementedError
+
+
+# *******************************************************************
+# ******************* HOMSequential *********************************
+# *******************************************************************
+
+
+class HOMSequential(HOM):
+    ops_ = None
+
+    def __init__(self, *modules):
+        nn.Sequential.__init__(self, *[self.replace_none(m) for m in modules])
+
+    def forward(self, input):
+        out = None
+        for mod in self:
+            out = self.output_or_none(mod, input)
+            if out is not None:
+                if self.ops_ is not None:
+                    self.ops_(input, out)
+                else:
+                    input = out
+        return out
+
+
+class SequentialAdd(HOMSequential):
+    @property
+    def ops_(self):
+        return torch.Tensor.add_
+
+
+class SequentialMul(HOMSequential):
+    @property
+    def ops_(self):
+        return torch.Tensor.mul_
+
+
+class SequentialCollect(HOMSequential):
+
+    def forward(self, input):
+        collected = []
+        for mod in self:
+            out = self.output_or_none(mod, input)
+            if out is not None:
+                if self.ops_ is not None:
+                    self.ops_(input, out)
+                else:
+                    input = out
+            collected += [out]
+        return tuple(collected)
+
+
+# *******************************************************************
+# ************************* Paths ***********************************
+# *******************************************************************
+
+
+class Paths(HOM):
+
+    def __init__(self, *modules):
+        super(Paths, self).__init__(*modules)
+
+    def forward(self, *inputs):
+        inputs = self.flat_inputs(*inputs)
+        return tuple(self.output_or_none(mod, x) for mod, x in zip(self, inputs))
+
+
+class FPaths(Paths):
+    @property
+    def ops(self):
+        return None
+
+    def forward(self, *inputs):
+        # flatten ((x, x, ....), ) to (x, x, ...)
+        inputs = self.flat_inputs(*inputs)
+        inputs = (x for x in inputs)
+        modules = (m for m in self)
+        out = None
+        # if the firsts modules were Nones
+        while out is None:
+            out = self.output_or_none(next(modules), (next(inputs)))
+        # now accum the outputs
+        for mod, x in zip(modules, inputs):
+            output = self.output_or_none(mod, x)
+            if output is not None:
+                # all output shapes have to be the same!...
+                self.ops(out, output)
+        return out
+
+
+class AddPaths(FPaths):
+    @property
+    def ops(self):
+        return torch.Tensor.add
+
+
+class MulPaths(FPaths):
+    @property
+    def ops(self):
+        return torch.Tensor.mul
+
+
+# *******************************************************************
+# ************************ block ************************************
+# *******************************************************************
+
+
+def block(name, layer_func, container_cls=HOMSequential):
+    def init(self, n_layers, *args, **kwargs):
+        container_cls.__init__(self, *(layer_func(i, *args, **kwargs)
+                                       for sub_block in n_layers for i in range(sub_block)))
+    return type(name, (container_cls,), {"__init__": init})
+
+
+class Tiers(HOMSequential):
+
+    def forward(self, tiers_inputs):
+        output = None
+        for tier, inpt in zip(self, tiers_inputs):
+            output = tier(inpt, output)
+        return output
+
+
+# *******************************************************************
+# ************************* MISC ************************************
+# *******************************************************************
+
+class Skips(nn.Module):
+    def __init__(self, skips):
+        super().__init__()
+        self.skp = skips
+
+    def forward(self, x, skips=None):
+        h = self.skp(x)
+        if skips is None:
+            skips = torch.zeros_like(h, device=h.device)
+        skips.add_(h)
+        return skips
+
+
+class GatedUnit(MulPaths):
+
+    def __new__(cls, module):
+        return MulPaths(nn.Sequential(module, nn.Tanh()),
+                        nn.Sequential(deepcopy(module), nn.Sigmoid()))
+
+
+def propped(*bases):
+    def wrapper(cls):
+        cls = dataclass(cls, init=True, repr=False, eq=False, frozen=False)
+
+        @wraps(cls.__init__)
+        def new_init(self, *args, **kwargs):
+            for base in bases:
+                base.__init__(self)
+            cls.__init__(self, *args, **kwargs)
+
+        cls.__init__ = new_init
+        return cls
+
+    return wrapper
\ No newline at end of file
diff --git a/mimikit/modules/ops.py b/mimikit/modules/ops.py
new file mode 100644
index 0000000..b2d3a3a
--- /dev/null
+++ b/mimikit/modules/ops.py
@@ -0,0 +1,59 @@
+import torch.nn.functional as F
+from torch import nn
+
+__all__ = [
+    'Abs',
+    'Transpose',
+    'CausalPad',
+    'Slice',
+    'Clone'
+]
+
+
+class Abs(nn.Module):
+
+    def forward(self, x):
+        return x.abs_()
+
+
+class Transpose(nn.Module):
+
+    def __init__(self, dim1, dim2):
+        super(Transpose, self).__init__()
+        self.dims = (dim1, dim2)
+
+    def forward(self, *args):
+        if len(args) > 1:
+            return tuple(x.transpose(*self.dims).contiguous() if x is not None else x for x in args)
+        return args[0].transpose(*self.dims).contiguous() if args[0] is not None else None
+
+
+class CausalPad(nn.Module):
+
+    @staticmethod
+    def int_to_lr(i):
+        return (i, 0) if i >= 0 else (0, abs(i))
+
+    def __init__(self, pad, **kwargs):  # TODO: learn=True
+        super().__init__()
+        # for whatever reason torch decided dimensions should be in the reversed order...
+        self.pad = tuple(i for p in reversed(pad) for i in self.int_to_lr(p))
+        self.kwargs = kwargs
+
+    def forward(self, x):
+        return F.pad(x, self.pad, **self.kwargs)
+
+
+class Slice(nn.Module):
+
+    def __init__(self, slc):
+        super(Slice, self).__init__()
+        self.slc = slc
+
+    def forward(self, x):
+        return x[self.slc].contiguous()
+
+
+class Clone(nn.Module):
+    def forward(self, x):
+        return x.clone()
\ No newline at end of file
diff --git a/mimikit/networks/__init__.py b/mimikit/networks/__init__.py
new file mode 100644
index 0000000..de29a0b
--- /dev/null
+++ b/mimikit/networks/__init__.py
@@ -0,0 +1,10 @@
+from .freqnet import *
+from .parametrized_gaussian import *
+from .sample_rnn import *
+from .seq2seq_lstms import *
+from .tied_autoencoder import *
+from .transformers import *
+from .wavenet import *
+
+
+__all__ = [_ for _ in dir() if not _.startswith("_")]
\ No newline at end of file
diff --git a/mimikit/networks/freqnet.py b/mimikit/networks/freqnet.py
new file mode 100644
index 0000000..024bd2d
--- /dev/null
+++ b/mimikit/networks/freqnet.py
@@ -0,0 +1,97 @@
+import torch.nn as nn
+from dataclasses import dataclass
+from typing import Optional
+
+from .wavenet import WNNetwork, WaveNetLayer
+from ..modules import homs as H, ops as Ops
+
+__all__ = [
+    'FreqNetNetwork'
+]
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class FreqNetNetwork(WNNetwork):
+    """
+    adapts Wavenet to the frequency-domain
+
+    Parameters
+    ----------
+    n_layers: tuple
+
+    input_dim: int
+
+    n_cin_classes: Optional[int]
+
+    cin_dim: Optional[int]
+
+    n_gin_classes: Optional[int]
+
+    gin_dim: Optional[int]
+
+    gate_dim: int
+
+    kernel_size: int
+
+    groups: int
+
+    accum_outputs: int
+
+    pad_input: int
+
+    skip_dim: Optional[int]
+
+    residuals_dim: Optional[int]
+
+
+    """
+    n_layers: tuple = (4,)
+    input_dim: int = 256
+    n_cin_classes: Optional[int] = None
+    cin_dim: Optional[int] = None
+    n_gin_classes: Optional[int] = None
+    gin_dim: Optional[int] = None
+    gate_dim: int = 128
+    kernel_size: int = 2
+    groups: int = 1
+    accum_outputs: int = 0
+    pad_input: int = 0
+    skip_dim: Optional[int] = None
+    residuals_dim: Optional[int] = None
+
+    def inpt_(self):
+        return H.Paths(
+            nn.Sequential(
+                H.GatedUnit(nn.Linear(self.input_dim, self.gate_dim)), Ops.Transpose(1, 2)),
+            # conditioning parameters :
+            nn.Sequential(
+                nn.Embedding(self.n_cin_classes, self.cin_dim), Ops.Transpose(1, 2)) if self.cin_dim else None,
+            nn.Sequential(
+                nn.Embedding(self.n_gin_classes, self.gin_dim), Ops.Transpose(1, 2)) if self.gin_dim else None
+        )
+
+    def layers_(self):
+        return nn.Sequential(*[
+            WaveNetLayer(i,
+                         gate_dim=self.gate_dim,
+                         skip_dim=self.skip_dim,
+                         residuals_dim=self.residuals_dim,
+                         kernel_size=self.kernel_size,
+                         cin_dim=self.cin_dim,
+                         gin_dim=self.gin_dim,
+                         groups=self.groups,
+                         pad_input=self.pad_input,
+                         accum_outputs=self.accum_outputs,
+                         )
+            for block in self.n_layers for i in range(block)
+        ])
+
+    def outpt_(self):
+        return nn.Sequential(
+            Ops.Transpose(1, 2),
+            nn.Linear(self.gate_dim if self.skip_dim is None else self.skip_dim, self.input_dim), Ops.Abs()
+        )
+
+    @staticmethod
+    def predict_(outpt, temp=None):
+        return outpt
diff --git a/mimikit/networks/generating_net.py b/mimikit/networks/generating_net.py
new file mode 100644
index 0000000..e475e99
--- /dev/null
+++ b/mimikit/networks/generating_net.py
@@ -0,0 +1,30 @@
+import torch
+import torch.nn as nn
+import numpy as np
+
+from ..models.parts import tqdm
+from ..data import process_batch
+
+__all__ = [
+    'GeneratingNetwork'
+]
+
+
+class GeneratingNetwork(nn.Module):
+
+    def prepare_prompt(self, prompt, n_steps, at_least_nd=2):
+        def _prepare(prmpt):
+            if isinstance(prmpt, np.ndarray):
+                prmpt = torch.from_numpy(prmpt)
+            while len(prmpt.shape) < at_least_nd:
+                prmpt = prmpt.unsqueeze(0)
+            prmpt = prmpt.to(self.device)
+            blank_shapes = prmpt.size(0), n_steps, *prmpt.size()[2:]
+            return torch.cat((prmpt, torch.zeros(*blank_shapes).to(prmpt)), dim=1)
+
+        return process_batch(prompt, lambda x: isinstance(x, (np.ndarray, torch.Tensor)), _prepare)
+
+    @staticmethod
+    def generate_tqdm(rng):
+        return tqdm(rng, desc="Generate", dynamic_ncols=True,
+                    leave=False, unit="step", mininterval=0.25)
diff --git a/mimikit/networks/parametrized_gaussian.py b/mimikit/networks/parametrized_gaussian.py
new file mode 100644
index 0000000..e12e6c8
--- /dev/null
+++ b/mimikit/networks/parametrized_gaussian.py
@@ -0,0 +1,35 @@
+import torch
+from torch import nn as nn
+from typing import Optional
+from dataclasses import dataclass
+
+__all__ = [
+    'ParametrizedGaussian'
+]
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class ParametrizedGaussian(nn.Module):
+    """
+    Parametrized Gaussian (often found in variational auto-encoders)
+    """
+    input_dim: int
+    z_dim: int
+    bias: Optional[bool] = False
+    pre_activation: Optional[nn.Module] = nn.Identity()
+    return_params: Optional[bool] = True
+
+    def __postinit__(self):
+        nn.Module.__init__(self)
+        self.fc1 = nn.Linear(self.input_dim, self.z_dim, bias=self.bias)
+        self.fc2 = nn.Linear(self.input_dim, self.z_dim, bias=self.bias)
+        self.pre_act = self.pre_activation if self.pre_activation is not None else lambda x: x
+
+    def forward(self, h):
+        mu, logvar = self.fc1(self.pre_act(h)), self.fc2(self.pre_act(h))
+        std = logvar.mul(0.5).exp_()
+        eps = torch.randn(*mu.size()).to(h.device)
+        z = mu + std * eps
+        if self.return_params:
+            return z, mu, std
+        return z
diff --git a/mimikit/networks/sample_rnn.py b/mimikit/networks/sample_rnn.py
new file mode 100644
index 0000000..f0a7ab1
--- /dev/null
+++ b/mimikit/networks/sample_rnn.py
@@ -0,0 +1,235 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from dataclasses import dataclass
+
+from .generating_net import GeneratingNetwork
+
+__all__ = [
+    "SampleRNNTier",
+    "SampleRNNNetwork"
+]
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class SampleRNNTier(nn.Module):
+    tier_index: int
+    frame_size: int
+    dim: int
+    up_sampling: int = 1
+    n_rnn: int = 2
+    q_levels: int = 256
+    embedding_dim: Optional[int] = None
+    mlp_dim: Optional[int] = None
+
+    is_bottom = property(lambda self: self.up_sampling == 1)
+
+    def embeddings_(self):
+        if self.embedding_dim is not None:
+            return nn.Embedding(self.q_levels, self.embedding_dim)
+        return None
+
+    def input_proj_(self):
+
+        if not self.is_bottom:  # top & middle tiers
+            return nn.Sequential(nn.Linear(self.frame_size, self.dim), )
+
+        else:  # bottom tier
+            class BottomProjector(nn.Module):
+                def __init__(self, emb_dim, out_dim, frame_size):
+                    super(BottomProjector, self).__init__()
+                    self.cnv = nn.Conv1d(emb_dim, out_dim, kernel_size=frame_size)
+
+                def forward(self, hx):
+                    """ hx : (B x T x FS x E) """
+                    B, T, FS, E = hx.size()
+                    hx = self.cnv(hx.view(B * T, FS, E).transpose(1, 2).contiguous())
+                    # now hx : (B*T, DIM, 1)
+                    return hx.squeeze().reshape(B, T, -1)
+
+            return BottomProjector(self.embedding_dim, self.dim, self.frame_size)
+
+    def rnn_(self):
+        # no rnn for bottom tier
+        if self.is_bottom:
+            return None
+        # we do not learn hidden for practical reasons : storing them as params in the state dict
+        # would force the network to always generate with the train batch_size or to implement some logic
+        # to handle this shape change. We choose the simpler solution : hidden are initialized to zeros
+        # whenever we need fresh ones.
+        self.h0, self.c0 = None, None
+        return nn.LSTM(self.dim, self.dim, self.n_rnn, batch_first=True)
+
+    def up_sampling_net_(self):
+        # no up sampling for bottom tier
+        if self.is_bottom:
+            return None
+
+        class TimeUpscalerLinear(nn.Module):
+
+            def __init__(self, in_dim, out_dim, up_sampling, **kwargs):
+                super(TimeUpscalerLinear, self).__init__()
+                self.up_sampling = up_sampling
+                self.out_dim = out_dim
+                self.fc = nn.Linear(in_dim, out_dim * up_sampling, **kwargs)
+
+            def forward(self, x):
+                B, T, _ = x.size()
+                return self.fc(x).reshape(B, T * self.up_sampling, self.out_dim)
+
+        class ConvUpscaler(nn.Module):
+            def __init__(self, dim, up_sampling):
+                super(ConvUpscaler, self).__init__()
+                self.cv = nn.ConvTranspose1d(dim, dim, kernel_size=up_sampling, stride=up_sampling)
+
+            def forward(self, x):
+                return self.cv(x.transpose(1, 2).contiguous()).transpose(1, 2).contiguous()
+
+        return TimeUpscalerLinear(self.dim, self.dim, self.up_sampling)
+
+    def mlp_(self):
+        if self.mlp_dim is None:
+            return None
+        return nn.Sequential(
+            nn.Linear(self.dim, self.mlp_dim), nn.ReLU(),
+            nn.Linear(self.mlp_dim, self.mlp_dim), nn.ReLU(),
+            nn.Linear(self.mlp_dim, self.q_levels if self.is_bottom else self.dim),
+        )
+
+    def __post_init__(self):
+        nn.Module.__init__(self)
+        self.embeddings = self.embeddings_()
+        self.inpt_proj = self.input_proj_()
+        self.rnn = self.rnn_()
+        self.up_net = self.up_sampling_net_()
+        self.mlp = self.mlp_()
+
+    def reset_hidden(self, batch_size, device):
+        self.h0 = torch.zeros(self.n_rnn, batch_size, self.dim).to(device)
+        self.c0 = torch.zeros(self.n_rnn, batch_size, self.dim).to(device)
+        return self.h0, self.c0
+
+    def linearize(self, q_samples):
+        """ maps input samples (0 <= qx < 256) to floats (-2. <= x < 2.) """
+        return ((q_samples.float() / self.q_levels) - .5) * 4
+
+    def forward(self, input_samples, prev_tier_output=None, hidden=None):
+        if self.embeddings is None:
+            x = self.linearize(input_samples)
+        else:
+            x = self.embeddings(input_samples)
+
+        if self.inpt_proj is not None:
+            p = self.inpt_proj(x)
+            if prev_tier_output is not None:
+                x = p + prev_tier_output
+            else:
+                x = p
+
+        if self.rnn is not None:
+            if hidden is None or x.size(0) != hidden[0].size(1):
+                hidden = self.reset_hidden(x.size(0), x.device)
+            else:
+                # TRUNCATED back propagation through time == detach()!
+                hidden = tuple(h.detach() for h in hidden)
+            x, hidden = self.rnn(x, hidden)
+
+        if self.up_net is not None:
+            x = self.up_net(x)
+
+        if self.mlp is not None:
+            x = self.mlp(x)
+        return x, hidden
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class SampleRNNNetwork(GeneratingNetwork, nn.Module):
+
+    frame_sizes: tuple = (16, 8, 8)  # from top to bottom!
+    dim: int = 512
+    n_rnn: int = 2
+    q_levels: int = 256
+    embedding_dim: int = 256
+    mlp_dim: int = 512
+
+    def __post_init__(self):
+        nn.Module.__init__(self)
+        GeneratingNetwork.__init__(self)
+        n_tiers = len(self.frame_sizes)
+        tiers = []
+        if n_tiers > 2:
+            for i, fs in enumerate(self.frame_sizes[:-2]):
+                tiers += [SampleRNNTier(i, fs, self.dim,
+                                        up_sampling=fs // self.frame_sizes[i + 1],
+                                        n_rnn=self.n_rnn,
+                                        q_levels=self.q_levels,
+                                        embedding_dim=None,
+                                        mlp_dim=None)]
+        # before last tier
+        tiers += [SampleRNNTier(n_tiers - 2, self.frame_sizes[-2], self.dim,
+                                up_sampling=self.frame_sizes[-1],
+                                n_rnn=self.n_rnn,
+                                q_levels=self.q_levels,
+                                embedding_dim=None,
+                                mlp_dim=None)]
+        # bottom tier
+        tiers += [SampleRNNTier(n_tiers - 1, self.frame_sizes[-1], self.dim,
+                                up_sampling=1,
+                                n_rnn=self.n_rnn,
+                                q_levels=self.q_levels,
+                                embedding_dim=self.embedding_dim,
+                                mlp_dim=self.mlp_dim)]
+
+        self.tiers = nn.ModuleList(tiers)
+        self.hidden = [None] * n_tiers
+
+    def forward(self, tiers_inputs):
+        prev_out = None
+        for i, (tier, inpt) in enumerate(zip(self.tiers, tiers_inputs)):
+            prev_out, self.hidden[i] = tier(inpt, prev_out, self.hidden[i])
+        return prev_out
+
+    def reset_h0(self):
+        self.hidden = [None] * len(self.frame_sizes)
+
+    def generate_(self, prompt, n_steps, temperature=1.):
+        output = self.prepare_prompt(prompt, n_steps, at_least_nd=2)
+        # trim to start with a whole number of top frames
+        output = output[:, prompt.size(1) % self.frame_sizes[0]:]
+        prior_t = prompt.size(1) - (prompt.size(1) % self.frame_sizes[0])
+
+        # init variables
+        fs = [*self.frame_sizes]
+        outputs = [None] * (len(fs) - 1)
+        # hidden are reset if prompt.size(0) != self.hidden.size(0)
+        hiddens = self.hidden
+        tiers = self.tiers
+
+        for t in self.generate_tqdm(range(fs[0], n_steps + prior_t)):
+            for i in range(len(tiers) - 1):
+                if t % fs[i] == 0:
+                    inpt = output[:, t - fs[i]:t].unsqueeze(1)
+
+                    if i == 0:
+                        prev_out = None
+                    else:
+                        prev_out = outputs[i - 1][:, (t // fs[i]) % (fs[i - 1] // fs[i])].unsqueeze(1)
+
+                    out, h = tiers[i](inpt, prev_out, hiddens[i])
+                    hiddens[i] = h
+                    outputs[i] = out
+            if t < prior_t:  # only used for warming-up
+                continue
+            prev_out = outputs[-1]
+            inpt = output[:, t - fs[-1]:t].reshape(-1, 1, fs[-1])
+
+            out, _ = tiers[-1](inpt, prev_out[:, (t % fs[-1]) - fs[-1]].unsqueeze(1))
+            if temperature is None:
+                pred = (nn.Softmax(dim=-1)(out.squeeze(1))).argmax(dim=-1)
+            else:
+                # great place to implement dynamic cooling/heating !
+                pred = torch.multinomial(nn.Softmax(dim=-1)(out.squeeze(1) / temperature), 1)
+            output.data[:, t] = pred.squeeze()
+
+        return output
+
diff --git a/mimikit/networks/seq2seq_lstms.py b/mimikit/networks/seq2seq_lstms.py
new file mode 100644
index 0000000..7d70f53
--- /dev/null
+++ b/mimikit/networks/seq2seq_lstms.py
@@ -0,0 +1,133 @@
+import torch
+import torch.nn as nn
+from typing import Optional
+from dataclasses import dataclass
+
+from ..networks.parametrized_gaussian import ParametrizedGaussian
+
+__all__ = [
+    'EncoderLSTM',
+    'DecoderLSTM',
+    'Seq2SeqLSTM',
+]
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class EncoderLSTM(nn.Module):
+    input_d: int
+    model_dim: int
+    num_layers: int
+    n_lstm: Optional[int] = 1
+    bottleneck: Optional[str] = "add"
+    n_fc: Optional[int] = 1
+    bias: Optional[bool] = False
+    weight_norm: Optional[bool] = False
+
+    def __postinit__(self):
+        nn.Module.__init__(self)
+        self.lstms = nn.ModuleList([
+            nn.LSTM(self.input_d if i == 0 else self.model_dim,
+                    self.model_dim if self.bottleneck == "add" else self.model_dim // 2,
+                    bias=self.bias,
+                    num_layers=self.num_layers,
+                    batch_first=True, bidirectional=True)
+            for i in range(self.n_lstm)
+        ])
+        self.fc = nn.Sequential(
+            *[nn.Sequential(nn.Linear(self.model_dim, self.model_dim), nn.Tanh()) for _ in range(self.n_fc - 1)],
+            nn.Linear(self.model_dim, self.model_dim, bias=False),  # NO ACTIVATION !
+        )
+        if self.weight_norm:
+            for name, p in dict(self.lstms.named_parameters()).items():
+                if "weight" in name:
+                    torch.nn.utils.weight_norm(self.lstms, name)
+
+    def forward(self, x, h0=None, c0=None):
+        if h0 is None or c0 is None:
+            hidden = tuple()
+        else:
+            hidden = (h0, c0)
+        ht, ct = None, None
+        for i, lstm in enumerate(self.lstms):
+            out, (ht, ct) = lstm(x, *hidden)
+            # sum forward and backward nets
+            out = out.view(*out.size()[:-1], self.model_dim, 2).sum(dim=-1)
+            # take residuals AFTER the first lstm
+            x = out if i == 0 else x + out
+        states = self.first_and_last_states(x)
+        return self.fc(states), (ht, ct)
+
+    def first_and_last_states(self, sequence):
+        first_states = sequence[:, 0, :]
+        last_states = sequence[:, -1, :]
+        if self.bottleneck == "add":
+            return first_states + last_states
+        else:
+            return torch.cat((first_states, last_states), dim=-1)
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class DecoderLSTM(nn.Module):
+    model_dim: int
+    num_layers: int
+    bottleneck: Optional[str] = "add"
+    bias: Optional[bool] = False
+    weight_norm: Optional[tuple] = (False, False)
+
+    def __postinit__(self):
+        nn.Module.__init__(self)
+        self.lstm1 = nn.LSTM(self.model_dim, self.model_dim if self.bottleneck == "add" else self.model_dim // 2,
+                             bias=self.bias,
+                             num_layers=self.num_layers, batch_first=True, bidirectional=True)
+        self.lstm2 = nn.LSTM(self.model_dim, self.model_dim if self.bottleneck == "add" else self.model_dim // 2,
+                             bias=self.bias,
+                             num_layers=self.num_layers, batch_first=True, bidirectional=True)
+        for lstm, wn in zip([self.lstm1, self.lstm2], self.weight_norm):
+            if wn:
+                for name, p in dict(lstm.named_parameters()).items():
+                    if "weight" in name:
+                        torch.nn.utils.weight_norm(lstm, name)
+
+    def forward(self, x, hiddens, cells):
+        if hiddens is None or cells is None:
+            output, (_, _) = self.lstm1(x)
+        else:
+            # ALL decoders get hidden states from encoder
+            output, (_, _) = self.lstm1(x, (hiddens, cells))
+        # sum forward and backward nets
+        output = output.view(*output.size()[:-1], self.model_dim, 2).sum(dim=-1)
+
+        if hiddens is None or cells is None:
+            output2, (hiddens, cells) = self.lstm2(output)
+        else:
+            output2, (hiddens, cells) = self.lstm2(output, (hiddens, cells))
+        # sum forward and backward nets
+        output2 = output2.view(*output2.size()[:-1], self.model_dim, 2).sum(dim=-1)
+
+        # sum the outputs
+        return output + output2, (hiddens, cells)
+
+
+class Seq2SeqLSTM(nn.Module):
+    def __init__(self,
+                 input_dim,
+                 model_dim,
+                 num_layers=1,
+                 n_lstm=1,
+                 bottleneck="add",
+                 n_fc=1):
+        super(Seq2SeqLSTM, self).__init__()
+        self.enc = EncoderLSTM(input_dim, model_dim, num_layers, n_lstm, bottleneck, n_fc)
+        self.dec = DecoderLSTM(model_dim, num_layers, bottleneck)
+        self.sampler = ParametrizedGaussian(model_dim, model_dim)
+        self.fc_out = nn.Linear(model_dim, input_dim, bias=False)
+
+    def forward(self, x, output_length=None):
+        coded, (h_enc, c_enc) = self.enc(x)
+        if output_length is None:
+            output_length = x.size(1)
+        coded = coded.unsqueeze(1).repeat(1, output_length, 1)
+        residuals, _, _ = self.sampler(coded)
+        coded = coded + residuals
+        output, (_, _) = self.dec(coded, h_enc, c_enc)
+        return self.fc_out(output).abs()
diff --git a/mimikit/networks/tied_autoencoder.py b/mimikit/networks/tied_autoencoder.py
new file mode 100644
index 0000000..0d35e4a
--- /dev/null
+++ b/mimikit/networks/tied_autoencoder.py
@@ -0,0 +1,30 @@
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+
+__all__ = [
+    "TiedAE"
+]
+
+
+class TiedAE(nn.Module):
+
+    def __init__(self, **kwargs):
+        super(TiedAE, self).__init__()
+
+        B, K = 8, 4
+
+        self.inpt = nn.Conv1d(1025, B, K, padding=1, bias=False)
+        self.permute = lambda x: x.transpose(1, 2)
+
+    def forward(self, x):
+        x = self.permute(x)
+
+        x = F.conv1d(x, self.inpt.weight, padding=1).abs()
+        wwt = torch.matmul(self.inpt.weight.sum(dim=2), self.inpt.weight.sum(dim=2).t())
+        sim2 = 1 * F.l1_loss(wwt,
+                             torch.eye(wwt.size(0)).to(wwt) * 1)
+
+        x = F.conv_transpose1d(x, self.inpt.weight, padding=1)
+        x = self.permute(x).abs()
+        return x, 0 + sim2
diff --git a/mimikit/networks/transformers.py b/mimikit/networks/transformers.py
new file mode 100644
index 0000000..8f51b05
--- /dev/null
+++ b/mimikit/networks/transformers.py
@@ -0,0 +1,88 @@
+import torch.nn as nn
+import torch
+from torch.nn import Transformer, TransformerEncoder, TransformerEncoderLayer
+import math
+
+__all__ = [
+    'PositionalEncoding',
+    'SimpleTransformer'
+]
+
+
+class PositionalEncoding(nn.Module):
+    r"""Inject some information about the relative or absolute position of the tokens
+        in the sequence. The positional encodings have the same dimension as
+        the embeddings, so that the two can be summed. Here, we use sine and cosine
+        functions of different frequencies.
+    .. math::
+        \text{PosEncoder}(pos, 2i) = sin(pos/10000^(2i/d_model))
+        \text{PosEncoder}(pos, 2i+1) = cos(pos/10000^(2i/d_model))
+        \text{where pos is the word position and i is the embed idx)
+    Args:
+        d_model: the embed dim (required).
+        dropout: the dropout value (default=0.1).
+        max_len: the max. length of the incoming sequence (default=5000).
+    Examples:
+        >>> pos_encoder = PositionalEncoding(d_model)
+    """
+
+    def __init__(self, d_model, dropout=0.1, max_len=5000):
+        super(PositionalEncoding, self).__init__()
+        self.dropout = nn.Dropout(p=dropout)
+
+        pe = torch.zeros(max_len, d_model)
+        position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model))
+        pe[:, 0::2] = torch.sin(position * div_term)
+        pe[:, 1::2] = torch.cos(position * div_term)
+        pe = pe.unsqueeze(0).transpose(0, 1)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x):
+        r"""Inputs of forward function
+        Args:
+            x: the sequence fed to the positional encoder model (required).
+        Shape:
+            x: [sequence length, batch size, embed dim]
+            output: [sequence length, batch size, embed dim]
+        Examples:
+            >>> output = pos_encoder(x)
+        """
+
+        x = x + self.pe[:x.size(0), :]
+        return self.dropout(x)
+
+
+class SimpleTransformer(nn.Module):
+
+    def _generate_square_subsequent_mask(self, sz):
+        mask = (torch.triu(torch.ones(sz, sz)) == 1).transpose(0, 1)
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def _generate_square_context_mask(self, sz, k=8):
+        mask = torch.zeros(sz, sz)
+        rg = torch.arange(k)
+        for i in range(sz):
+            mask[i, torch.clamp(rg + i - (2 // k), min=0, max=sz - 1)] = 1.
+        mask = mask.float().masked_fill(mask == 0, float('-inf')).masked_fill(mask == 1, float(0.0))
+        return mask
+
+    def __init__(self, d_model, nhead, dim_feedforward, num_layers):
+        super(SimpleTransformer, self).__init__()
+        layer = TransformerEncoderLayer(d_model=d_model,
+                                        nhead=nhead,
+                                        dim_feedforward=dim_feedforward,
+                                        dropout=0.0,
+                                        activation="gelu")
+        norm = nn.LayerNorm(d_model)
+        self.model = TransformerEncoder(layer, num_layers=num_layers, norm=None)
+        self.src_mask = None
+
+    def forward(self, src):
+        src = src.permute(1, 0, 2).contiguous()
+        device = src.device
+        if self.src_mask is None or self.src_mask.size(0) != len(src):
+            mask = self._generate_square_context_mask(len(src)).to(device)
+            self.src_mask = mask
+        return self.model(src, self.src_mask).permute(1, 0, 2).contiguous()
\ No newline at end of file
diff --git a/mimikit/networks/wavenet.py b/mimikit/networks/wavenet.py
new file mode 100644
index 0000000..690e943
--- /dev/null
+++ b/mimikit/networks/wavenet.py
@@ -0,0 +1,333 @@
+import math
+import torch
+import torch.nn as nn
+from dataclasses import dataclass
+from typing import Optional
+
+from ..modules import homs as H, ops as Ops
+from .generating_net import GeneratingNetwork
+
+__all__ = [
+    'WaveNetLayer',
+    'WNNetwork'
+]
+
+
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class WaveNetLayer(nn.Module):
+
+    layer_i: int
+    gate_dim: int = 128
+    residuals_dim: Optional[int] = None
+    skip_dim: Optional[int] = None
+    kernel_size: int = 2
+    groups: int = 1
+    cin_dim: Optional[int] = None
+    gin_dim: Optional[int] = None
+    pad_input: int = 1
+    accum_outputs: int = -1
+
+    stride: int = 1
+    bias: bool = True
+
+    dilation = property(lambda self: self.kernel_size ** self.layer_i)
+
+    shift_diff = property(
+        lambda self:
+        (self.kernel_size - 1) * self.dilation if self.pad_input == 0 else 0
+    )
+    input_padding = property(
+        lambda self:
+        self.pad_input * (self.kernel_size - 1) * self.dilation if self.pad_input else 0
+    )
+    output_padding = property(
+        lambda self:
+        - self.accum_outputs * self.shift_diff
+    )
+    receptive_field = property(
+        lambda self:
+        self.kernel_size * self.dilation
+    )
+
+    @property
+    def conv_kwargs(self):
+        return dict(kernel_size=self.kernel_size, dilation=self.dilation,
+                    stride=self.stride, bias=self.bias, groups=self.groups)
+
+    @property
+    def kwargs_1x1(self):
+        return dict(kernel_size=1, bias=self.bias, groups=self.groups)
+
+    def gcu_(self):
+        return H.GatedUnit(
+            H.AddPaths(
+                # core
+                nn.Conv1d(self.gate_dim, self.residuals_dim, **self.conv_kwargs),
+                # conditioning parameters :
+                nn.Conv1d(self.cin_dim, self.residuals_dim, **self.kwargs_1x1) if self.cin_dim else None,
+                nn.Conv1d(self.gin_dim, self.residuals_dim, **self.kwargs_1x1) if self.gin_dim else None
+            ))
+
+    def residuals_(self):
+        return nn.Conv1d(self.residuals_dim, self.gate_dim, **self.kwargs_1x1)
+
+    def skips_(self):
+        return H.Skips(
+            nn.Conv1d(self.residuals_dim, self.skip_dim, **self.kwargs_1x1)
+        )
+
+    def accumulator(self):
+        return H.AddPaths(nn.Identity(), nn.Identity())
+
+    def __post_init__(self):
+        nn.Module.__init__(self)
+
+        with_residuals = self.residuals_dim is not None
+        if not with_residuals:
+            self.residuals_dim = self.gate_dim
+
+        self.gcu = self.gcu_()
+
+        if with_residuals:
+            self.residuals = self.residuals_()
+        else:  # keep the signature but just pass through
+            self.residuals = nn.Identity()
+
+        if self.skip_dim is not None:
+            self.skips = self.skips_()
+        else:
+            self.skips = lambda y, skp: skp
+
+        if self.accum_outputs:
+            self.accum = self.accumulator()
+        else:
+            self.accum = lambda y, x: y
+
+    def forward(self, inputs):
+        x, cin, gin, skips = inputs
+        if self.pad_input == 0:
+            cause = self.shift_diff if x.size(2) > self.kernel_size else self.kernel_size - 1
+            slc = slice(cause, None) if self.accum_outputs <= 0 else slice(None, -cause)
+            padder = nn.Identity()
+        else:
+            slc = slice(None)
+            padder = Ops.CausalPad((0, 0, self.input_padding))
+        y = self.gcu(padder(x), cin, gin)
+        if skips is not None and y.size(-1) != skips.size(-1):
+            skips = skips[:, :, slc]
+        skips = self.skips(y, skips)
+        y = self.accum(self.residuals(y), x[:, :, slc])
+        return y, cin, gin, skips
+
+    def output_length(self, input_length):
+        if bool(self.pad_input):
+            # no matter what, padding input gives the same output shape
+            return input_length
+        # output is gonna be less than input
+        numerator = input_length - self.dilation * (self.kernel_size - 1) - 1
+        return math.floor(1 + numerator / self.stride)
+
+
+# Finally : define the network class
+@dataclass(init=True, repr=False, eq=False, frozen=False, unsafe_hash=True)
+class WNNetwork(GeneratingNetwork):
+
+    n_layers: tuple = (4,)
+    q_levels: int = 256
+    n_cin_classes: Optional[int] = None
+    cin_dim: Optional[int] = None
+    n_gin_classes: Optional[int] = None
+    gin_dim: Optional[int] = None
+    gate_dim: int = 128
+    kernel_size: int = 2
+    groups: int = 1
+    accum_outputs: int = 0
+    pad_input: int = 0
+    skip_dim: Optional[int] = None
+    residuals_dim: Optional[int] = None
+    head_dim: Optional[int] = None
+
+    def inpt_(self):
+        return H.Paths(
+            nn.Sequential(nn.Embedding(self.q_levels, self.gate_dim), Ops.Transpose(1, 2)),
+            nn.Sequential(nn.Embedding(self.n_cin_classes, self.cin_dim),
+                          Ops.Transpose(1, 2)) if self.cin_dim else None,
+            nn.Sequential(nn.Embedding(self.n_gin_classes, self.gin_dim), Ops.Transpose(1, 2)) if self.gin_dim else None
+        )
+
+    def layers_(self):
+        return nn.Sequential(*[
+            WaveNetLayer(i,
+                         gate_dim=self.gate_dim,
+                         skip_dim=self.skip_dim,
+                         residuals_dim=self.residuals_dim,
+                         kernel_size=self.kernel_size,
+                         cin_dim=self.cin_dim,
+                         gin_dim=self.gin_dim,
+                         groups=self.groups,
+                         pad_input=self.pad_input,
+                         accum_outputs=self.accum_outputs,
+                         )
+            for block in self.n_layers for i in range(block)
+        ])
+
+    def outpt_(self):
+        return nn.Sequential(
+            nn.ReLU(),
+            nn.Conv1d(self.gate_dim if self.skip_dim is None else self.skip_dim,
+                      self.gate_dim if self.head_dim is None else self.head_dim, kernel_size=1),
+            nn.ReLU(),
+            nn.Conv1d(self.gate_dim if self.head_dim is None else self.head_dim,
+                      self.q_levels, kernel_size=1),
+            Ops.Transpose(1, 2)
+        )
+
+    def __post_init__(self):
+        nn.Module.__init__(self)
+        self.inpt = self.inpt_()
+        self.layers = self.layers_()
+        self.outpt = self.outpt_()
+
+        rf = 0
+        for i, layer in enumerate(self.layers):
+            if i == (len(self.layers) - 1):
+                rf += layer.receptive_field
+            elif self.layers[i + 1].layer_i == 0:
+                rf += layer.receptive_field - 1
+        self.receptive_field = rf
+
+        if self.pad_input == 1:
+            self.shift = 1
+        elif self.pad_input == -1:
+            self.shift = self.receptive_field
+        else:
+            self.shift = sum(layer.shift_diff for layer in self.layers) + 1
+
+    def forward(self, xi, cin=None, gin=None):
+        x, cin, gin = self.inpt(xi, cin, gin)
+        y, _, _, skips = self.layers((x, cin, gin, None))
+        return self.outpt(skips if skips is not None else y)
+
+    def output_shape(self, input_shape):
+        return input_shape[0], self.all_output_lengths(input_shape[1])[-1], input_shape[-1]
+
+    def all_output_lengths(self, input_length):
+        out_length = input_length
+        lengths = []
+        for layer in self.layers:
+            out_length = layer.output_length(out_length)
+            lengths += [out_length]
+        return tuple(lengths)
+
+    def generation_slices(self):
+        # input is always the last receptive field
+        input_slice = slice(-self.receptive_field, None)
+        if self.pad_input == 1:
+            # then there's only one future time-step : the last output
+            output_slice = slice(-1, None)
+        else:
+            # for all other cases, the first output has the shift of the whole network
+            output_slice = slice(None, 1)
+        return input_slice, output_slice
+
+    @staticmethod
+    def predict_(outpt, temp):
+        if temp is None:
+            return nn.Softmax(dim=-1)(outpt).argmax(dim=-1, keepdims=True)
+        else:
+            return torch.multinomial(nn.Softmax(dim=-1)(outpt / temp), 1)
+
+    def generate_(self, prompt, n_steps, temperature=0.5, benchmark=False):
+        return self.generate_slow(prompt, n_steps, temperature)
+
+    def generate_slow(self, prompt, n_steps, temperature=0.5):
+
+        output = self.prepare_prompt(prompt, n_steps, at_least_nd=2)
+        prior_t = prompt.size(1)
+        rf = self.receptive_field
+        _, out_slc = self.generation_slices()
+
+        for t in self.generate_tqdm(range(prior_t, prior_t + n_steps)):
+            output.data[:, t:t + 1] = self.predict_(self.forward(output[:, t - rf:t])[:, out_slc],
+                                                    temperature)
+        return output
+
+    def generate_fast(self, prompt, n_steps, temperature=0.5):
+
+        output = self.prepare_prompt(prompt, n_steps, at_least_nd=2)
+        prior_t = prompt.size(1)
+
+        inpt = output[:, prior_t - self.receptive_field:prior_t]
+        z, cin, gin = self.inpt(inpt, None, None)
+        qs = [(z.clone(), None)]
+        # initialize queues with one full forward pass
+        skips = None
+        for layer in self.layers:
+            z, _, _, skips = layer((z, cin, gin, skips))
+            qs += [(z.clone(), skips.clone() if skips is not None else skips)]
+
+        outpt = self.outpt(skips if skips is not None else z)[:, -1:].squeeze()
+        outpt = self.predict_(outpt, temperature)
+        output.data[:, prior_t:prior_t + 1] = outpt
+
+        qs = {i: q for i, q in enumerate(qs)}
+
+        # disable padding and dilation
+        dilations = {}
+        for mod in self.modules():
+            if isinstance(mod, nn.Conv1d) and mod.dilation != (1,):
+                dilations[mod] = mod.dilation
+                mod.dilation = (1,)
+        for layer in self.layers:
+            layer.pad_input = 0
+
+        # cache the indices of the inputs for each layer
+        lyr_slices = {}
+        for l, layer in enumerate(self.layers):
+            d, k = layer.dilation, layer.kernel_size
+            rf = d * k
+            indices = [qs[l][0].size(2) - 1 - n for n in range(0, rf, d)][::-1]
+            lyr_slices[l] = torch.tensor(indices).long().to(self.device)
+
+        for t in self.generate_tqdm(range(prior_t + 1, prior_t + n_steps)):
+
+            x, cin, gin = self.inpt(output[:, t - 1:t], None, None)
+            q, _ = qs[0]
+            q = q.roll(-1, 2)
+            q[:, :, -1:] = x
+            qs[0] = (q, None)
+
+            for l, layer in enumerate(self.layers):
+                z, skips = qs[l]
+                zi = torch.index_select(z, 2, lyr_slices[l])
+                if skips is not None:
+                    # we only need one skip : the first or the last of the kernel's indices
+                    i = lyr_slices[l][0 if layer.accum_outputs > 0 else -1].item()
+                    skips = skips[:, :, i:i + 1]
+
+                z, _, _, skips = layer((zi, cin, gin, skips))
+
+                if l < len(self.layers) - 1:
+                    q, skp = qs[l + 1]
+
+                    q = q.roll(-1, 2)
+                    q[:, :, -1:] = z
+                    if skp is not None:
+                        skp = skp.roll(-1, 2)
+                        skp[:, :, -1:] = skips
+
+                    qs[l + 1] = (q, skp)
+                else:
+                    y, skips = z, skips
+
+            outpt = self.outpt(skips if skips is not None else y).squeeze()
+            outpt = self.predict_(outpt, temperature)
+            output.data[:, t:t + 1] = outpt
+
+        # reset the layers' parameters
+        for mod, d in dilations.items():
+            mod.dilation = d
+        for layer in self.layers:
+            layer.pad_input = self.pad_input
+
+        return output
\ No newline at end of file
diff --git a/mimikit/utils.py b/mimikit/utils.py
index bb901d8..2321567 100644
--- a/mimikit/utils.py
+++ b/mimikit/utils.py
@@ -4,7 +4,12 @@
 import IPython.display as ipd
 import matplotlib.pyplot as plt
 
-from .data.transforms import HOP_LENGTH, SR
+__all__ = [
+    'audio',
+    'show'
+]
+
+HOP_LENGTH, SR = 512, 22050
 
 # Conversion
 
diff --git a/requirements.txt b/requirements.txt
index 0ee2e4c..7687685 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,10 +1,11 @@
 numpy==1.19.1
 pandas==1.1.3
 librosa==0.8
+torchaudio>=0.8.0
 h5py==2.10.0
-tables==3.5.1
+tables>=3.6
 torch>=1.6.0
-pytorch-lightning==1.1.2
+pytorch-lightning==1.2.4
 tqdm==4.42.1
 matplotlib
 scipy==1.4.1
@@ -14,3 +15,6 @@ neptune-client==0.4.130
 neptune-contrib==0.24.9
 test-tube>=0.7.5
 ipython>=7.12.0
+muspy==0.3.0
+PyYAML!=5.4.*,>=5.1
+chardet==3.0.4
\ No newline at end of file
diff --git a/setup.py b/setup.py
index 2c2189e..948216e 100644
--- a/setup.py
+++ b/setup.py
@@ -37,8 +37,8 @@
     "long_description_content_type": "text/markdown",
     'author': 'Antoine Daurat',
     'author_email': 'ktonalberlin@gmail.com',
-    'url': 'https://github.com/k-tonal/mimikit',
-    'download_url': 'https://github.com/k-tonal/mimikit',
+    'url': 'https://github.com/ktonal/mimikit',
+    'download_url': 'https://github.com/ktonal/mimikit',
     'license': 'GNU General Public License v3 (GPLv3)',
     'classifiers': [
         'Development Status :: 3 - Alpha',
@@ -60,7 +60,8 @@
     'packages': PACKAGES,
     "entry_points": {
         'console_scripts': [
-            'freqnet-db=mimikit.data.freqnet_db:main'
+            'samplernn=mimikit.models.sample_rnn:main',
+            'freqnet=mimikit.models.freqnet:main'
         ]}
 
 }


Test patch:
diff --git a/tests/test_data.py b/tests/test_data.py
index 43950f2..18795ee 100644
--- a/tests/test_data.py
+++ b/tests/test_data.py
@@ -1,11 +1,12 @@
-import pytest
 import numpy as np
 import pandas as pd
+import pytest
 import soundfile
-import h5py
 
-from mimikit.data.factory import AudioFileWalker, file_to_db, make_root_db
-from mimikit.data import default_extract_func, DataObject, Database
+import mimikit.audios.fmodules as A
+from mimikit.file_walker import FileWalker
+from mimikit.data import Database
+import mimikit.audios.features as F
 
 
 @pytest.fixture
@@ -29,49 +30,37 @@ def audio_tree(tmp_path):
 
 
 def test_audio_file_walker(audio_tree):
-    walker = AudioFileWalker(roots=audio_tree)
+    walker = FileWalker('audio', sources=audio_tree)
     assert len(list(walker)) == 4
-    walker = AudioFileWalker(files=[audio_tree + "/dir2/test1.wav",
-                                    audio_tree + "/dir1/test4.notaudio"])
+    walker = FileWalker('audio', sources=[audio_tree + "/dir2/test1.wav",
+                                          audio_tree + "/dir1/test4.notaudio"])
     assert len(list(walker)) == 1
 
 
-def test_default_extract_func(audio_tree):
-    dic = default_extract_func(audio_tree + "/dir2/test1.wav")
-    # we must have regions
-    assert "regions" in dic and isinstance(dic["regions"][1], pd.DataFrame), dic["regions"]
-    # regions + feature
-    assert len(dic) == 2
-
-
-def test_file_to_db(audio_tree):
+class TestDB(Database):
+    y = None
+    features = ["y"]
 
-    file_path = audio_tree + "/dir2/test1.wav"
+    @staticmethod
+    def extract(path, **kwargs):
+        return dict(y=({}, A.FileToSignal(sr=16000)(path), None))
 
-    def faulty_extract_func(path):
-        dic = default_extract_func(path)
-        dic.pop("regions")
-        return dic
 
-    with pytest.raises(ValueError, match=r".*regions.*"):
-        file_to_db(file_path, faulty_extract_func)
+def test_Database_make(audio_tree):
 
-    file_to_db(file_path)
-    with h5py.File(file_path.split(".")[0] + ".h5", "r") as f:
-        assert f["regions"] is not None
+    TestDB.make(audio_tree + "/test_db.h5", items=audio_tree)
 
+    db = TestDB(audio_tree + "/test_db.h5")
+    assert isinstance(db.y.files, pd.DataFrame)
+    assert len(db.y.files) == 4, len(db.y.files)
+    assert np.any(db.y[:40] != 0), db.y[:40]
 
-def test_make_root_db_and_Database(audio_tree):
-    make_root_db(audio_tree + "/test_db.h5", roots=audio_tree, extract_func=default_extract_func, n_cores=1)
 
-    db = Database(audio_tree + "/test_db.h5")
-    assert db.regions is not pd.DataFrame()
-    assert len(db.regions) == 4, len(db.regions)
-    assert np.any(db.fft[:4] != 0), db.fft[:4]
-    assert isinstance(db.fft.get(db.regions.iloc[:1]), np.ndarray), db.fft.get(db.regions.iloc[:1])
+def test_Database_build(audio_tree):
+    fdict = dict(y=F.AudioSignal())
+    TestDB.create(audio_tree + "/test_db.h5", sources=audio_tree, schema=fdict)
 
-    # test Dataset Integration
-    ds = DataObject(db.fft)
-    assert ds is not None, ds
-    assert len(ds) == len(db.fft), (len(ds), len(db.fft))
-    assert np.all(ds[:10] == db.fft[:10])
+    db = TestDB(audio_tree + "/test_db.h5")
+    assert isinstance(db.y.files, pd.DataFrame)
+    assert len(db.y.files) == 4, len(db.y.files)
+    assert np.any(db.y[:40] != 0), db.y[:40]
diff --git a/tests/test_dataobject.py b/tests/test_dataobject.py
deleted file mode 100644
index 8c7662e..0000000
--- a/tests/test_dataobject.py
+++ /dev/null
@@ -1,144 +0,0 @@
-import pytest
-import numpy as np
-import torch
-
-from mimikit.data import DataObject
-
-
-class Case:
-
-    @property
-    def ds(self):
-        if self._ds is None:
-            self._ds = DataObject(self.data_object)
-        return self._ds
-
-    def __init__(self, data_object, expected):
-        self.data_object = data_object
-        self.expected = expected
-        self._ds = None
-
-    def __repr__(self):
-        return str(self.data_object)
-
-
-valid_data_objects = [
-    Case(np.random.randn(10, 10),
-         dict(shape=((10, 10),),
-              style=("map",),
-              dtype=(torch.float64,),
-              device=("cpu",),
-              )),
-    Case(torch.randn(10, 10, requires_grad=False),
-         dict(shape=((10, 10),),
-              style=("map",),
-              dtype=(torch.float32,),
-              device=("cpu",),
-              )),
-    Case((np.random.randn(10, 10), np.random.randn(10, 10)),
-         dict(shape=((10, 10), (10, 10)),
-              style=("map", "map"),
-              dtype=(torch.float64, torch.float64),
-              device=("cpu", "cpu"),
-              )),
-    Case([list(range(10)) for _ in range(10)],
-         dict(shape=((10, 10),),
-              style=("map",),
-              dtype=(torch.int64,),
-              device=("cpu",),
-              )),
-    Case((list(range(10)) for _ in range(10)),
-         dict(shape=((None, 10),),
-              style=("iter",),
-              dtype=(torch.int64,),
-              device=("cpu",),
-              ))
-]
-
-
-@pytest.fixture(params=valid_data_objects)
-def init_valid_case(request):
-    return request.param
-
-
-def test_computed_properties_of_valid_cases(init_valid_case):
-    case = init_valid_case
-    to_check = case.expected.keys()
-    for attr in to_check:
-        computed = getattr(case.ds, attr)
-        expected = case.expected[attr]
-        assert computed == expected, (attr, computed, expected)
-
-    assert case.ds.n_features == len(case.expected["shape"]), (case.ds.n_features, len(case.expected["shape"]))
-
-
-def test_core_methods(init_valid_case):
-    case = init_valid_case
-    func = torch.all if isinstance(case.ds.get_elem()[0], torch.Tensor) else np.all
-    if "map" in case.expected["style"]:
-        assert func(case.ds.get_elem()[0] == case.ds[0][0])
-        assert len(case.ds) == 10
-        assert len(case.ds[(0, 1, 2)]) == case.ds.n_features
-    else:
-        assert func(case.ds.get_elem()[0] == next(iter(case.ds)))
-
-
-def test_to_tensor(init_valid_case):
-    case = init_valid_case
-    if "iter" in case.ds.style:
-        with pytest.raises(TypeError):
-            case.ds.to_tensor()
-        return
-    case.ds.to_tensor()
-    if isinstance(case.ds.data, tuple):
-        assert all(isinstance(dataset.data, torch.Tensor) for dataset in case.ds.data)
-    else:
-        assert isinstance(case.ds.data, torch.Tensor), type(case.ds.data)
-
-
-def test_to_device():
-    pass
-
-
-def test_select():
-    pass
-
-
-def test_split():
-    pass
-
-
-def test_to_datamodule():
-    pass
-
-
-invalid_data_objects = [
-    Case(None,
-         dict(style=ValueError, msg="None")
-         ),
-    Case((np.random.randn(10, 10), (x for x in range(10))),
-         dict(style=TypeError, msg="same style")
-         ),
-    Case(np.random.randn,
-         dict(style=ValueError, msg="implement")
-         ),
-    Case((np.random.randn(12, 10), np.random.randn(10, 10)),
-         dict(style=ValueError, msg="same lengths")
-         ),
-    Case(["this string's dtype isn't recognized"],
-         dict(dtype=TypeError, msg="numerical")
-         )
-]
-
-
-@pytest.fixture(params=invalid_data_objects)
-def init_invalid_case(request):
-    return request.param
-
-
-def test_constructor_exceptions(init_invalid_case):
-    case = init_invalid_case
-    exception = [e for key, e in case.expected.items() if isinstance(e, type)][0]
-    msg = case.expected["msg"]
-    with pytest.raises(exception, match=r".* " + msg + r".*"):
-        assert case.ds is not None
diff --git a/tests/test_freqnets.py b/tests/test_freqnets.py
deleted file mode 100644
index f72b9be..0000000
--- a/tests/test_freqnets.py
+++ /dev/null
@@ -1,230 +0,0 @@
-import pytest
-import numpy as np
-import torch
-import soundfile
-
-from mimikit.kit import get_trainer
-from mimikit.freqnet import *
-from mimikit.data import freqnet_db
-
-
-@pytest.fixture
-def audio_tree(tmp_path):
-    root = (tmp_path / "audios")
-    root.mkdir()
-    dir1 = (root / "dir1")
-    dir1.mkdir()
-    dir2 = (root / "dir2")
-    dir2.mkdir()
-    extensions = [".aif",
-                  ".wav",
-                  ".mp3",
-                  ".aiff",
-                  ".notaudio",
-                  ""]
-    tests = [np.random.randn(10000) for _ in range(len(extensions))]
-    for i, arr, ext in zip(range(len(extensions)), tests, extensions):
-        soundfile.write(str([dir1, dir2][i % 2]) + "/test" + str(i) + ext, arr, 22050, 'PCM_24', format="WAV")
-    return str(root)
-
-
-def test_freqnet_db(audio_tree):
-    db = freqnet_db(audio_tree + "/test_db.h5", roots=audio_tree)
-    assert db is not None
-
-
-def test_base(tmp_path):
-    class TestModel(FreqNetModel):
-        def __init__(self, **kwargs):
-            super(TestModel, self).__init__(**kwargs)
-            self.fc = torch.nn.Linear(1025, 1025)
-            self.loss_fn = torch.nn.MSELoss()
-            # torch automatically converts numpy to torch.float64
-            # but then the parameters are torch.float32 by default...
-            self.datamodule.ds.to_tensor()
-            self.datamodule.ds.to(torch.float32)
-
-        def forward(self, x):
-            return self.fc(x)
-
-        def targets_shifts_and_lengths(self, input_length):
-            # shift and length of each target
-            return [(4, input_length), ]
-
-    base = TestModel(data_object=np.random.randn(32, 1025),
-                     input_seq_length=16,
-                     batch_size=10,
-                     in_mem_data=False,
-                     splits=[1.])
-
-    assert isinstance(base.datamodule, FreqData)
-    assert isinstance(base.optim, FreqOptim)
-
-    trainer = get_trainer(root_dir=tmp_path, max_epochs=1)
-    trainer.fit(base)
-
-    assert base.datamodule.train_dataloader() is not None
-    loader = base.datamodule.train_dataloader()
-    batch = next(iter(loader))
-    assert isinstance(batch, list), type(batch)
-    assert len(batch) == 2, len(batch)
-    assert isinstance(batch[0], torch.Tensor), type(batch[0])
-    assert isinstance(batch[1], torch.Tensor), type(batch[1])
-
-    optims = base.configure_optimizers()
-    assert optims is not None
-    assert len(optims[0][0].param_groups[0]) > 0
-
-
-def test_layer_computed_properties():
-    cases = [
-        (FreqLayer(layer_index=0, ),
-         dict(shift=2, receptive_field=2, rel_shift=1)),
-        (FreqLayer(layer_index=0, strict=True),
-         dict(shift=2, receptive_field=2, rel_shift=2)),
-        (FreqLayer(layer_index=1, ),
-         dict(shift=4, receptive_field=4, rel_shift=2)),
-        (FreqLayer(layer_index=1, strict=True),
-         dict(shift=5, receptive_field=4, rel_shift=3)),
-        (FreqLayer(layer_index=2, ),
-         dict(shift=8, receptive_field=8, rel_shift=4)),
-        (FreqLayer(layer_index=2, strict=True),
-         dict(shift=10, receptive_field=8, rel_shift=5)),
-    ]
-
-    for i, (layer, expected) in enumerate(cases):
-
-        for attr in expected.keys():
-            assert getattr(layer, attr)() == expected[attr], \
-                ("case_" + str(i), attr, getattr(layer, attr)(), expected[attr])
-
-
-def test_single_layer_output_length():
-    input_length = 32
-    cases = [
-        (FreqLayer(layer_index=0, ), input_length - 1),
-        (FreqLayer(layer_index=0, strict=True), input_length - 1),
-        (FreqLayer(layer_index=1, ), input_length - 2),
-        (FreqLayer(layer_index=2, strict=True), input_length - 4),
-        (FreqLayer(layer_index=0, concat_outputs="left"), input_length),
-        # strict && concat always ADDS a time-step (concat appends MORE past)
-        (FreqLayer(layer_index=1, strict=True, concat_outputs="left"), input_length + 1),
-        (FreqLayer(layer_index=1, pad_input="left"), input_length),
-        # strict && pad returns same shape : NOT LIKE CONCAT !
-        (FreqLayer(layer_index=2, strict=True, pad_input="left"), input_length),
-    ]
-    for i, (layer, expected) in enumerate(cases):
-        assert layer.output_length(input_length) == expected, \
-            ("computed case_" + str(i), layer.output_length(input_length), expected)
-
-        inpt = torch.randn(1, layer.layer_dim, input_length)
-        outpt, _ = layer(inpt)
-        assert outpt.size(-1) == expected, \
-            ("forward case_" + str(i), outpt.size(-1), expected, layer.padding)
-
-
-def test_freqnet_computed_properties():
-    input_length = 16
-    data_object = torch.randn(1, input_length, 1025)
-
-    cases = [
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(1, 1, 1),
-                 ),
-         dict(shift=4,
-              receptive_field=4,
-              all_rel_shifts=(1, 1, 1),
-              all_shifts=(2, 3, 4),
-              targets=[(4, input_length-3)],
-              out_length=input_length-3,
-              )),
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(1, 1, 1),
-                 strict=True),
-         dict(shift=6,
-              receptive_field=4,
-              all_rel_shifts=(2, 2, 2),
-              all_shifts=(2, 4, 6),
-              targets=[(6, input_length-3)],
-              out_length=input_length-3)),
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(3,),
-                 ),
-         dict(shift=8,
-              receptive_field=8,
-              all_rel_shifts=(1, 2, 4),
-              all_shifts=(2, 4, 8),
-              targets=[(8, input_length-7)],
-              out_length=input_length-7)),
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(3,),
-                 strict=True),
-         dict(shift=8+2,  # strict shifts (n_layers -1) MORE than not strict
-              receptive_field=8,
-              all_rel_shifts=(2, 3, 5),
-              all_shifts=(2, 5, 10),
-              targets=[(8+2, input_length-7)],
-              out_length=input_length-7)),
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(3, 3),
-                 ),
-         dict(shift=15,
-              receptive_field=15,
-              all_rel_shifts=(1, 2, 4, 1, 2, 4),
-              all_shifts=(2, 4, 8, 9, 11, 15),
-              targets=[(15, input_length-14)],
-              out_length=input_length-14)),
-        (FreqNet(data_object=data_object, input_seq_length=input_length,
-                 n_layers=(3, 3),
-                 strict=True),
-         dict(shift=20,
-              receptive_field=15,
-              all_rel_shifts=(2, 3, 5, 2, 3, 5),
-              all_shifts=(2, 5, 10, 12, 15, 20),
-              targets=[(20, input_length-14)],
-              out_length=input_length-14)),
-    ]
-
-    for i, (net, expected) in enumerate(cases):
-
-        for attr in expected.keys():
-            if attr not in ("out_length", "targets"):
-                assert getattr(net, attr)() == expected[attr], \
-                    ("case_" + str(i), attr, getattr(net, attr)(), expected[attr])
-            if attr == "out_length":
-                # verify that computed and returned lengths are both equal to the one that is expected
-                # computed
-                assert net.output_length(input_length) == expected[attr], \
-                    ("computed case_" + str(i), net.output_length(input_length), expected[attr])
-                # returned
-                outpt = net(data_object)
-                assert outpt.size(1) == expected[attr], \
-                    ("forward case_" + str(i), outpt.size(1), expected[attr])
-
-            if attr == "targets":
-                result = net.targets_shifts_and_lengths(input_length)
-                assert result == expected[attr], \
-                    (attr, "case_" + str(i), result, expected[attr])
-
-
-def test_models_train(tmp_path):
-
-    data_object = torch.randn(32, 1025)
-    kwargs = dict(splits=None, input_seq_length=8)
-    models = [
-
-        FreqNet(data_object=data_object, **kwargs),
-        FreaksNet(data_object=data_object, **kwargs),
-        HKFreqNet(data_object=data_object, n_layers=2, **kwargs),
-        FrexpertMixture(data_object=data_object, **kwargs),
-        LayerWiseLossFreqNet(data_object=data_object, **kwargs)
-    ]
-
-    for mdl in models:
-        print("CLASS : ", mdl.__class__.__name__,
-              mdl.targets_shifts_and_lengths(kwargs["input_seq_length"]))
-        trainer = get_trainer(root_dir=tmp_path, max_epochs=1)
-
-        trainer.fit(mdl)
-        # print("FREQSHAPES:", mdl.targets_shifts_and_lengths(8), mdl(data_object[:8].unsqueeze(0)).shape)
-        assert trainer.global_step > 0, type(mdl)
\ No newline at end of file
diff --git a/tests/test_kit/test_ds_wrappers.py b/tests/test_kit/test_ds_wrappers.py
deleted file mode 100644
index 893ab90..0000000
--- a/tests/test_kit/test_ds_wrappers.py
+++ /dev/null
@@ -1,60 +0,0 @@
-import numpy as np
-
-from mimikit.kit import DataObject, DSWrapper, InputEqualTarget, ShiftedSeqsPair
-
-
-class TestBaseWrapper:
-    data = np.random.randn(10, 10)
-    ds = DataObject(data)
-    wrapper = DSWrapper()
-
-    def test_wraps_without_error(self):
-        ds = self.wrapper(self.ds)
-        assert ds.data is self.data
-
-    def test_wraps_returns_new_copy(self):
-        ds = self.wrapper(self.ds)
-        assert self.ds.__class__ is not ds.__class__
-
-
-class TestInputEqualTarget:
-    data = np.random.randn(10, 10)
-    ds = DataObject(data)
-    wrapper = InputEqualTarget()
-
-    def test_wraps_without_error(self):
-        ds = self.wrapper(self.ds)
-        assert ds.data is self.data
-
-    def test_overrides_getitem(self):
-        ds = self.wrapper(self.ds)
-        returned = ds[0]
-        assert len(returned) == 2, returned
-        assert np.all(returned[0] == returned[1])
-
-
-class TestShiftedSeqsPair:
-    data = np.random.randn(10, 4)
-    ds = DataObject(data)
-    wrapper = ShiftedSeqsPair(input_length=2, targets=[(1, 3)])  # ta
-
-    def test_wraps_without_error(self):
-        ds = self.wrapper(self.ds)
-        assert ds.data is self.data
-
-    def test_overrides_getitem(self):
-        ds = self.wrapper(self.ds)
-        returned = ds[0]
-        # returns input and target
-        assert len(returned) == 2, returned
-        # returns correct input shape
-        assert returned[0].shape == (self.wrapper.lengths[0], self.data.shape[-1]), returned[0].shape
-        # returns correct target shape
-        assert returned[1].shape == (self.wrapper.lengths[1], self.data.shape[-1]), returned[1].shape
-        # shifts sequences correctly
-        assert np.all(returned[0][self.wrapper.shifts[1]] == returned[1][0]), (returned[0][self.wrapper.shifts[1]],
-                                                                            returned[1][0])
-
-    def test_overrides_len(self):
-        ds = self.wrapper(self.ds)
-        assert len(ds) < len(self.ds), (len(ds), len(self.ds))
diff --git a/tests/test_kit/test_get_trainer.py b/tests/test_kit/test_get_trainer.py
index cf7ce00..b6f2e2a 100644
--- a/tests/test_kit/test_get_trainer.py
+++ b/tests/test_kit/test_get_trainer.py
@@ -6,9 +6,8 @@
 import os
 import shutil
 
-from mimikit.kit.get_trainer import get_trainer
-from mimikit.kit.loggers import MMKDefaultLogger
-from mimikit.kit.checkpoint import MMKCheckpoint
+from mimikit.get_trainer import get_trainer
+from mimikit.models.parts import MMKDefaultLogger, MMKCheckpoint
 from mimikit.connectors.neptune import NeptuneConnector
 
 
diff --git a/tests/test_kit/test_loggers_and_checkpoints.py b/tests/test_kit/test_loggers_and_checkpoints.py
index 9e6203a..0367781 100644
--- a/tests/test_kit/test_loggers_and_checkpoints.py
+++ b/tests/test_kit/test_loggers_and_checkpoints.py
@@ -4,7 +4,8 @@
 import torch
 import torch.nn as nn
 
-from mimikit.kit import MMKHooks, LoggingHooks, get_trainer
+from mimikit.models.parts import MMKHooks, LoggingHooks
+from mimikit.get_trainer import get_trainer
 
 
 class TestModel(MMKHooks,
@@ -12,7 +13,9 @@ class TestModel(MMKHooks,
                 pl.LightningModule):
 
     def __init__(self, model_dim=10):
-        super(TestModel, self).__init__()
+        super(pl.LightningModule, self).__init__()
+        MMKHooks.__init__(self)
+        LoggingHooks.__init__(self)
         self.fc = nn.Linear(model_dim, model_dim)
         self.save_hyperparameters()
 
@@ -26,6 +29,13 @@ def training_step(self, batch, batch_idx):
         self.log("recon", L, on_step=False, on_epoch=True)
         return {"loss": L, "recon_loss": L}
 
+    def validation_step(self, batch, batch_idx):
+        inpt, target = batch
+        pred = self.forward(inpt)
+        L = nn.MSELoss()(pred, target)
+        self.log("recon", L, on_step=False, on_epoch=True)
+        return {"loss": L, "recon_loss": L}
+
     def configure_optimizers(self):
         return torch.optim.Adam(self.parameters())
 
@@ -53,7 +63,7 @@ def test_model(tmp_path):
 
     states = os.listdir(os.path.join(root, "states"))
     assert 'last_optim_state.ckpt' in states, states
-    assert 'epoch=0.ckpt' in states, states
+    assert 'epoch=1.ckpt' in states, states
 
     ckpt_path = str(root / "states" / 'epoch=1.ckpt')
     ckpt = torch.load(ckpt_path)
@@ -67,7 +77,7 @@ def test_model(tmp_path):
 
     from_ckpt = TestModel.load_from_checkpoint(ckpt_path)
     assert isinstance(from_ckpt, TestModel)
-    new_trainer = pl.Trainer(resume_from_checkpoint=ckpt_path, max_epochs=4)
+    new_trainer = get_trainer(resume_from_checkpoint=ckpt_path, max_epochs=4)
     print("SECOND RUN")
     new_trainer.fit(from_ckpt)
     assert new_trainer.current_epoch == 3
@@ -82,4 +92,3 @@ def test_model(tmp_path):
     assert new_trainer.global_step > trainer.global_step, (new_trainer.global_step, trainer.global_step)
 
     shutil.rmtree(root)
-    shutil.rmtree("./lightning_logs/")