Spaces:
Sleeping
Sleeping
# Copyright (c) Meta Platforms, Inc. and affiliates. | |
# All rights reserved. | |
# | |
# This source code is licensed under the license found in the | |
# LICENSE file in the root directory of this source tree. | |
from pathlib import Path | |
import time | |
import typing as tp | |
import flashy | |
import math | |
import omegaconf | |
import torch | |
from torch.nn import functional as F | |
from . import base, builders | |
from .compression import CompressionSolver | |
from .. import metrics as eval_metrics | |
from .. import models | |
from ..data.audio_dataset import AudioDataset | |
from ..data.music_dataset import MusicDataset, MusicInfo, AudioInfo | |
from ..data.audio_utils import normalize_audio | |
from ..modules.conditioners import JointEmbedCondition, SegmentWithAttributes, WavCondition | |
from ..utils.cache import CachedBatchWriter, CachedBatchLoader | |
from ..utils.samples.manager import SampleManager | |
from ..utils.utils import get_dataset_from_loader, is_jsonable, warn_once | |
class MusicGenSolver(base.StandardSolver): | |
"""Solver for MusicGen training task. | |
Used in: https://arxiv.org/abs/2306.05284 | |
""" | |
DATASET_TYPE: builders.DatasetType = builders.DatasetType.MUSIC | |
def __init__(self, cfg: omegaconf.DictConfig): | |
super().__init__(cfg) | |
# easier access to sampling parameters | |
self.generation_params = { | |
'use_sampling': self.cfg.generate.lm.use_sampling, | |
'temp': self.cfg.generate.lm.temp, | |
'top_k': self.cfg.generate.lm.top_k, | |
'top_p': self.cfg.generate.lm.top_p, | |
} | |
self._best_metric_name: tp.Optional[str] = 'ce' | |
self._cached_batch_writer = None | |
self._cached_batch_loader = None | |
if cfg.cache.path: | |
if cfg.cache.write: | |
self._cached_batch_writer = CachedBatchWriter(Path(cfg.cache.path)) | |
if self.cfg.cache.write_num_shards: | |
self.logger.warning("Multiple shard cache, best_metric_name will be set to None.") | |
self._best_metric_name = None | |
else: | |
self._cached_batch_loader = CachedBatchLoader( | |
Path(cfg.cache.path), cfg.dataset.batch_size, cfg.dataset.num_workers, | |
min_length=self.cfg.optim.updates_per_epoch or 1) | |
self.dataloaders['original_train'] = self.dataloaders['train'] | |
self.dataloaders['train'] = self._cached_batch_loader # type: ignore | |
def get_eval_solver_from_sig(sig: str, dtype: tp.Optional[str] = None, | |
device: tp.Optional[str] = None, autocast: bool = True, | |
batch_size: tp.Optional[int] = None, | |
override_cfg: tp.Optional[tp.Union[dict, omegaconf.DictConfig]] = None, | |
**kwargs): | |
"""Mostly a convenience function around magma.train.get_solver_from_sig, | |
populating all the proper param, deactivating EMA, FSDP, loading the best state, | |
basically all you need to get a solver ready to "play" with in single GPU mode | |
and with minimal memory overhead. | |
Args: | |
sig (str): signature to load. | |
dtype (str or None): potential dtype, as a string, i.e. 'float16'. | |
device (str or None): potential device, as a string, i.e. 'cuda'. | |
override_cfg (dict or omegaconf.DictConfig or None): potential device, as a string, i.e. 'cuda'. | |
""" | |
from audiocraft import train | |
our_override_cfg: tp.Dict[str, tp.Any] = {'optim': {'ema': {'use': False}}} | |
our_override_cfg['autocast'] = autocast | |
if dtype is not None: | |
our_override_cfg['dtype'] = dtype | |
if device is not None: | |
our_override_cfg['device'] = device | |
if batch_size is not None: | |
our_override_cfg['dataset'] = {'batch_size': batch_size} | |
if override_cfg is None: | |
override_cfg = {} | |
override_cfg = omegaconf.OmegaConf.merge( | |
omegaconf.DictConfig(override_cfg), omegaconf.DictConfig(our_override_cfg)) # type: ignore | |
solver = train.get_solver_from_sig( | |
sig, override_cfg=override_cfg, | |
load_best=True, disable_fsdp=True, | |
ignore_state_keys=['optimizer', 'ema'], **kwargs) | |
solver.model.eval() | |
return solver | |
def get_formatter(self, stage_name: str) -> flashy.Formatter: | |
return flashy.Formatter({ | |
'lr': '.2E', | |
'ce': '.3f', | |
'ppl': '.3f', | |
'grad_norm': '.3E', | |
}, exclude_keys=['ce_q*', 'ppl_q*']) | |
def best_metric_name(self) -> tp.Optional[str]: | |
return self._best_metric_name | |
def build_model(self) -> None: | |
"""Instantiate models and optimizer.""" | |
# we can potentially not use all quantizers with which the EnCodec model was trained | |
# (e.g. we trained the model with quantizers dropout) | |
self.compression_model = CompressionSolver.wrapped_model_from_checkpoint( | |
self.cfg, self.cfg.compression_model_checkpoint, device=self.device) | |
assert self.compression_model.sample_rate == self.cfg.sample_rate, ( | |
f"Compression model sample rate is {self.compression_model.sample_rate} but " | |
f"Solver sample rate is {self.cfg.sample_rate}." | |
) | |
# ensure we have matching configuration between LM and compression model | |
assert self.cfg.transformer_lm.card == self.compression_model.cardinality, ( | |
"Cardinalities of the LM and compression model don't match: ", | |
f"LM cardinality is {self.cfg.transformer_lm.card} vs ", | |
f"compression model cardinality is {self.compression_model.cardinality}" | |
) | |
assert self.cfg.transformer_lm.n_q == self.compression_model.num_codebooks, ( | |
"Numbers of codebooks of the LM and compression models don't match: ", | |
f"LM number of codebooks is {self.cfg.transformer_lm.n_q} vs ", | |
f"compression model numer of codebooks is {self.compression_model.num_codebooks}" | |
) | |
self.logger.info("Compression model has %d codebooks with %d cardinality, and a framerate of %d", | |
self.compression_model.num_codebooks, self.compression_model.cardinality, | |
self.compression_model.frame_rate) | |
# instantiate LM model | |
self.model: models.LMModel = models.builders.get_lm_model(self.cfg).to(self.device) | |
if self.cfg.fsdp.use: | |
assert not self.cfg.autocast, "Cannot use autocast with fsdp" | |
self.model = self.wrap_with_fsdp(self.model) | |
self.register_ema('model') | |
# initialize optimization | |
self.optimizer = builders.get_optimizer(builders.get_optim_parameter_groups(self.model), self.cfg.optim) | |
self.lr_scheduler = builders.get_lr_scheduler(self.optimizer, self.cfg.schedule, self.total_updates) | |
self.register_stateful('compression_model', 'model', 'optimizer', 'lr_scheduler') | |
self.register_best_state('model') | |
self.autocast_dtype = { | |
'float16': torch.float16, 'bfloat16': torch.bfloat16 | |
}[self.cfg.autocast_dtype] | |
self.scaler: tp.Optional[torch.cuda.amp.GradScaler] = None | |
if self.cfg.fsdp.use: | |
need_scaler = self.cfg.fsdp.param_dtype == 'float16' | |
else: | |
need_scaler = self.cfg.autocast and self.autocast_dtype is torch.float16 | |
if need_scaler: | |
if self.cfg.fsdp.use: | |
from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler | |
self.scaler = ShardedGradScaler() # type: ignore | |
else: | |
self.scaler = torch.cuda.amp.GradScaler() | |
self.register_stateful('scaler') | |
def build_dataloaders(self) -> None: | |
"""Instantiate audio dataloaders for each stage.""" | |
self.dataloaders = builders.get_audio_datasets(self.cfg, dataset_type=self.DATASET_TYPE) | |
def show(self) -> None: | |
"""Show the compression model and LM model.""" | |
self.logger.info("Compression model:") | |
self.log_model_summary(self.compression_model) | |
self.logger.info("LM model:") | |
self.log_model_summary(self.model) | |
def load_state_dict(self, state: dict) -> None: | |
if 'condition_provider' in state: | |
model_state = state['model'] | |
condition_provider_state = state.pop('condition_provider') | |
prefix = 'condition_provider.' | |
for key, value in condition_provider_state.items(): | |
key = prefix + key | |
assert key not in model_state | |
model_state[key] = value | |
super().load_state_dict(state) | |
def load_from_pretrained(self, name: str): | |
# TODO: support native HF versions of MusicGen. | |
lm_pkg = models.loaders.load_lm_model_ckpt(name) | |
state: dict = { | |
'best_state': { | |
'model': lm_pkg['best_state'], | |
}, | |
} | |
return state | |
def _compute_cross_entropy( | |
self, logits: torch.Tensor, targets: torch.Tensor, mask: torch.Tensor | |
) -> tp.Tuple[torch.Tensor, tp.List[torch.Tensor]]: | |
"""Compute cross entropy between multi-codebook targets and model's logits. | |
The cross entropy is computed per codebook to provide codebook-level cross entropy. | |
Valid timesteps for each of the codebook are pulled from the mask, where invalid | |
timesteps are set to 0. | |
Args: | |
logits (torch.Tensor): Model's logits of shape [B, K, T, card]. | |
targets (torch.Tensor): Target codes, of shape [B, K, T]. | |
mask (torch.Tensor): Mask for valid target codes, of shape [B, K, T]. | |
Returns: | |
ce (torch.Tensor): Cross entropy averaged over the codebooks | |
ce_per_codebook (list of torch.Tensor): Cross entropy per codebook (detached). | |
""" | |
B, K, T = targets.shape | |
assert logits.shape[:-1] == targets.shape | |
assert mask.shape == targets.shape | |
ce = torch.zeros([], device=targets.device) | |
ce_per_codebook: tp.List[torch.Tensor] = [] | |
for k in range(K): | |
logits_k = logits[:, k, ...].contiguous().view(-1, logits.size(-1)) # [B x T, card] | |
targets_k = targets[:, k, ...].contiguous().view(-1) # [B x T] | |
mask_k = mask[:, k, ...].contiguous().view(-1) # [B x T] | |
ce_targets = targets_k[mask_k] | |
ce_logits = logits_k[mask_k] | |
q_ce = F.cross_entropy(ce_logits, ce_targets) | |
ce += q_ce | |
ce_per_codebook.append(q_ce.detach()) | |
# average cross entropy across codebooks | |
ce = ce / K | |
return ce, ce_per_codebook | |
def _prepare_tokens_and_attributes( | |
self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], | |
check_synchronization_points: bool = False | |
) -> tp.Tuple[dict, torch.Tensor, torch.Tensor]: | |
"""Prepare input batchs for language model training. | |
Args: | |
batch (tuple[torch.Tensor, list[SegmentWithAttributes]]): Input batch with audio tensor of shape [B, C, T] | |
and corresponding metadata as SegmentWithAttributes (with B items). | |
check_synchronization_points (bool): Whether to check for synchronization points slowing down training. | |
Returns: | |
Condition tensors (dict[str, any]): Preprocessed condition attributes. | |
Tokens (torch.Tensor): Audio tokens from compression model, of shape [B, K, T_s], | |
with B the batch size, K the number of codebooks, T_s the token timesteps. | |
Padding mask (torch.Tensor): Mask with valid positions in the tokens tensor, of shape [B, K, T_s]. | |
""" | |
if self._cached_batch_loader is None or self.current_stage != "train": | |
audio, infos = batch | |
audio = audio.to(self.device) | |
audio_tokens = None | |
assert audio.size(0) == len(infos), ( | |
f"Mismatch between number of items in audio batch ({audio.size(0)})", | |
f" and in metadata ({len(infos)})" | |
) | |
else: | |
audio = None | |
# In that case the batch will be a tuple coming from the _cached_batch_writer bit below. | |
infos, = batch # type: ignore | |
assert all([isinstance(info, AudioInfo) for info in infos]) | |
assert all([info.audio_tokens is not None for info in infos]) # type: ignore | |
audio_tokens = torch.stack([info.audio_tokens for info in infos]).to(self.device) # type: ignore | |
audio_tokens = audio_tokens.long() | |
for info in infos: | |
if isinstance(info, MusicInfo): | |
# Careful here, if you want to use this condition_wav (e.b. chroma conditioning), | |
# then you must be using the chroma cache! otherwise the code will try | |
# to use this segment and fail (by that I mean you will see NaN everywhere). | |
info.self_wav = WavCondition( | |
torch.full([1, info.channels, info.total_frames], float('NaN')), | |
length=torch.tensor([info.n_frames]), | |
sample_rate=[info.sample_rate], | |
path=[info.meta.path], | |
seek_time=[info.seek_time]) | |
dataset = get_dataset_from_loader(self.dataloaders['original_train']) | |
assert isinstance(dataset, MusicDataset), type(dataset) | |
if dataset.paraphraser is not None and info.description is not None: | |
# Hackingly reapplying paraphraser when using cache. | |
info.description = dataset.paraphraser.sample_paraphrase( | |
info.meta.path, info.description) | |
# prepare attributes | |
attributes = [info.to_condition_attributes() for info in infos] | |
attributes = self.model.cfg_dropout(attributes) | |
attributes = self.model.att_dropout(attributes) | |
tokenized = self.model.condition_provider.tokenize(attributes) | |
# Now we should be synchronization free. | |
if self.device == "cuda" and check_synchronization_points: | |
torch.cuda.set_sync_debug_mode("warn") | |
if audio_tokens is None: | |
with torch.no_grad(): | |
audio_tokens, scale = self.compression_model.encode(audio) | |
assert scale is None, "Scaled compression model not supported with LM." | |
with self.autocast: | |
condition_tensors = self.model.condition_provider(tokenized) | |
# create a padding mask to hold valid vs invalid positions | |
padding_mask = torch.ones_like(audio_tokens, dtype=torch.bool, device=audio_tokens.device) | |
# replace encodec tokens from padded audio with special_token_id | |
if self.cfg.tokens.padding_with_special_token: | |
audio_tokens = audio_tokens.clone() | |
padding_mask = padding_mask.clone() | |
token_sample_rate = self.compression_model.frame_rate | |
B, K, T_s = audio_tokens.shape | |
for i in range(B): | |
n_samples = infos[i].n_frames | |
audio_sample_rate = infos[i].sample_rate | |
# take the last token generated from actual audio frames (non-padded audio) | |
valid_tokens = math.floor(float(n_samples) / audio_sample_rate * token_sample_rate) | |
audio_tokens[i, :, valid_tokens:] = self.model.special_token_id | |
padding_mask[i, :, valid_tokens:] = 0 | |
if self.device == "cuda" and check_synchronization_points: | |
torch.cuda.set_sync_debug_mode("default") | |
if self._cached_batch_writer is not None and self.current_stage == 'train': | |
assert self._cached_batch_loader is None | |
assert audio_tokens is not None | |
for info, one_audio_tokens in zip(infos, audio_tokens): | |
assert isinstance(info, AudioInfo) | |
if isinstance(info, MusicInfo): | |
assert not info.joint_embed, "joint_embed and cache not supported yet." | |
info.self_wav = None | |
assert one_audio_tokens.max() < 2**15, one_audio_tokens.max().item() | |
info.audio_tokens = one_audio_tokens.short().cpu() | |
self._cached_batch_writer.save(infos) | |
return condition_tensors, audio_tokens, padding_mask | |
def run_step(self, idx: int, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], metrics: dict) -> dict: | |
"""Perform one training or valid step on a given batch.""" | |
check_synchronization_points = idx == 1 and self.device == 'cuda' | |
condition_tensors, audio_tokens, padding_mask = self._prepare_tokens_and_attributes( | |
batch, check_synchronization_points) | |
self.deadlock_detect.update('tokens_and_conditions') | |
if check_synchronization_points: | |
torch.cuda.set_sync_debug_mode('warn') | |
with self.autocast: | |
model_output = self.model.compute_predictions(audio_tokens, [], condition_tensors) # type: ignore | |
logits = model_output.logits | |
mask = padding_mask & model_output.mask | |
ce, ce_per_codebook = self._compute_cross_entropy(logits, audio_tokens, mask) | |
loss = ce | |
self.deadlock_detect.update('loss') | |
if check_synchronization_points: | |
torch.cuda.set_sync_debug_mode('default') | |
if self.is_training: | |
metrics['lr'] = self.optimizer.param_groups[0]['lr'] | |
if self.scaler is not None: | |
loss = self.scaler.scale(loss) | |
self.deadlock_detect.update('scale') | |
if self.cfg.fsdp.use: | |
loss.backward() | |
flashy.distrib.average_tensors(self.model.buffers()) | |
elif self.cfg.optim.eager_sync: | |
with flashy.distrib.eager_sync_model(self.model): | |
loss.backward() | |
else: | |
# this should always be slower but can be useful | |
# for weird use cases like multiple backwards. | |
loss.backward() | |
flashy.distrib.sync_model(self.model) | |
self.deadlock_detect.update('backward') | |
if self.scaler is not None: | |
self.scaler.unscale_(self.optimizer) | |
if self.cfg.optim.max_norm: | |
if self.cfg.fsdp.use: | |
metrics['grad_norm'] = self.model.clip_grad_norm_(self.cfg.optim.max_norm) # type: ignore | |
else: | |
metrics['grad_norm'] = torch.nn.utils.clip_grad_norm_( | |
self.model.parameters(), self.cfg.optim.max_norm | |
) | |
if self.scaler is None: | |
self.optimizer.step() | |
else: | |
self.scaler.step(self.optimizer) | |
self.scaler.update() | |
if self.lr_scheduler: | |
self.lr_scheduler.step() | |
self.optimizer.zero_grad() | |
self.deadlock_detect.update('optim') | |
if self.scaler is not None: | |
scale = self.scaler.get_scale() | |
metrics['grad_scale'] = scale | |
if not loss.isfinite().all(): | |
raise RuntimeError("Model probably diverged.") | |
metrics['ce'] = ce | |
metrics['ppl'] = torch.exp(ce) | |
for k, ce_q in enumerate(ce_per_codebook): | |
metrics[f'ce_q{k + 1}'] = ce_q | |
metrics[f'ppl_q{k + 1}'] = torch.exp(ce_q) | |
return metrics | |
def run_generate_step(self, batch: tp.Tuple[torch.Tensor, tp.List[SegmentWithAttributes]], | |
gen_duration: float, prompt_duration: tp.Optional[float] = None, | |
remove_prompt: bool = False, | |
**generation_params) -> dict: | |
"""Run generate step on a batch of optional audio tensor and corresponding attributes. | |
Args: | |
batch (tuple[torch.Tensor, list[SegmentWithAttributes]]): | |
use_prompt (bool): Whether to do audio continuation generation with prompt from audio batch. | |
gen_duration (float): Target audio duration for the generation. | |
prompt_duration (float, optional): Duration for the audio prompt to use for continuation. | |
remove_prompt (bool, optional): Whether to remove the prompt from the generated audio. | |
generation_params: Additional generation parameters. | |
Returns: | |
gen_outputs (dict): Generation outputs, consisting in audio, audio tokens from both the generation | |
and the prompt along with additional information. | |
""" | |
bench_start = time.time() | |
audio, meta = batch | |
assert audio.size(0) == len(meta), ( | |
f"Mismatch between number of items in audio batch ({audio.size(0)})", | |
f" and in metadata ({len(meta)})" | |
) | |
# prepare attributes | |
attributes = [x.to_condition_attributes() for x in meta] | |
# TODO: Add dropout for chroma? | |
# prepare audio prompt | |
if prompt_duration is None: | |
prompt_audio = None | |
else: | |
assert prompt_duration < gen_duration, "Prompt duration must be lower than target generation duration" | |
prompt_audio_frames = int(prompt_duration * self.compression_model.sample_rate) | |
prompt_audio = audio[..., :prompt_audio_frames] | |
# get audio tokens from compression model | |
if prompt_audio is None or prompt_audio.nelement() == 0: | |
num_samples = len(attributes) | |
prompt_tokens = None | |
else: | |
num_samples = None | |
prompt_audio = prompt_audio.to(self.device) | |
prompt_tokens, scale = self.compression_model.encode(prompt_audio) | |
assert scale is None, "Compression model in MusicGen should not require rescaling." | |
# generate by sampling from the LM | |
with self.autocast: | |
total_gen_len = math.ceil(gen_duration * self.compression_model.frame_rate) | |
gen_tokens = self.model.generate( | |
prompt_tokens, attributes, max_gen_len=total_gen_len, | |
num_samples=num_samples, **self.generation_params) | |
# generate audio from tokens | |
assert gen_tokens.dim() == 3 | |
gen_audio = self.compression_model.decode(gen_tokens, None) | |
bench_end = time.time() | |
gen_outputs = { | |
'rtf': (bench_end - bench_start) / gen_duration, | |
'ref_audio': audio, | |
'gen_audio': gen_audio, | |
'gen_tokens': gen_tokens, | |
'prompt_audio': prompt_audio, | |
'prompt_tokens': prompt_tokens, | |
} | |
return gen_outputs | |
def generate_audio(self) -> dict: | |
"""Audio generation stage.""" | |
generate_stage_name = f'{self.current_stage}' | |
sample_manager = SampleManager(self.xp) | |
self.logger.info(f"Generating samples in {sample_manager.base_folder}") | |
loader = self.dataloaders['generate'] | |
updates = len(loader) | |
lp = self.log_progress(generate_stage_name, loader, total=updates, updates=self.log_updates) | |
dataset = get_dataset_from_loader(loader) | |
dataset_duration = dataset.segment_duration | |
assert dataset_duration is not None | |
assert isinstance(dataset, AudioDataset) | |
target_duration = self.cfg.generate.lm.gen_duration | |
prompt_duration = self.cfg.generate.lm.prompt_duration | |
if target_duration is None: | |
target_duration = dataset_duration | |
if prompt_duration is None: | |
prompt_duration = dataset_duration / 4 | |
assert prompt_duration < dataset_duration, ( | |
f"Specified prompt duration ({prompt_duration}s) is longer", | |
f" than reference audio duration ({dataset_duration}s)" | |
) | |
def get_hydrated_conditions(meta: tp.List[SegmentWithAttributes]): | |
hydrated_conditions = [] | |
for sample in [x.to_condition_attributes() for x in meta]: | |
cond_dict = {} | |
for cond_type in sample.__annotations__.keys(): | |
for cond_key, cond_val in getattr(sample, cond_type).items(): | |
if cond_key not in self.model.condition_provider.conditioners.keys(): | |
continue | |
if is_jsonable(cond_val): | |
cond_dict[cond_key] = cond_val | |
elif isinstance(cond_val, WavCondition): | |
cond_dict[cond_key] = cond_val.path | |
elif isinstance(cond_val, JointEmbedCondition): | |
cond_dict[cond_key] = cond_val.text # only support text at inference for now | |
else: | |
# if we reached this point, it is not clear how to log the condition | |
# so we just log the type. | |
cond_dict[cond_key] = str(type(cond_val)) | |
continue | |
hydrated_conditions.append(cond_dict) | |
return hydrated_conditions | |
metrics: dict = {} | |
average = flashy.averager() | |
for batch in lp: | |
audio, meta = batch | |
# metadata for sample manager | |
hydrated_conditions = get_hydrated_conditions(meta) | |
sample_generation_params = { | |
**{f'classifier_free_guidance_{k}': v for k, v in self.cfg.classifier_free_guidance.items()}, | |
**self.generation_params | |
} | |
if self.cfg.generate.lm.unprompted_samples: | |
if self.cfg.generate.lm.gen_gt_samples: | |
# get the ground truth instead of generation | |
self.logger.warn( | |
"Use ground truth instead of audio generation as generate.lm.gen_gt_samples=true") | |
gen_unprompted_audio = audio | |
rtf = 1. | |
else: | |
gen_unprompted_outputs = self.run_generate_step( | |
batch, gen_duration=target_duration, prompt_duration=prompt_duration, | |
**self.generation_params) | |
gen_unprompted_audio = gen_unprompted_outputs['gen_audio'].cpu() | |
rtf = gen_unprompted_outputs['rtf'] | |
sample_manager.add_samples( | |
gen_unprompted_audio, self.epoch, hydrated_conditions, | |
ground_truth_wavs=audio, generation_args=sample_generation_params) | |
if self.cfg.generate.lm.prompted_samples: | |
gen_outputs = self.run_generate_step( | |
batch, gen_duration=target_duration, prompt_duration=prompt_duration, | |
**self.generation_params) | |
gen_audio = gen_outputs['gen_audio'].cpu() | |
prompt_audio = gen_outputs['prompt_audio'].cpu() | |
sample_manager.add_samples( | |
gen_audio, self.epoch, hydrated_conditions, | |
prompt_wavs=prompt_audio, ground_truth_wavs=audio, | |
generation_args=sample_generation_params) | |
metrics['rtf'] = rtf | |
metrics = average(metrics) | |
flashy.distrib.barrier() | |
return metrics | |
def generate(self) -> dict: | |
"""Generate stage.""" | |
self.model.eval() | |
with torch.no_grad(): | |
return self.generate_audio() | |
def run_epoch(self): | |
if self.cfg.cache.write: | |
if ((self.epoch - 1) % self.cfg.cache.write_num_shards) != self.cfg.cache.write_shard: | |
return | |
super().run_epoch() | |
def train(self): | |
"""Train stage. | |
""" | |
if self._cached_batch_writer is not None: | |
self._cached_batch_writer.start_epoch(self.epoch) | |
if self._cached_batch_loader is None: | |
dataset = get_dataset_from_loader(self.dataloaders['train']) | |
assert isinstance(dataset, AudioDataset) | |
dataset.current_epoch = self.epoch | |
else: | |
self._cached_batch_loader.start_epoch(self.epoch) | |
return super().train() | |
def evaluate_audio_generation(self) -> dict: | |
"""Evaluate audio generation with off-the-shelf metrics.""" | |
evaluate_stage_name = f'{self.current_stage}_generation' | |
# instantiate evaluation metrics, if at least one metric is defined, run audio generation evaluation | |
fad: tp.Optional[eval_metrics.FrechetAudioDistanceMetric] = None | |
kldiv: tp.Optional[eval_metrics.KLDivergenceMetric] = None | |
text_consistency: tp.Optional[eval_metrics.TextConsistencyMetric] = None | |
chroma_cosine: tp.Optional[eval_metrics.ChromaCosineSimilarityMetric] = None | |
should_run_eval = False | |
eval_chroma_wavs: tp.Optional[torch.Tensor] = None | |
if self.cfg.evaluate.metrics.fad: | |
fad = builders.get_fad(self.cfg.metrics.fad).to(self.device) | |
should_run_eval = True | |
if self.cfg.evaluate.metrics.kld: | |
kldiv = builders.get_kldiv(self.cfg.metrics.kld).to(self.device) | |
should_run_eval = True | |
if self.cfg.evaluate.metrics.text_consistency: | |
text_consistency = builders.get_text_consistency(self.cfg.metrics.text_consistency).to(self.device) | |
should_run_eval = True | |
if self.cfg.evaluate.metrics.chroma_cosine: | |
chroma_cosine = builders.get_chroma_cosine_similarity(self.cfg.metrics.chroma_cosine).to(self.device) | |
# if we have predefind wavs for chroma we should purge them for computing the cosine metric | |
has_predefined_eval_chromas = 'self_wav' in self.model.condition_provider.conditioners and \ | |
self.model.condition_provider.conditioners['self_wav'].has_eval_wavs() | |
if has_predefined_eval_chromas: | |
warn_once(self.logger, "Attempting to run cosine eval for config with pre-defined eval chromas! " | |
'Resetting eval chromas to None for evaluation.') | |
eval_chroma_wavs = self.model.condition_provider.conditioners.self_wav.eval_wavs # type: ignore | |
self.model.condition_provider.conditioners.self_wav.reset_eval_wavs(None) # type: ignore | |
should_run_eval = True | |
def get_compressed_audio(audio: torch.Tensor) -> torch.Tensor: | |
audio_tokens, scale = self.compression_model.encode(audio.to(self.device)) | |
compressed_audio = self.compression_model.decode(audio_tokens, scale) | |
return compressed_audio[..., :audio.shape[-1]] | |
metrics: dict = {} | |
if should_run_eval: | |
loader = self.dataloaders['evaluate'] | |
updates = len(loader) | |
lp = self.log_progress(f'{evaluate_stage_name} inference', loader, total=updates, updates=self.log_updates) | |
average = flashy.averager() | |
dataset = get_dataset_from_loader(loader) | |
assert isinstance(dataset, AudioDataset) | |
self.logger.info(f"Computing evaluation metrics on {len(dataset)} samples") | |
for idx, batch in enumerate(lp): | |
audio, meta = batch | |
assert all([self.cfg.sample_rate == m.sample_rate for m in meta]) | |
target_duration = audio.shape[-1] / self.cfg.sample_rate | |
if self.cfg.evaluate.fixed_generation_duration: | |
target_duration = self.cfg.evaluate.fixed_generation_duration | |
gen_outputs = self.run_generate_step( | |
batch, gen_duration=target_duration, | |
**self.generation_params | |
) | |
y_pred = gen_outputs['gen_audio'].detach() | |
y_pred = y_pred[..., :audio.shape[-1]] | |
normalize_kwargs = dict(self.cfg.generate.audio) | |
normalize_kwargs.pop('format', None) | |
y_pred = torch.stack([normalize_audio(w, **normalize_kwargs) for w in y_pred], dim=0).cpu() | |
y = audio.cpu() # should already be on CPU but just in case | |
sizes = torch.tensor([m.n_frames for m in meta]) # actual sizes without padding | |
sample_rates = torch.tensor([m.sample_rate for m in meta]) # sample rates for audio samples | |
audio_stems = [Path(m.meta.path).stem + f"_{m.seek_time}" for m in meta] | |
if fad is not None: | |
if self.cfg.metrics.fad.use_gt: | |
y_pred = get_compressed_audio(y).cpu() | |
fad.update(y_pred, y, sizes, sample_rates, audio_stems) | |
if kldiv is not None: | |
if self.cfg.metrics.kld.use_gt: | |
y_pred = get_compressed_audio(y).cpu() | |
kldiv.update(y_pred, y, sizes, sample_rates) | |
if text_consistency is not None: | |
texts = [m.description for m in meta] | |
if self.cfg.metrics.text_consistency.use_gt: | |
y_pred = y | |
text_consistency.update(y_pred, texts, sizes, sample_rates) | |
if chroma_cosine is not None: | |
if self.cfg.metrics.chroma_cosine.use_gt: | |
y_pred = get_compressed_audio(y).cpu() | |
chroma_cosine.update(y_pred, y, sizes, sample_rates) | |
# restore chroma conditioner's eval chroma wavs | |
if eval_chroma_wavs is not None: | |
self.model.condition_provider.conditioners['self_wav'].reset_eval_wavs(eval_chroma_wavs) | |
flashy.distrib.barrier() | |
if fad is not None: | |
metrics['fad'] = fad.compute() | |
if kldiv is not None: | |
kld_metrics = kldiv.compute() | |
metrics.update(kld_metrics) | |
if text_consistency is not None: | |
metrics['text_consistency'] = text_consistency.compute() | |
if chroma_cosine is not None: | |
metrics['chroma_cosine'] = chroma_cosine.compute() | |
metrics = average(metrics) | |
metrics = flashy.distrib.average_metrics(metrics, len(loader)) | |
return metrics | |
def evaluate(self) -> dict: | |
"""Evaluate stage.""" | |
self.model.eval() | |
with torch.no_grad(): | |
metrics: dict = {} | |
if self.cfg.evaluate.metrics.base: | |
metrics.update(self.common_train_valid('evaluate')) | |
gen_metrics = self.evaluate_audio_generation() | |
return {**metrics, **gen_metrics} | |