Module audiocraft.solvers.builders

All the functions to build the relevant solvers and used objects from the Hydra config.

Functions

def get_adversarial_losses(cfg) ‑> torch.nn.modules.container.ModuleDict

Expand source code

def get_adversarial_losses(cfg) -> nn.ModuleDict:
    """Initialize dict of adversarial losses from config."""
    device = cfg.device
    adv_cfg = getattr(cfg, 'adversarial')
    adversaries = adv_cfg.get('adversaries', [])
    adv_loss_name = adv_cfg['adv_loss']
    feat_loss_name = adv_cfg.get('feat_loss')
    normalize = adv_cfg.get('normalize', True)
    feat_loss: tp.Optional[adversarial.FeatureMatchingLoss] = None
    if feat_loss_name:
        assert feat_loss_name in ['l1', 'l2'], f"Feature loss only support L1 or L2 but {feat_loss_name} found."
        loss = get_loss(feat_loss_name, cfg)
        feat_loss = adversarial.FeatureMatchingLoss(loss, normalize)
    loss = adversarial.get_adv_criterion(adv_loss_name)
    loss_real = adversarial.get_real_criterion(adv_loss_name)
    loss_fake = adversarial.get_fake_criterion(adv_loss_name)
    adv_losses = nn.ModuleDict()
    for adv_name in adversaries:
        adversary = get_adversary(adv_name, cfg).to(device)
        optimizer = get_optimizer(adversary.parameters(), cfg.optim)
        adv_loss = adversarial.AdversarialLoss(
            adversary,
            optimizer,
            loss=loss,
            loss_real=loss_real,
            loss_fake=loss_fake,
            loss_feat=feat_loss,
            normalize=normalize
        )
        adv_losses[adv_name] = adv_loss
    return adv_losses

Initialize dict of adversarial losses from config.

def get_adversary(name: str, cfg: omegaconf.dictconfig.DictConfig) ‑> torch.nn.modules.module.Module

Expand source code

def get_adversary(name: str, cfg: omegaconf.DictConfig) -> nn.Module:
    """Initialize adversary from config."""
    klass = {
        'msd': adversarial.MultiScaleDiscriminator,
        'mpd': adversarial.MultiPeriodDiscriminator,
        'msstftd': adversarial.MultiScaleSTFTDiscriminator,
    }[name]
    adv_cfg: tp.Dict[str, tp.Any] = dict(getattr(cfg, name))
    return klass(**adv_cfg)

Initialize adversary from config.

def get_audio_datasets(cfg: omegaconf.dictconfig.DictConfig,
dataset_type: DatasetType = DatasetType.AUDIO) ‑> Dict[str, torch.utils.data.dataloader.DataLoader]

Expand source code

def get_audio_datasets(cfg: omegaconf.DictConfig,
                       dataset_type: DatasetType = DatasetType.AUDIO) -> tp.Dict[str, torch.utils.data.DataLoader]:
    """Build AudioDataset from configuration.

    Args:
        cfg (omegaconf.DictConfig): Configuration.
        dataset_type: The type of dataset to create.
    Returns:
        dict[str, torch.utils.data.DataLoader]: Map of dataloader for each data split.
    """
    dataloaders: dict = {}

    sample_rate = cfg.sample_rate
    channels = cfg.channels
    seed = cfg.seed
    max_sample_rate = cfg.datasource.max_sample_rate
    max_channels = cfg.datasource.max_channels

    assert cfg.dataset is not None, "Could not find dataset definition in config"

    dataset_cfg = dict_from_config(cfg.dataset)
    splits_cfg: dict = {}
    splits_cfg['train'] = dataset_cfg.pop('train')
    splits_cfg['valid'] = dataset_cfg.pop('valid')
    splits_cfg['evaluate'] = dataset_cfg.pop('evaluate')
    splits_cfg['generate'] = dataset_cfg.pop('generate')
    execute_only_stage = cfg.get('execute_only', None)

    for split, path in cfg.datasource.items():
        if not isinstance(path, str):
            continue  # skipping this as not a path
        if execute_only_stage is not None and split != execute_only_stage:
            continue
        logger.info(f"Loading audio data split {split}: {str(path)}")
        assert (
            cfg.sample_rate <= max_sample_rate
        ), f"Expecting a max sample rate of {max_sample_rate} for datasource but {sample_rate} found."
        assert (
            cfg.channels <= max_channels
        ), f"Expecting a max number of channels of {max_channels} for datasource but {channels} found."

        split_cfg = splits_cfg[split]
        split_kwargs = {k: v for k, v in split_cfg.items()}
        kwargs = {**dataset_cfg, **split_kwargs}  # split kwargs overrides default dataset_cfg
        kwargs['sample_rate'] = sample_rate
        kwargs['channels'] = channels

        if kwargs.get('permutation_on_files') and cfg.optim.updates_per_epoch:
            kwargs['num_samples'] = (
                flashy.distrib.world_size() * cfg.dataset.batch_size * cfg.optim.updates_per_epoch)

        num_samples = kwargs['num_samples']
        shuffle = kwargs['shuffle']

        return_info = kwargs.pop('return_info')
        batch_size = kwargs.pop('batch_size', None)
        num_workers = kwargs.pop('num_workers')

        if dataset_type == DatasetType.MUSIC:
            dataset = data.music_dataset.MusicDataset.from_meta(path, **kwargs)
        elif dataset_type == DatasetType.SOUND:
            dataset = data.sound_dataset.SoundDataset.from_meta(path, **kwargs)
        elif dataset_type == DatasetType.AUDIO:
            dataset = data.info_audio_dataset.InfoAudioDataset.from_meta(path, return_info=return_info, **kwargs)
        elif dataset_type == DatasetType.JASCO:
            dataset = data.jasco_dataset.JascoDataset.from_meta(path, return_info=return_info, **kwargs)
        else:
            raise ValueError(f"Dataset type is unsupported: {dataset_type}")

        loader = get_loader(
            dataset,
            num_samples,
            batch_size=batch_size,
            num_workers=num_workers,
            seed=seed,
            collate_fn=dataset.collater if return_info else None,
            shuffle=shuffle,
        )
        dataloaders[split] = loader

    return dataloaders

