unwrap nn~ models from neutone models

unwrap.neutone.py

from pathlib import Path

from fire import Fire
import torch

def main(model:str, outfile:str=None):
    """
    Args:
        model: path to neutone model file
        outfile: path to write nn~ model file
    """
    # load the neutone model
    m = torch.jit.load(model)
    if not hasattr(m, 'model'):
        print('not recognized as a neutone model')
        return
      
    # unwrap the nn~ model from the neutone model
    m = m.model
    
    if not hasattr(m, 'latent_pca'):
        print('warning: this does not appear to be a RAVE model')
    if outfile is None:
        old_path = Path(model)
        new_path = old_path.with_stem(f'{old_path.stem}_nntilde')
        outfile = str(new_path)
        
    # save again as nn~ / NN.ar model
    torch.jit.save(m, outfile)

if __name__=='__main__':
    Fire(main)