speechbrain/speechbrain

Unable to use model trained from enhancement template

Zahuczky opened this issue · 1 comments

Describe the bug

I trained with my custom dataset with minimal modifications to https://github.com/speechbrain/speechbrain/tree/develop/templates/enhancement (0.5.16 version)

After, I end up with a result folder structure such as

results
└───4234
    │   env.log
    │   hyperparams.yaml
    │   log.txt
    │   train.py
    │   train_log.txt
    │
    └───save
        └───CKPT+2024-02-15+17-49-49+00
                brain.ckpt
                CKPT.yaml
                counter.ckpt
                dataloader-TRAIN.ckpt
                model.ckpt
                optimizer.ckpt

How could I load/use the model I trained for enhancement in my python script, like

model.enhance_file("example.wav", "output.wav")

Expected behaviour

Being able to use it like

model.enhance_file("example.wav", "output.wav")

To Reproduce

No response

Environment Details

No response

Relevant Log Output

No response

Additional Context

No response

Currently, the process is manual. Typically, you would:

  • Look up a relevant interface in speechbrain/inference (as of 1.0), e.g. WaveformEnhancement
  • Check what hparams/mods it attempts to use
  • Copy over the relevant files in a directory (e.g. here likely just model.ckpt)
  • Build an inference hyperparams.yaml in the same directory from your training hyperparams (optionally taking inspiration from a similar published model on HF)
  • Finally, write and use a test script that uses the inference interface e.g. enhancement = WaveformEnhancement.from_hparams("./mymodeldirectory/", run_opts={"device": device})