Build AudioDataset from configuration.

Args

cfg : omegaconf.DictConfig

Configuration.

dataset_type

The type of dataset to create.

Returns

dict[str, torch.utils.data.DataLoader]

Map of dataloader for each data split.

def get_balancer(loss_weights: Dict[str, float], cfg: omegaconf.dictconfig.DictConfig) ‑> Balancer

Expand source code

def get_balancer(loss_weights: tp.Dict[str, float], cfg: omegaconf.DictConfig) -> losses.Balancer:
    """Instantiate loss balancer from configuration for the provided weights."""
    kwargs: tp.Dict[str, tp.Any] = dict_from_config(cfg)
    return losses.Balancer(loss_weights, **kwargs)

Instantiate loss balancer from configuration for the provided weights.

def get_chroma_cosine_similarity(cfg: omegaconf.dictconfig.DictConfig) ‑> ChromaCosineSimilarityMetric

Expand source code

def get_chroma_cosine_similarity(cfg: omegaconf.DictConfig) -> metrics.ChromaCosineSimilarityMetric:
    """Instantiate Chroma Cosine Similarity metric from config."""
    assert cfg.model == 'chroma_base', "Only support 'chroma_base' method for chroma cosine similarity metric"
    kwargs = dict_from_config(cfg.get(cfg.model))
    return metrics.ChromaCosineSimilarityMetric(**kwargs)

Instantiate Chroma Cosine Similarity metric from config.

def get_ema(module_dict: torch.nn.modules.container.ModuleDict,
cfg: omegaconf.dictconfig.DictConfig) ‑> ModuleDictEMA | None

Expand source code

def get_ema(module_dict: nn.ModuleDict, cfg: omegaconf.DictConfig) -> tp.Optional[optim.ModuleDictEMA]:
    """Initialize Exponential Moving Average.

    Args:
        module_dict (nn.ModuleDict): ModuleDict for which to compute the EMA.
        cfg (omegaconf.DictConfig): Optim EMA configuration.
    Returns:
        optim.ModuleDictEMA: EMA version of the ModuleDict.
    """
    kw: tp.Dict[str, tp.Any] = dict(cfg)
    use = kw.pop('use', False)
    decay = kw.pop('decay', None)
    device = kw.pop('device', None)
    if not use:
        return None
    if len(module_dict) == 0:
        raise ValueError("Trying to build EMA but an empty module_dict source is provided!")
    ema_module = optim.ModuleDictEMA(module_dict, decay=decay, device=device)
    return ema_module

Initialize Exponential Moving Average.

Args

module_dict : nn.ModuleDict

ModuleDict for which to compute the EMA.

cfg : omegaconf.DictConfig

Optim EMA configuration.

Returns

optim.ModuleDictEMA

EMA version of the ModuleDict.

def get_fad(cfg: omegaconf.dictconfig.DictConfig) ‑> FrechetAudioDistanceMetric

Expand source code

def get_fad(cfg: omegaconf.DictConfig) -> metrics.FrechetAudioDistanceMetric:
    """Instantiate Frechet Audio Distance metric from config."""
    kwargs = dict_from_config(cfg.tf)
    xp = dora.get_xp()
    kwargs['log_folder'] = xp.folder
    return metrics.FrechetAudioDistanceMetric(**kwargs)

