#====================================== From CompressionSolver.py
# 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.
import logging
import multiprocessing
from pathlib import Path
import typing as tp
import flashy
import omegaconf
import torch
from torch import nn
# from . import base, builders
from .. import models, quantization
from ..utils import checkpoint
from ..utils.samples.manager import SampleManager
from ..utils.utils import get_pool_executor
class CompressionSolver(): #base.StandardSolver):
"""Solver for compression task.
The compression task combines a set of perceptual and objective losses
to train an EncodecModel (composed of an encoder-decoder and a quantizer)
to perform high fidelity audio reconstruction.
"""
def __init__(self, cfg: omegaconf.DictConfig):
# super().__init__(cfg)
self.cfg = cfg
self.rng: torch.Generator # set at each epoch
self.adv_losses = builders.get_adversarial_losses(self.cfg)
self.aux_losses = nn.ModuleDict()
self.info_losses = nn.ModuleDict()
assert not cfg.fsdp.use, "FSDP not supported by CompressionSolver."
loss_weights = dict()
for loss_name, weight in self.cfg.losses.items():
if loss_name in ['adv', 'feat']:
for adv_name, _ in self.adv_losses.items():
loss_weights[f'{loss_name}_{adv_name}'] = weight
elif weight > 0:
self.aux_losses[loss_name] = builders.get_loss(loss_name, self.cfg)
loss_weights[loss_name] = weight
else:
self.info_losses[loss_name] = builders.get_loss(loss_name, self.cfg)
self.balancer = builders.get_balancer(loss_weights, self.cfg.balancer)
self.register_stateful('adv_losses')
@property
def best_metric_name(self) -> tp.Optional[str]:
# best model is the last for the compression model
return None
def build_model(self):
"""Instantiate model and optimizer."""
# Model and optimizer
self.model = models.builders.get_compression_model(self.cfg).to(self.device)
self.optimizer = builders.get_optimizer(self.model.parameters(), self.cfg.optim)
self.register_stateful('model', 'optimizer')
self.register_best_state('model')
self.register_ema('model')
def evaluate(self):
"""Evaluate stage. Runs audio reconstruction evaluation."""
self.model.eval()
evaluate_stage_name = str(self.current_stage)
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()
pendings = []
ctx = multiprocessing.get_context('spawn')
with get_pool_executor(self.cfg.evaluate.num_workers, mp_context=ctx) as pool:
for idx, batch in enumerate(lp):
x = batch.to(self.device)
with torch.no_grad():
qres = self.model(x)
y_pred = qres.x.cpu()
y = batch.cpu() # should already be on CPU but just in case
pendings.append(pool.submit(evaluate_audio_reconstruction, y_pred, y, self.cfg))
metrics_lp = self.log_progress(f'{evaluate_stage_name} metrics', pendings, updates=self.log_updates)
for pending in metrics_lp:
metrics = pending.result()
metrics = average(metrics)
metrics = flashy.distrib.average_metrics(metrics, len(loader))
return metrics
def generate(self):
"""Generate stage."""
self.model.eval()
sample_manager = SampleManager(self.xp, map_reference_to_sample_id=True)
generate_stage_name = str(self.current_stage)
loader = self.dataloaders['generate']
updates = len(loader)
lp = self.log_progress(generate_stage_name, loader, total=updates, updates=self.log_updates)
for batch in lp:
reference, _ = batch
reference = reference.to(self.device)
with torch.no_grad():
qres = self.model(reference)
assert isinstance(qres, quantization.QuantizedResult)
reference = reference.cpu()
estimate = qres.x.cpu()
sample_manager.add_samples(estimate, self.epoch, ground_truth_wavs=reference)
flashy.distrib.barrier()
def load_from_pretrained(self, name: str) -> dict:
model = models.CompressionModel.get_pretrained(name)
if isinstance(model, models.DAC):
raise RuntimeError("Cannot fine tune a DAC model.")
elif isinstance(model, models.HFEncodecCompressionModel):
self.logger.warning('Trying to automatically convert a HuggingFace model '
'to AudioCraft, this might fail!')
state = model.model.state_dict()
new_state = {}
for k, v in state.items():
if k.startswith('decoder.layers') and '.conv.' in k and '.block.' not in k:
# We need to determine if this a convtr or a regular conv.
layer = int(k.split('.')[2])
if isinstance(model.model.decoder.layers[layer].conv, torch.nn.ConvTranspose1d):
k = k.replace('.conv.', '.convtr.')
k = k.replace('encoder.layers.', 'encoder.model.')
k = k.replace('decoder.layers.', 'decoder.model.')
k = k.replace('conv.', 'conv.conv.')
k = k.replace('convtr.', 'convtr.convtr.')
k = k.replace('quantizer.layers.', 'quantizer.vq.layers.')
k = k.replace('.codebook.', '._codebook.')
new_state[k] = v
state = new_state
elif isinstance(model, models.EncodecModel):
state = model.state_dict()
else:
raise RuntimeError(f"Cannot fine tune model type {type(model)}.")
return {
'best_state': {'model': state}
}
@staticmethod
def model_from_checkpoint(checkpoint_path: tp.Union[Path, str],
device: tp.Union[torch.device, str] = 'cpu') -> models.CompressionModel:
"""Instantiate a CompressionModel from a given checkpoint path or dora sig.
This method is a convenient endpoint to load a CompressionModel to use in other solvers.
Args:
checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.
This also supports pre-trained models by using a path of the form //pretrained/NAME.
See `model_from_pretrained` for a list of supported pretrained models.
use_ema (bool): Use EMA variant of the model instead of the actual model.
device (torch.device or str): Device on which the model is loaded.
"""
checkpoint_path = str(checkpoint_path)
if checkpoint_path.startswith('//pretrained/'):
name = checkpoint_path.split('/', 3)[-1]
return models.CompressionModel.get_pretrained(name, device)
logger = logging.getLogger(__name__)
logger.info(f"Loading compression model from checkpoint: {checkpoint_path}")
_checkpoint_path = checkpoint.resolve_checkpoint_path(checkpoint_path, use_fsdp=False)
assert _checkpoint_path is not None, f"Could not resolve compression model checkpoint path: {checkpoint_path}"
state = checkpoint.load_checkpoint(_checkpoint_path)
assert state is not None and 'xp.cfg' in state, f"Could not load compression model from ckpt: {checkpoint_path}"
cfg = state['xp.cfg']
cfg.device = device
compression_model = models.builders.get_compression_model(cfg).to(device)
assert compression_model.sample_rate == cfg.sample_rate, "Compression model sample rate should match"
assert 'best_state' in state and state['best_state'] != {}
assert 'exported' not in state, "When loading an exported checkpoint, use the //pretrained/ prefix."
compression_model.load_state_dict(state['best_state']['model'])
compression_model.eval()
logger.info("Compression model loaded!")
return compression_model
@staticmethod
def wrapped_model_from_checkpoint(cfg: omegaconf.DictConfig,
checkpoint_path: tp.Union[Path, str],
device: tp.Union[torch.device, str] = 'cpu') -> models.CompressionModel:
"""Instantiate a wrapped CompressionModel from a given checkpoint path or dora sig.
Args:
cfg (omegaconf.DictConfig): Configuration to read from for wrapped mode.
checkpoint_path (Path or str): Path to checkpoint or dora sig from where the checkpoint is resolved.
use_ema (bool): Use EMA variant of the model instead of the actual model.
device (torch.device or str): Device on which the model is loaded.
"""
compression_model = CompressionSolver.model_from_checkpoint(checkpoint_path, device)
compression_model = models.builders.get_wrapped_compression_model(compression_model, cfg)
return compression_model
#=========================================================================== ORIG
import typing as tp
import torch
import julius
from .unet import DiffusionUnet
from ..modules.diffusion_schedule import NoiseSchedule
from .encodec import CompressionModel
from .loaders import load_compression_model, load_diffusion_models
class DiffusionProcess:
"""Sampling for a diffusion Model.
Args:
model (DiffusionUnet): Diffusion U-Net model.
noise_schedule (NoiseSchedule): Noise schedule for diffusion process.
"""
def __init__(self, model: DiffusionUnet, noise_schedule: NoiseSchedule) -> None:
self.model = model
self.schedule = noise_schedule
def generate(self, condition: torch.Tensor, initial_noise: torch.Tensor,
step_list: tp.Optional[tp.List[int]] = None):
"""Perform one diffusion process to generate one of the bands.
Args:
condition (torch.Tensor): The embeddings from the compression model.
initial_noise (torch.Tensor): The initial noise to start the process.
"""
return self.schedule.generate_subsampled(model=self.model, initial=initial_noise, step_list=step_list,
condition=condition)
class MultiBandDiffusion:
"""Sample from multiple diffusion models.
Args:
DPs (list of DiffusionProcess): Diffusion processes.
codec_model (CompressionModel): Underlying compression model used to obtain discrete tokens.
"""
def __init__(self, DPs: tp.List[DiffusionProcess], codec_model: CompressionModel) -> None:
self.DPs = DPs
self.codec_model = codec_model
self.device = next(self.codec_model.parameters()).device
@property
def sample_rate(self) -> int:
return self.codec_model.sample_rate
@staticmethod
def get_mbd_musicgen(device=None):
"""Load our diffusion models trained for MusicGen."""
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
path = 'facebook/multiband-diffusion'
filename = 'mbd_musicgen_32khz.th'
name = 'facebook/musicgen-small'
codec_model = load_compression_model(name, device=device)
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
DPs = []
for i in range(len(models)):
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)
@staticmethod
def get_mbd_24khz(bw: float = 3.0,
device: tp.Optional[tp.Union[torch.device, str]] = None,
n_q: tp.Optional[int] = None):
"""Get the pretrained Models for MultibandDiffusion.
Args:
bw (float): Bandwidth of the compression model.
device (torch.device or str, optional): Device on which the models are loaded.
n_q (int, optional): Number of quantizers to use within the compression model.
"""
if device is None:
device = 'cuda' if torch.cuda.is_available() else 'cpu'
assert bw in [1.5, 3.0, 6.0], f"bandwidth {bw} not available"
if n_q is not None:
assert n_q in [2, 4, 8]
assert {1.5: 2, 3.0: 4, 6.0: 8}[bw] == n_q, \
f"bandwidth and number of codebooks missmatch to use n_q = {n_q} bw should be {n_q * (1.5 / 2)}"
n_q = {1.5: 2, 3.0: 4, 6.0: 8}[bw]
codec_model = CompressionSolver.model_from_checkpoint(
'//pretrained/facebook/encodec_24khz', device=device)
codec_model.set_num_codebooks(n_q)
codec_model = codec_model.to(device)
path = 'facebook/multiband-diffusion'
filename = f'mbd_comp_{n_q}.pt'
models, processors, cfgs = load_diffusion_models(path, filename=filename, device=device)
DPs = []
for i in range(len(models)):
schedule = NoiseSchedule(**cfgs[i].schedule, sample_processor=processors[i], device=device)
DPs.append(DiffusionProcess(model=models[i], noise_schedule=schedule))
return MultiBandDiffusion(DPs=DPs, codec_model=codec_model)
@torch.no_grad()
def get_condition(self, wav: torch.Tensor, sample_rate: int) -> torch.Tensor:
"""Get the conditioning (i.e. latent representations of the compression model) from a waveform.
Args:
wav (torch.Tensor): The audio that we want to extract the conditioning from.
sample_rate (int): Sample rate of the audio."""
if sample_rate != self.sample_rate:
wav = julius.resample_frac(wav, sample_rate, self.sample_rate)
codes, scale = self.codec_model.encode(wav)
assert scale is None, "Scaled compression models not supported."
emb = self.get_emb(codes)
return emb
@torch.no_grad()
def get_emb(self, codes: torch.Tensor):
"""Get latent representation from the discrete codes.
Args:
codes (torch.Tensor): Discrete tokens."""
emb = self.codec_model.decode_latent(codes)
return emb
def generate(self, emb: torch.Tensor, size: tp.Optional[torch.Size] = None,
step_list: tp.Optional[tp.List[int]] = None):
"""Generate waveform audio from the latent embeddings of the compression model.
Args:
emb (torch.Tensor): Conditioning embeddings
size (None, torch.Size): Size of the output
if None this is computed from the typical upsampling of the model.
step_list (list[int], optional): list of Markov chain steps, defaults to 50 linearly spaced step.
"""
if size is None:
upsampling = int(self.codec_model.sample_rate / self.codec_model.frame_rate)
size = torch.Size([emb.size(0), self.codec_model.channels, emb.size(-1) * upsampling])
assert size[0] == emb.size(0)
out = torch.zeros(size).to(self.device)
for DP in self.DPs:
out += DP.generate(condition=emb, step_list=step_list, initial_noise=torch.randn_like(out))
return out
def re_eq(self, wav: torch.Tensor, ref: torch.Tensor, n_bands: int = 32, strictness: float = 1):
"""Match the eq to the encodec output by matching the standard deviation of some frequency bands.
Args:
wav (torch.Tensor): Audio to equalize.
ref (torch.Tensor): Reference audio from which we match the spectrogram.
n_bands (int): Number of bands of the eq.
strictness (float): How strict the matching. 0 is no matching, 1 is exact matching.
"""
split = julius.SplitBands(n_bands=n_bands, sample_rate=self.codec_model.sample_rate).to(wav.device)
bands = split(wav)
bands_ref = split(ref)
out = torch.zeros_like(ref)
for i in range(n_bands):
out += bands[i] * (bands_ref[i].std() / bands[i].std()) ** strictness
return out
def regenerate(self, wav: torch.Tensor, sample_rate: int):
"""Regenerate a waveform through compression and diffusion regeneration.
Args:
wav (torch.Tensor): Original 'ground truth' audio.
sample_rate (int): Sample rate of the input (and output) wav.
"""
if sample_rate != self.codec_model.sample_rate:
wav = julius.resample_frac(wav, sample_rate, self.codec_model.sample_rate)
emb = self.get_condition(wav, sample_rate=self.codec_model.sample_rate)
size = wav.size()
out = self.generate(emb, size=size)
if sample_rate != self.codec_model.sample_rate:
out = julius.resample_frac(out, self.codec_model.sample_rate, sample_rate)
return out
def tokens_to_wav(self, tokens: torch.Tensor, n_bands: int = 32):
"""Generate Waveform audio with diffusion from the discrete codes.
Args:
tokens (torch.Tensor): Discrete codes.
n_bands (int): Bands for the eq matching.
"""
wav_encodec = self.codec_model.decode(tokens)
condition = self.get_emb(tokens)
wav_diffusion = self.generate(emb=condition, size=wav_encodec.size())
return self.re_eq(wav=wav_diffusion, ref=wav_encodec, n_bands=n_bands)