|
|
|
|
|
|
|
|
|
|
|
|
|
""" |
|
Main model for using AudioGen. This will combine all the required components |
|
and provide easy access to the generation API. |
|
""" |
|
|
|
import typing as tp |
|
|
|
import torch |
|
|
|
from .encodec import CompressionModel |
|
from .genmodel import BaseGenModel |
|
from .lm import LMModel |
|
from .builders import get_debug_compression_model, get_debug_lm_model |
|
from .loaders import load_compression_model, load_lm_model |
|
|
|
|
|
class AudioGen(BaseGenModel): |
|
"""AudioGen main model with convenient generation API. |
|
|
|
Args: |
|
name (str): name of the model. |
|
compression_model (CompressionModel): Compression model |
|
used to map audio to invertible discrete representations. |
|
lm (LMModel): Language model over discrete representations. |
|
max_duration (float, optional): maximum duration the model can produce, |
|
otherwise, inferred from the training params. |
|
""" |
|
def __init__(self, name: str, compression_model: CompressionModel, lm: LMModel, |
|
max_duration: tp.Optional[float] = None): |
|
super().__init__(name, compression_model, lm, max_duration) |
|
self.set_generation_params(duration=5) |
|
|
|
@staticmethod |
|
def get_pretrained(name: str = 'facebook/audiogen-medium', device=None): |
|
"""Return pretrained model, we provide a single model for now: |
|
- facebook/audiogen-medium (1.5B), text to sound, |
|
# see: https://huggingface.co/facebook/audiogen-medium |
|
""" |
|
if device is None: |
|
if torch.cuda.device_count(): |
|
device = 'cuda' |
|
else: |
|
device = 'cpu' |
|
|
|
if name == 'debug': |
|
|
|
compression_model = get_debug_compression_model(device, sample_rate=16000) |
|
lm = get_debug_lm_model(device) |
|
return AudioGen(name, compression_model, lm, max_duration=10) |
|
|
|
compression_model = load_compression_model(name, device=device) |
|
lm = load_lm_model(name, device=device) |
|
assert 'self_wav' not in lm.condition_provider.conditioners, \ |
|
"AudioGen do not support waveform conditioning for now" |
|
return AudioGen(name, compression_model, lm) |
|
|
|
def set_generation_params(self, use_sampling: bool = True, top_k: int = 250, |
|
top_p: float = 0.0, temperature: float = 1.0, |
|
duration: float = 10.0, cfg_coef: float = 3.0, |
|
two_step_cfg: bool = False, extend_stride: float = 2): |
|
"""Set the generation parameters for AudioGen. |
|
|
|
Args: |
|
use_sampling (bool, optional): Use sampling if True, else do argmax decoding. Defaults to True. |
|
top_k (int, optional): top_k used for sampling. Defaults to 250. |
|
top_p (float, optional): top_p used for sampling, when set to 0 top_k is used. Defaults to 0.0. |
|
temperature (float, optional): Softmax temperature parameter. Defaults to 1.0. |
|
duration (float, optional): Duration of the generated waveform. Defaults to 10.0. |
|
cfg_coef (float, optional): Coefficient used for classifier free guidance. Defaults to 3.0. |
|
two_step_cfg (bool, optional): If True, performs 2 forward for Classifier Free Guidance, |
|
instead of batching together the two. This has some impact on how things |
|
are padded but seems to have little impact in practice. |
|
extend_stride: when doing extended generation (i.e. more than 10 seconds), by how much |
|
should we extend the audio each time. Larger values will mean less context is |
|
preserved, and shorter value will require extra computations. |
|
""" |
|
assert extend_stride < self.max_duration, "Cannot stride by more than max generation duration." |
|
self.extend_stride = extend_stride |
|
self.duration = duration |
|
self.generation_params = { |
|
'use_sampling': use_sampling, |
|
'temp': temperature, |
|
'top_k': top_k, |
|
'top_p': top_p, |
|
'cfg_coef': cfg_coef, |
|
'two_step_cfg': two_step_cfg, |
|
} |
|
|
