mirror of
https://github.com/YuzuZensai/spleeter.git
synced 2026-01-31 14:58:23 +00:00
🐛 fix backend resolution
This commit is contained in:
@@ -12,6 +12,11 @@
|
|||||||
|
|
||||||
from enum import Enum
|
from enum import Enum
|
||||||
|
|
||||||
|
# pyright: reportMissingImports=false
|
||||||
|
# pylint: disable=import-error
|
||||||
|
import tensorflow as tf
|
||||||
|
# pylint: enable=import-error
|
||||||
|
|
||||||
__email__ = 'spleeter@deezer.com'
|
__email__ = 'spleeter@deezer.com'
|
||||||
__author__ = 'Deezer Research'
|
__author__ = 'Deezer Research'
|
||||||
__license__ = 'MIT License'
|
__license__ = 'MIT License'
|
||||||
@@ -34,3 +39,12 @@ class STFTBackend(str, Enum):
|
|||||||
AUTO: str = 'auto'
|
AUTO: str = 'auto'
|
||||||
TENSORFLOW: str = 'tensorflow'
|
TENSORFLOW: str = 'tensorflow'
|
||||||
LIBROSA: str = 'librosa'
|
LIBROSA: str = 'librosa'
|
||||||
|
|
||||||
|
def resolve(cls: type, backend: str) -> str:
|
||||||
|
if backend not in cls.__members__.items():
|
||||||
|
raise ValueError(f'Unsupported backend {backend}')
|
||||||
|
if backend == cls.AUTO:
|
||||||
|
if len(tf.config.list_physical_devices('GPU')):
|
||||||
|
return cls.TENSORFLOW
|
||||||
|
return STFTBackend.LIBROSA
|
||||||
|
return backend
|
||||||
|
|||||||
@@ -129,8 +129,7 @@ class Separator(object):
|
|||||||
else:
|
else:
|
||||||
self._pool = None
|
self._pool = None
|
||||||
self._tasks = []
|
self._tasks = []
|
||||||
# NOTE: provide type check here ?
|
self._params['stft_backend'] = STFTBackend.resolve(stft_backend)
|
||||||
self._params['stft_backend'] = stft_backend
|
|
||||||
self._data_generator = DataGenerator()
|
self._data_generator = DataGenerator()
|
||||||
|
|
||||||
def __del__(self) -> None:
|
def __del__(self) -> None:
|
||||||
@@ -333,11 +332,6 @@ class Separator(object):
|
|||||||
(Optional) string describing the waveform (e.g. filename).
|
(Optional) string describing the waveform (e.g. filename).
|
||||||
"""
|
"""
|
||||||
backend: str = self._params['stft_backend']
|
backend: str = self._params['stft_backend']
|
||||||
if backend == STFTBackend.AUTO:
|
|
||||||
if len(tf.config.list_physical_devices('GPU')):
|
|
||||||
backend = STFTBackend.TENSORFLOW
|
|
||||||
else:
|
|
||||||
backend = STFTBackend.LIBROSA
|
|
||||||
if backend == STFTBackend.TENSORFLOW:
|
if backend == STFTBackend.TENSORFLOW:
|
||||||
return self._separate_tensorflow(waveform, audio_descriptor)
|
return self._separate_tensorflow(waveform, audio_descriptor)
|
||||||
elif backend == STFTBackend.LIBROSA:
|
elif backend == STFTBackend.LIBROSA:
|
||||||
|
|||||||
Reference in New Issue
Block a user