Instantiate Frechet Audio Distance metric from config.

def get_kldiv(cfg: omegaconf.dictconfig.DictConfig) ‑> KLDivergenceMetric

Expand source code

def get_kldiv(cfg: omegaconf.DictConfig) -> metrics.KLDivergenceMetric:
    """Instantiate KL-Divergence metric from config."""
    kld_metrics = {
        'passt': metrics.PasstKLDivergenceMetric,
    }
    klass = kld_metrics[cfg.model]
    kwargs = dict_from_config(cfg.get(cfg.model))
    return klass(**kwargs)

Instantiate KL-Divergence metric from config.

def get_loss(loss_name: str, cfg: omegaconf.dictconfig.DictConfig)

Expand source code

def get_loss(loss_name: str, cfg: omegaconf.DictConfig):
    """Instantiate loss from configuration."""
    klass = {
        'l1': torch.nn.L1Loss,
        'l2': torch.nn.MSELoss,
        'mel': losses.MelSpectrogramL1Loss,
        'mrstft': losses.MRSTFTLoss,
        'msspec': losses.MultiScaleMelSpectrogramLoss,
        'sisnr': losses.SISNR,
        'wm_detection': losses.WMDetectionLoss,
        'wm_mb': losses.WMMbLoss,
        'tf_loudnessratio': losses.TFLoudnessRatio
    }[loss_name]
    kwargs = dict(getattr(cfg, loss_name))
    return klass(**kwargs)

Instantiate loss from configuration.

def get_lr_scheduler(optimizer: torch.optim.optimizer.Optimizer,
cfg: omegaconf.dictconfig.DictConfig,
total_updates: int) ‑> torch.optim.lr_scheduler.LRScheduler | None

Expand source code

def get_lr_scheduler(optimizer: torch.optim.Optimizer,
                     cfg: omegaconf.DictConfig,
                     total_updates: int) -> tp.Optional[LRScheduler]:
    """Build torch learning rate scheduler from config and associated optimizer.
    Supported learning rate schedulers: ExponentialLRScheduler, PlateauLRScheduler

    Args:
        optimizer (torch.optim.Optimizer): Optimizer.
        cfg (DictConfig): Schedule-related configuration.
        total_updates (int): Total number of updates.
    Returns:
        torch.optim.Optimizer.
    """
    if 'lr_scheduler' not in cfg:
        raise KeyError("LR Scheduler not found in config")

    lr_sched: tp.Optional[LRScheduler] = None
    if cfg.lr_scheduler == 'step':
        lr_sched = torch.optim.lr_scheduler.StepLR(optimizer, **cfg.step)
    elif cfg.lr_scheduler == 'exponential':
        lr_sched = torch.optim.lr_scheduler.ExponentialLR(optimizer, gamma=cfg.exponential)
    elif cfg.lr_scheduler == 'cosine':
        kwargs = dict_from_config(cfg.cosine)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.CosineLRScheduler(
            optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)
    elif cfg.lr_scheduler == 'polynomial_decay':
        kwargs = dict_from_config(cfg.polynomial_decay)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.PolynomialDecayLRScheduler(
            optimizer, warmup_steps=warmup_steps, total_steps=total_updates, **kwargs)
    elif cfg.lr_scheduler == 'inverse_sqrt':
        kwargs = dict_from_config(cfg.inverse_sqrt)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.InverseSquareRootLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)
    elif cfg.lr_scheduler == 'linear_warmup':
        kwargs = dict_from_config(cfg.linear_warmup)
        warmup_steps = kwargs.pop('warmup')
        lr_sched = optim.LinearWarmupLRScheduler(optimizer, warmup_steps=warmup_steps, **kwargs)
    elif cfg.lr_scheduler is not None:
        raise ValueError(f"Unsupported LR Scheduler: {cfg.lr_scheduler}")
    return lr_sched

Build torch learning rate scheduler from config and associated optimizer. Supported learning rate schedulers: ExponentialLRScheduler, PlateauLRScheduler

Args

optimizer : torch.optim.Optimizer

Optimizer.

cfg : DictConfig

Schedule-related configuration.

total_updates : int

Total number of updates.

Returns

torch.optim.Optimizer.

def get_optim_parameter_groups(model: torch.nn.modules.module.Module)

Expand source code

def get_optim_parameter_groups(model: nn.Module):
    """Create parameter groups for the model using the appropriate method
    if defined for each modules, to create the different groups.

    Args:
        model (nn.Module): torch model
    Returns:
        List of parameter groups
    """
    seen_params: tp.Set[nn.parameter.Parameter] = set()
    other_params = []
    groups = []
    for name, module in model.named_modules():
        if hasattr(module, 'make_optim_group'):
            group = module.make_optim_group()
            params = set(group['params'])
            assert params.isdisjoint(seen_params)
            seen_params |= set(params)
            groups.append(group)
    for param in model.parameters():
        if param not in seen_params:
            other_params.append(param)
    groups.insert(0, {'params': other_params})
    parameters = groups
    return parameters

