softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6 7from concurrent.futures import ProcessPoolExecutor8from contextlib import contextmanager9from functools import wraps, lru_cache10import hashlib11import json12import logging13from pathlib import Path14import typing as tp15 16import flashy17import flashy.distrib18import omegaconf19import torch20from torch.nn.utils.rnn import pad_sequence21 22 23logger = logging.getLogger(__name__)24 25 26def model_hash(model: torch.nn.Module) -> str:27    """Return a model hash. This should allow us to track regressions in model init28    from the logs of past experiments.29    """30    hasher = hashlib.sha1()31    for p in model.parameters():32        hasher.update(p.data.cpu().numpy().tobytes())33    return hasher.hexdigest()34 35 36def dict_from_config(cfg: omegaconf.DictConfig) -> dict:37    """Convenience function to map an omegaconf configuration to a dictionary.38 39    Args:40        cfg (omegaconf.DictConfig): Original configuration to map to dict.41    Returns:42        dict: Config as dictionary object.43    """44    dct = omegaconf.OmegaConf.to_container(cfg, resolve=True)45    assert isinstance(dct, dict)46    return dct47 48 49def random_subset(dataset, max_samples: int, seed: int = 42) -> torch.utils.data.Subset:50    if max_samples >= len(dataset):51        return dataset52 53    generator = torch.Generator().manual_seed(seed)54    perm = torch.randperm(len(dataset), generator=generator)55    return torch.utils.data.Subset(dataset, perm[:max_samples].tolist())56 57 58def get_loader(dataset, num_samples: tp.Optional[int], batch_size: int,59               num_workers: int, seed: int, **kwargs) -> torch.utils.data.DataLoader:60    """Convenience function to load dataset into a dataloader with optional subset sampling.61 62    Args:63        dataset: Dataset to load.64        num_samples (Optional[int]): Number of samples to limit subset size.65        batch_size (int): Batch size.66        num_workers (int): Number of workers for data loading.67        seed (int): Random seed.68    """69    if num_samples is not None:70        dataset = random_subset(dataset, num_samples, seed)71 72    dataloader = flashy.distrib.loader(73        dataset,74        batch_size=batch_size,75        num_workers=num_workers,76        **kwargs77    )78    return dataloader79 80 81def get_dataset_from_loader(dataloader):82    dataset = dataloader.dataset83    if isinstance(dataset, torch.utils.data.Subset):84        return dataset.dataset85    else:86        return dataset87 88 89def multinomial(input: torch.Tensor, num_samples: int, replacement=False, *, generator=None):90    """torch.multinomial with arbitrary number of dimensions, and number of candidates on the last dimension.91 92    Args:93        input (torch.Tensor): The input tensor containing probabilities.94        num_samples (int): Number of samples to draw.95        replacement (bool): Whether to draw with replacement or not.96    Keywords args:97        generator (torch.Generator): A pseudorandom number generator for sampling.98    Returns:99        torch.Tensor: Last dimension contains num_samples indices100            sampled from the multinomial probability distribution101            located in the last dimension of tensor input.102    """103    input_ = input.reshape(-1, input.shape[-1])104    output_ = torch.multinomial(input_, num_samples=num_samples, replacement=replacement, generator=generator)105    output = output_.reshape(*list(input.shape[:-1]), -1)106    return output107 108 109def sample_top_k(probs: torch.Tensor, k: int) -> torch.Tensor:110    """Sample next token from top K values along the last dimension of the input probs tensor.111 112    Args:113        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.114        k (int): The k in “top-k”.115    Returns:116        torch.Tensor: Sampled tokens.117    """118    top_k_value, _ = torch.topk(probs, k, dim=-1)119    min_value_top_k = top_k_value[..., [-1]]120    probs *= (probs >= min_value_top_k).float()121    probs.div_(probs.sum(dim=-1, keepdim=True))122    next_token = multinomial(probs, num_samples=1)123    return next_token124 125 126def sample_top_p(probs: torch.Tensor, p: float) -> torch.Tensor:127    """Sample next token from top P probabilities along the last dimension of the input probs tensor.128 129    Args:130        probs (torch.Tensor): Input probabilities with token candidates on the last dimension.131        p (int): The p in “top-p”.132    Returns:133        torch.Tensor: Sampled tokens.134    """135    probs_sort, probs_idx = torch.sort(probs, dim=-1, descending=True)136    probs_sum = torch.cumsum(probs_sort, dim=-1)137    mask = probs_sum - probs_sort > p138    probs_sort *= (~mask).float()139    probs_sort.div_(probs_sort.sum(dim=-1, keepdim=True))140    next_token = multinomial(probs_sort, num_samples=1)141    next_token = torch.gather(probs_idx, -1, next_token)142    return next_token143 144 145class DummyPoolExecutor:146    """Dummy pool executor to use when we actually have only 1 worker.147    (e.g. instead of ProcessPoolExecutor).148    """149    class DummyResult:150        def __init__(self, func, *args, **kwargs):151            self.func = func152            self.args = args153            self.kwargs = kwargs154 155        def result(self):156            return self.func(*self.args, **self.kwargs)157 158    def __init__(self, workers, mp_context=None):159        pass160 161    def submit(self, func, *args, **kwargs):162        return DummyPoolExecutor.DummyResult(func, *args, **kwargs)163 164    def __enter__(self):165        return self166 167    def __exit__(self, exc_type, exc_value, exc_tb):168        return169 170 171def get_pool_executor(num_workers: int, mp_context=None):172    return ProcessPoolExecutor(num_workers, mp_context) if num_workers > 1 else DummyPoolExecutor(1)173 174 175def length_to_mask(lengths: torch.Tensor, max_len: tp.Optional[int] = None) -> torch.Tensor:176    """Utility function to convert a tensor of sequence lengths to a mask (useful when working on padded sequences).177    For example: [3, 5] => [[1, 1, 1, 0, 0], [1, 1, 1, 1, 1]]178 179    Args:180        lengths (torch.Tensor): tensor with lengths181        max_len (int): can set the max length manually. Defaults to None.182    Returns:183        torch.Tensor: mask with 0s where there is pad tokens else 1s184    """185    assert len(lengths.shape) == 1, "Length shape should be 1 dimensional."186    final_length = lengths.max().item() if not max_len else max_len187    final_length = max(final_length, 1)  # if all seqs are of len zero we don't want a zero-size tensor188    return torch.arange(final_length, device=lengths.device)[None, :] < lengths[:, None]189 190 191def hash_trick(word: str, vocab_size: int) -> int:192    """Hash trick to pair each word with an index193 194    Args:195        word (str): word we wish to convert to an index196        vocab_size (int): size of the vocabulary197    Returns:198        int: index of the word in the embedding LUT199    """200    hash = int(hashlib.sha256(word.encode("utf-8")).hexdigest(), 16)201    return hash % vocab_size202 203 204def with_rank_rng(base_seed: int = 1234):205    """Decorator for a function so that the function will use a Random Number Generator206    whose state depend on the GPU rank. The original RNG state is restored upon returning.207 208    Args:209        base_seed (int): Random seed.210    """211    def _decorator(fun: tp.Callable):212        @wraps(fun)213        def _decorated(*args, **kwargs):214            state = torch.get_rng_state()215            seed = base_seed ^ flashy.distrib.rank()216            torch.manual_seed(seed)217            logger.debug('Rank dependent seed set to %d', seed)218            try:219                return fun(*args, **kwargs)220            finally:221                torch.set_rng_state(state)222                logger.debug('RNG state restored.')223        return _decorated224    return _decorator225 226 227def collate(tensors: tp.List[torch.Tensor], dim: int = 0) -> tp.Tuple[torch.Tensor, torch.Tensor]:228    """Get a list of tensors and collate them to a single tensor. according to the following logic:229    - `dim` specifies the time dimension which will be stacked and padded.230    - The output will contain 1 new dimension (dimension index 0) which will be the size of231    of the original list.232 233    Args:234        tensors (tp.List[torch.Tensor]): List of tensors to collate.235        dim (int): Dimension which will be stacked and padded.236    Returns:237        tp.Tuple[torch.Tensor, torch.Tensor]:238            torch.Tensor: Stacked and padded tensor. The output will contain 1 new dimension239                (dimension index 0) which will be the size of the original list.240            torch.Tensor: Tensor containing length of original tensor sizes (without padding).241    """242    tensors = [x.transpose(0, dim) for x in tensors]243    lens = torch.LongTensor([len(x) for x in tensors])244    padded_tensors = pad_sequence(tensors)245    padded_tensors = padded_tensors.transpose(0, 1)246    padded_tensors = padded_tensors.transpose(1, dim + 1)247    return padded_tensors, lens248 249 250# TODO: Move to flashy?251def copy_state(state: tp.Any, device: tp.Union[torch.device, str] = 'cpu',252               dtype: tp.Optional[torch.dtype] = None) -> tp.Any:253    if isinstance(state, torch.Tensor):254        if dtype is None or not state.is_floating_point():255            dtype = state.dtype256        return state.detach().to(device=device, dtype=dtype, copy=True)257    elif isinstance(state, dict):258        return {k: copy_state(v, device, dtype) for k, v in state.items()}259    elif isinstance(state, list):260        return [copy_state(v, device, dtype) for v in state]261 262 263# TODO: Move to flashy?264@contextmanager265def swap_state(model, state, **kwargs):266    old_state = copy_state(model.state_dict())267    model.load_state_dict(state, **kwargs)268    try:269        yield270    finally:271        model.load_state_dict(old_state)272 273 274@lru_cache(None)275def warn_once(logger, msg):276    """Warn about a given message only once."""277    logger.warning(msg)278 279 280def is_jsonable(x: tp.Any):281    """Check if an object can be serialized into a json:"""282    try:283        json.dumps(x)284        return True285    except (TypeError, OverflowError):286        return False287 288 289def load_clap_state_dict(clap_model, path: tp.Union[str, Path]):290    """Wrapper around state dict loading of CLAP model291    addressing compatibility issues between CLAP and AudioCraft292    HuggingFace transformer version.293    See: https://github.com/LAION-AI/CLAP/issues/118294    """295    from clap_module.factory import load_state_dict  # type: ignore296    pkg = load_state_dict(path)297    pkg.pop('text_branch.embeddings.position_ids', None)298    clap_model.model.load_state_dict(pkg)299