File size: 8,550 Bytes
d72b2c3 54adc39 d72b2c3 54adc39 d9889a1 d72b2c3 0a8807e e366cd5 86b9ce4 fe62fb4 d912185 3dfbf53 d72b2c3 54adc39 d912185 d72b2c3 54adc39 0a8807e 54adc39 86b9ce4 dcfe0d4 df63ff0 3ac9f34 54adc39 d72b2c3 54adc39 d72b2c3 54adc39 d72b2c3 54adc39 86b9ce4 dcfe0d4 54adc39 dcfe0d4 54adc39 86b9ce4 54adc39 86b9ce4 54adc39 e366cd5 54adc39 86b9ce4 54adc39 a84b206 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 |
import typing as tp
import omegaconf
from torch import nn
import torch
from huggingface_hub import hf_hub_download
import os
from omegaconf import OmegaConf, DictConfig
from .encodec import EncodecModel
from .lm import LMModel
from .seanet import SEANetDecoder
from .codebooks_patterns import DelayedPatternProvider
from .conditioners import T5Conditioner
from .vq import ResidualVectorQuantizer
def _delete_param(cfg: DictConfig, full_name: str):
parts = full_name.split('.')
for part in parts[:-1]:
if part in cfg:
cfg = cfg[part]
else:
return
OmegaConf.set_struct(cfg, False)
if parts[-1] in cfg:
del cfg[parts[-1]]
OmegaConf.set_struct(cfg, True)
def dict_from_config(cfg):
dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)
return dct
# ============================================== DEFINE AUDIOGEN
class AudioGen(nn.Module):
# https://huggingface.co/facebook/audiogen-medium
def __init__(self,
duration=0.024,
device='cpu'):
super().__init__()
self.device = device # needed for loading & select float16 LM
self.load_compression_model()
self.load_lm_model()
self.duration = duration
@property
def frame_rate(self):
return self.compression_model.frame_rate
def generate(self,
descriptions):
with torch.no_grad():
gen_tokens = self.lm.generate(
descriptions=descriptions,
max_gen_len=int(self.duration * self.frame_rate)) # [bs, 4, 37 * self.lm.n_draw]
x = self.compression_model.decode(gen_tokens, None) #[bs, 1, 11840]
# print('______________\nAudioGen Tokens', gen_tokens)
return x / x.abs().max(2, keepdims=True)[0] + 1e-7
# == BUILD Fn
def get_quantizer(self, quantizer, cfg, dimension):
klass = {
'no_quant': None,
'rvq': ResidualVectorQuantizer
}[quantizer]
kwargs = dict_from_config(getattr(cfg, quantizer))
if quantizer != 'no_quant':
kwargs['dimension'] = dimension
return klass(**kwargs)
def get_encodec_autoencoder(self, cfg):
kwargs = dict_from_config(getattr(cfg, 'seanet'))
_ = kwargs.pop('encoder')
decoder_override_kwargs = kwargs.pop('decoder')
decoder_kwargs = {**kwargs, **decoder_override_kwargs}
decoder = SEANetDecoder(**decoder_kwargs)
return decoder
def get_compression_model(self, cfg):
"""Instantiate a compression model."""
if cfg.compression_model == 'encodec':
kwargs = dict_from_config(getattr(cfg, 'encodec'))
quantizer_name = kwargs.pop('quantizer')
decoder = self.get_encodec_autoencoder(cfg)
quantizer = self.get_quantizer(quantizer_name, cfg, 128)
renormalize = kwargs.pop('renormalize', False)
# deprecated params
# print(f'{frame_rate=} {encoder.dimension=}') frame_rate=50 encoder.dimension=128
kwargs.pop('renorm', None)
# print('\n______!____________\n', kwargs, '\n______!____________\n')
# ______!____________
# {'autoencoder': 'seanet', 'sample_rate': 16000, 'channels': 1, 'causal': False}
# ______!____________
return EncodecModel(decoder=decoder,
quantizer=quantizer,
frame_rate=50,
renormalize=renormalize,
sample_rate=16000,
channels=1,
causal=False
).to(cfg.device)
else:
raise KeyError(f"Unexpected compression model {cfg.compression_model}")
def get_lm_model(self, cfg):
"""Instantiate a transformer LM."""
if cfg.lm_model in ['transformer_lm',
'transformer_lm_magnet']:
kwargs = dict_from_config(getattr(cfg, 'transformer_lm'))
n_q = kwargs['n_q']
q_modeling = kwargs.pop('q_modeling', None)
codebooks_pattern_cfg = getattr(cfg, 'codebooks_pattern')
attribute_dropout = dict_from_config(getattr(cfg, 'attribute_dropout'))
cls_free_guidance = dict_from_config(getattr(cfg, 'classifier_free_guidance'))
cfg_prob, cfg_coef = cls_free_guidance['training_dropout'], cls_free_guidance['inference_coef']
# if len(fuser.fuse2cond['cross']) > 0: # enforce cross-att programmatically
kwargs['cross_attention'] = True
if codebooks_pattern_cfg.modeling is None:
print('Q MODELING\n=\n=><')
assert q_modeling is not None, \
"LM model should either have a codebook pattern defined or transformer_lm.q_modeling"
codebooks_pattern_cfg = omegaconf.OmegaConf.create(
{'modeling': q_modeling, 'delay': {'delays': list(range(n_q))}}
)
pattern_provider = self.get_codebooks_pattern_provider(n_q, codebooks_pattern_cfg)
return LMModel(
pattern_provider=pattern_provider,
condition_provider=T5Conditioner(name='t5-large', output_dim=kwargs["dim"], device=self.device),
cfg_dropout=cfg_prob,
cfg_coef=cfg_coef,
attribute_dropout=attribute_dropout,
dtype=getattr(torch, cfg.dtype),
device=self.device,
**kwargs
).to(cfg.device)
else:
raise KeyError(f"Unexpected LM model {cfg.lm_model}")
def get_codebooks_pattern_provider(self, n_q, cfg):
pattern_providers = {
'delay': DelayedPatternProvider, # THIS
}
name = cfg.modeling
kwargs = dict_from_config(cfg.get(name)) if hasattr(cfg, name) else {}
klass = pattern_providers[name]
return klass(n_q, **kwargs)
# ======================
def load_compression_model(self):
file = hf_hub_download(
repo_id='facebook/audiogen-medium',
filename="compression_state_dict.bin",
cache_dir=os.environ.get('AUDIOCRAFT_CACHE_DIR', None),
library_name="audiocraft",
library_version= '1.3.0a1') # Found at __init__.py #audiocraft.__version__)
pkg = torch.load(file, map_location='cpu')
# if 'pretrained' in pkg:
# print('NO RPtrained\n=\n=\n=\n=\n=')
# return EncodecModel.get_pretrained(pkg['pretrained'], device='cpu')
cfg = OmegaConf.create(pkg['xp.cfg'])
cfg.device = 'cpu'
model = self.get_compression_model(cfg)
model.load_state_dict(pkg['best_state'], strict=False) # ckpt has also unused encoder weights
# return model
self.compression_model = model
def load_lm_model(self):
file = hf_hub_download(
repo_id='facebook/audiogen-medium',
filename="state_dict.bin",
cache_dir=os.environ.get('AUDIOCRAFT_CACHE_DIR', None),
library_name="audiocraft",
library_version= '1.3.0a1') # Found at __init__.py #audiocraft.__version__)
pkg = torch.load(file,
map_location=self.device) #'cpu')
cfg = OmegaConf.create(pkg['xp.cfg'])
# cfg.device = 'cpu'
if self.device == 'cpu':
cfg.dtype = 'float32'
else:
cfg.dtype = 'float16'
_delete_param(cfg, 'conditioners.self_wav.chroma_stem.cache_path')
_delete_param(cfg, 'conditioners.args.merge_text_conditions_p')
_delete_param(cfg, 'conditioners.args.drop_desc_p')
model = self.get_lm_model(cfg)
_best = pkg['best_state']
_best['condition_provider.output_proj.weight'] = _best.pop('condition_provider.conditioners.description.output_proj.weight')
_best['condition_provider.output_proj.bias'] = _best.pop('condition_provider.conditioners.description.output_proj.bias')
model.load_state_dict(pkg['best_state'])
model.cfg = cfg
# return model
self.lm = model.to(torch.float)
# def _flush(self):
# self.lm._flush() # already done in lm generate at end
|