Create parameter groups for the model using the appropriate method if defined for each modules, to create the different groups.

Args

model : nn.Module

torch model

Returns

List of parameter groups

def get_optimizer(params: torch.nn.modules.module.Module | Iterable[torch.Tensor],
cfg: omegaconf.dictconfig.DictConfig) ‑> torch.optim.optimizer.Optimizer

Expand source code

def get_optimizer(params: tp.Union[nn.Module, tp.Iterable[torch.Tensor]], cfg: omegaconf.DictConfig) -> Optimizer:
    """Build torch optimizer from config and set of parameters.
    Supported optimizers: Adam, AdamW

    Args:
        params (nn.Module or iterable of torch.Tensor): Parameters to optimize.
        cfg (DictConfig): Optimization-related configuration.
    Returns:
        torch.optim.Optimizer.
    """
    if 'optimizer' not in cfg:
        if getattr(cfg, 'optim', None) is not None:
            raise KeyError("Optimizer not found in config. Try instantiating optimizer from cfg.optim?")
        else:
            raise KeyError("Optimizer not found in config.")

    parameters = get_optim_parameter_groups(params) if isinstance(params, nn.Module) else params
    optimizer: torch.optim.Optimizer
    if cfg.optimizer == 'adam':
        optimizer = torch.optim.Adam(parameters, lr=cfg.lr, **cfg.adam)
    elif cfg.optimizer == 'adamw':
        optimizer = torch.optim.AdamW(parameters, lr=cfg.lr, **cfg.adam)
    elif cfg.optimizer == 'dadam':
        optimizer = optim.DAdaptAdam(parameters, lr=cfg.lr, **cfg.adam)
    else:
        raise ValueError(f"Unsupported Optimizer: {cfg.optimizer}")
    return optimizer

Build torch optimizer from config and set of parameters. Supported optimizers: Adam, AdamW

Args

params : nn.Module or iterable of torch.Tensor

Parameters to optimize.

cfg : DictConfig

Optimization-related configuration.

Returns

torch.optim.Optimizer.

def get_solver(cfg: omegaconf.dictconfig.DictConfig) ‑> StandardSolver

Expand source code

def get_solver(cfg: omegaconf.DictConfig) -> StandardSolver:
    """Instantiate solver from config."""
    from .audiogen import AudioGenSolver
    from .compression import CompressionSolver
    from .musicgen import MusicGenSolver
    from .diffusion import DiffusionSolver
    from .magnet import MagnetSolver, AudioMagnetSolver
    from .watermark import WatermarkSolver
    from .jasco import JascoSolver
    klass = {
        'compression': CompressionSolver,
        'musicgen': MusicGenSolver,
        'audiogen': AudioGenSolver,
        'magnet': MagnetSolver,
        'audio_magnet': AudioMagnetSolver,
        'lm': MusicGenSolver,  # backward compatibility
        'diffusion': DiffusionSolver,
        'sound_lm': AudioGenSolver,  # backward compatibility
        'watermarking': WatermarkSolver,
        'jasco': JascoSolver,
    }[cfg.solver]
    return klass(cfg)  # type: ignore

Instantiate solver from config.

def get_text_consistency(cfg: omegaconf.dictconfig.DictConfig) ‑> TextConsistencyMetric

Expand source code

def get_text_consistency(cfg: omegaconf.DictConfig) -> metrics.TextConsistencyMetric:
    """Instantiate Text Consistency metric from config."""
    text_consistency_metrics = {
        'clap': metrics.CLAPTextConsistencyMetric
    }
    klass = text_consistency_metrics[cfg.model]
    kwargs = dict_from_config(cfg.get(cfg.model))
    return klass(**kwargs)

Instantiate Text Consistency metric from config.

def get_visqol(cfg: omegaconf.dictconfig.DictConfig) ‑> ViSQOL

Expand source code

def get_visqol(cfg: omegaconf.DictConfig) -> metrics.ViSQOL:
    """Instantiate ViSQOL metric from config."""
    kwargs = dict_from_config(cfg)
    return metrics.ViSQOL(**kwargs)

Instantiate ViSQOL metric from config.

Classes

class DatasetType (value, names=None, *, module=None, qualname=None, type=None, start=1)

Expand source code

class DatasetType(Enum):
    AUDIO = "audio"
    MUSIC = "music"
    SOUND = "sound"
    JASCO = "jasco"

An enumeration.

Ancestors

• enum.Enum

Class variables

var AUDIO

var JASCO

var MUSIC

var SOUND