Upload 9 files

code/code_adaptations_audiocraft/audiocraft_lm_generate.py

@@ -0,0 +1,534 @@
+# 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 dataclasses import dataclass
+from functools import partial
+import logging
+import math
+import typing as tp
+
+import torch
+from torch import nn
+
+from audiocraft.utils import utils
+from audiocraft.modules.streaming import StreamingModule, State
+from audiocraft.modules.transformer import StreamingTransformer, create_norm_fn
+from audiocraft.modules.conditioners import (
+    ConditionFuser,
+    ClassifierFreeGuidanceDropout,
+    AttributeDropout,
+    ConditioningProvider,
+    ConditioningAttributes,
+    ConditionType,
+)
+from audiocraft.modules.codebooks_patterns import CodebooksPatternProvider
+from audiocraft.modules.activations import get_activation_fn
+
+
+logger = logging.getLogger(__name__)
+ConditionTensors = tp.Dict[str, ConditionType]
+CFGConditions = tp.Union[ConditionTensors, tp.Tuple[ConditionTensors, ConditionTensors]]
+
+
+def get_init_fn(method: str, input_dim: int, init_depth: tp.Optional[int] = None):
+    """LM layer initialization.
+    Inspired from xlformers: https://github.com/fairinternal/xlformers
+
+    Args:
+        method (str): Method name for init function. Valid options are:
+            'gaussian', 'uniform'.
+        input_dim (int): Input dimension of the initialized module.
+        init_depth (int, optional): Optional init depth value used to rescale
+            the standard deviation if defined.
+    """
+    # Compute std
+    std = 1 / math.sqrt(input_dim)
+    # Rescale with depth
+    if init_depth is not None:
+        std = std / math.sqrt(2 * init_depth)
+
+    if method == 'gaussian':
+        return partial(
+            torch.nn.init.trunc_normal_, mean=0.0, std=std, a=-3 * std, b=3 * std
+        )
+    elif method == 'uniform':
+        bound = math.sqrt(3) * std  # ensure the standard deviation is `std`
+        return partial(torch.nn.init.uniform_, a=-bound, b=bound)
+    else:
+        raise ValueError("Unsupported layer initialization method")
+
+
+def init_layer(m: nn.Module,
+               method: str,
+               init_depth: tp.Optional[int] = None,
+               zero_bias_init: bool = False):
+    """Wrapper around ``get_init_fn`` for proper initialization of LM modules.
+
+    Args:
+        m (nn.Module): Module to initialize.
+        method (str): Method name for the init function.
+        init_depth (int, optional): Optional init depth value used to rescale
+            the standard deviation if defined.
+        zero_bias_init (bool): Whether to initialize the bias to 0 or not.
+    """
+    if isinstance(m, nn.Linear):
+        init_fn = get_init_fn(method, m.in_features, init_depth=init_depth)
+        if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16:
+            weight = m.weight.float()
+            init_fn(weight)
+            m.weight.data[:] = weight.half()
+        else:
+            init_fn(m.weight)
+        if zero_bias_init and m.bias is not None:
+            nn.init.constant_(m.bias, 0)
+    elif isinstance(m, nn.Embedding):
+        init_fn = get_init_fn(method, m.embedding_dim, init_depth=None)
+        if m.weight.device.type == 'cpu' and m.weight.dtype == torch.float16:
+            weight = m.weight.float()
+            init_fn(weight)
+            m.weight.data[:] = weight.half()
+        else:
+            init_fn(m.weight)
+
+
+class ScaledEmbedding(nn.Embedding):
+    """Boost learning rate for embeddings (with `scale`).
+    """
+    def __init__(self, *args, lr=None, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.lr = lr
+
+    def make_optim_group(self):
+        group = {"params": list(self.parameters())}
+        if self.lr is not None:
+            group["lr"] = self.lr
+        return group
+
+
+@dataclass
+class LMOutput:
+    # The logits are already re-aligned with the input codes
+    # hence no extra shift is required, e.g. when computing CE
+    logits: torch.Tensor  # [B, K, T, card]
+    mask: torch.Tensor  # [B, K, T]
+
+
+class LMModel(StreamingModule):
+    """Transformer-based language model on multiple streams of codes.
+
+    Args:
+        pattern_provider (CodebooksPatternProvider): Pattern provider for codebook interleaving.
+        condition_provider (MusicConditioningProvider): Conditioning provider from metadata.
+        fuser (ConditionFuser): Fuser handling the fusing of conditions with language model input.
+        n_q (int): Number of parallel streams to model.
+        card (int): Cardinality, vocabulary size.
+        dim (int): Dimension of the transformer encoder.
+        num_heads (int): Number of heads for the transformer encoder.
+        hidden_scale (int): Scale for hidden feed forward dimension of the transformer encoder.
+        norm (str): Normalization method.
+        norm_first (bool): Use pre-norm instead of post-norm.
+        emb_lr (float, optional): Embedding-specific learning rate.
+        bias_proj (bool): Use bias for output projections.
+        weight_init (str, optional): Method for weight initialization.
+        depthwise_init (str, optional): Method for depthwise weight initialization.
+        zero_bias_init (bool): If true and bias in Linears, initialize bias to zeros.
+        cfg_dropout (float): Classifier-free guidance dropout.
+        cfg_coef (float): Classifier-free guidance coefficient.
+        attribute_dropout (dict): Attribute dropout probabilities.
+        two_step_cfg (bool): Whether to run classifier free-guidance with 2 distinct steps.
+        **kwargs: Additional parameters for the transformer encoder.
+    """
+    def __init__(self, pattern_provider: CodebooksPatternProvider, condition_provider: ConditioningProvider,
+                 fuser: ConditionFuser, n_q: int = 8, card: int = 1024, dim: int = 128, num_heads: int = 8,
+                 hidden_scale: int = 4, norm: str = 'layer_norm', norm_first: bool = False,
+                 emb_lr: tp.Optional[float] = None, bias_proj: bool = True,
+                 weight_init: tp.Optional[str] = None, depthwise_init: tp.Optional[str] = None,
+                 zero_bias_init: bool = False, cfg_dropout: float = 0, cfg_coef: float = 1.0,
+                 attribute_dropout: tp.Dict[str, tp.Dict[str, float]] = {}, two_step_cfg: bool = False,
+                 **kwargs):
+        super().__init__()
+        self.cfg_coef = cfg_coef
+        self.cfg_dropout = ClassifierFreeGuidanceDropout(p=cfg_dropout)
+        self.att_dropout = AttributeDropout(p=attribute_dropout)
+        self.condition_provider = condition_provider
+        self.fuser = fuser
+        self.card = card
+        embed_dim = self.card + 1
+        self.n_q = n_q
+        self.dim = dim
+        self.pattern_provider = pattern_provider
+        self.two_step_cfg = two_step_cfg
+        self.emb = nn.ModuleList([ScaledEmbedding(embed_dim, dim, lr=emb_lr) for _ in range(n_q)])
+        if 'activation' in kwargs:
+            kwargs['activation'] = get_activation_fn(kwargs['activation'])
+        self.transformer = StreamingTransformer(
+            d_model=dim, num_heads=num_heads, dim_feedforward=int(hidden_scale * dim),
+            norm=norm, norm_first=norm_first, **kwargs)
+        self.out_norm: tp.Optional[nn.Module] = None
+        if norm_first:
+            self.out_norm = create_norm_fn(norm, dim)
+        self.linears = nn.ModuleList([nn.Linear(dim, self.card, bias=bias_proj) for _ in range(n_q)])
+        self._init_weights(weight_init, depthwise_init, zero_bias_init)
+        self._fsdp: tp.Optional[nn.Module]
+        self.__dict__['_fsdp'] = None
+
+    def _init_weights(self, weight_init: tp.Optional[str], depthwise_init: tp.Optional[str], zero_bias_init: bool):
+        """Initialization of the transformer module weights.
+
+        Args:
+            weight_init (str, optional): Weight initialization strategy. See ``get_init_fn`` for valid options.
+            depthwise_init (str, optional): Depthwise initialization strategy. The following options are valid:
+                'current' where the depth corresponds to the current layer index or 'global' where the total number
+                of layer is used as depth. If not set, no depthwise initialization strategy is used.
+            zero_bias_init (bool): Whether to initialize bias to zero or not.
+        """
+        assert depthwise_init is None or depthwise_init in ['current', 'global']
+        assert depthwise_init is None or weight_init is not None, \
+            "If 'depthwise_init' is defined, a 'weight_init' method should be provided."
+        assert not zero_bias_init or weight_init is not None, \
+            "If 'zero_bias_init', a 'weight_init' method should be provided"
+
+        if weight_init is None:
+            return
+
+        for emb_layer in self.emb:
+            init_layer(emb_layer, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)
+
+        for layer_idx, tr_layer in enumerate(self.transformer.layers):
+            depth = None
+            if depthwise_init == 'current':
+                depth = layer_idx + 1
+            elif depthwise_init == 'global':
+                depth = len(self.transformer.layers)
+            init_fn = partial(init_layer, method=weight_init, init_depth=depth, zero_bias_init=zero_bias_init)
+            tr_layer.apply(init_fn)
+
+        for linear in self.linears:
+            init_layer(linear, method=weight_init, init_depth=None, zero_bias_init=zero_bias_init)
+
+    @property
+    def special_token_id(self) -> int:
+        return self.card
+
+    @property
+    def num_codebooks(self) -> int:
+        return self.n_q
+
+    def forward(self, sequence: torch.Tensor,
+                conditions: tp.List[ConditioningAttributes],
+                condition_tensors: tp.Optional[ConditionTensors] = None) -> torch.Tensor:
+        """Apply language model on sequence and conditions.
+        Given a tensor of sequence of shape [B, K, S] with K the number of codebooks and
+        S the sequence steps, return the logits with shape [B, card, K, S].
+
+        Args:
+            indices (torch.Tensor): Indices of the codes to model.
+            conditions (list of ConditioningAttributes): Conditions to use when modeling
+                the given codes. Note that when evaluating multiple time with the same conditioning
+                you should pre-compute those and pass them as `condition_tensors`.
+            condition_tensors (dict[str, ConditionType], optional): Pre-computed conditioning
+                tensors, see `conditions`.
+        Returns:
+            torch.Tensor: Logits.
+        """
+        B, K, S = sequence.shape
+        assert K == self.num_codebooks, "Sequence shape must match the specified number of codebooks"
+        input_ = sum([self.emb[k](sequence[:, k]) for k in range(K)])
+        if condition_tensors is None:
+            assert not self._is_streaming, "Conditions tensors should be precomputed when streaming."
+            # apply dropout modules
+            conditions = self.cfg_dropout(conditions)
+            conditions = self.att_dropout(conditions)
+            tokenized = self.condition_provider.tokenize(conditions)
+            # encode conditions and fuse, both have a streaming cache to not recompute when generating.
+            condition_tensors = self.condition_provider(tokenized)
+        else:
+            assert not conditions, "Shouldn't pass both conditions and condition_tensors."
+
+        input_, cross_attention_input = self.fuser(input_, condition_tensors)
+
+        out = self.transformer(input_, cross_attention_src=cross_attention_input)
+        if self.out_norm:
+            out = self.out_norm(out)
+        logits = torch.stack([self.linears[k](out) for k in range(K)], dim=1)  # [B, K, S, card]
+
+        # remove the prefix from the model outputs
+        if len(self.fuser.fuse2cond['prepend']) > 0:
+            logits = logits[:, :, -S:]
+
+        return logits  # [B, K, S, card]
+
+    def compute_predictions(
+            self, codes: torch.Tensor,
+            conditions: tp.List[ConditioningAttributes],
+            condition_tensors: tp.Optional[ConditionTensors] = None) -> LMOutput:
+        """Given an input tensor of codes [B, K, T] and list of conditions, runs the model
+        forward using the specified codes interleaving pattern.
+
+        Args:
+            codes (torch.Tensor): Input codes of shape [B, K, T] with B the batch size,
+                K the number of codebooks and T the number of timesteps.
+            conditions (list of ConditioningAttributes): conditionings to use when modeling
+                the given codes. Note that when evaluating multiple time with the same conditioning
+                you should pre-compute those and pass them as `condition_tensors`.
+            condition_tensors (dict[str, ConditionType], optional): pre-computed conditioning
+                tensors, see `conditions`.
+        Returns:
+            LMOutput: Language model outputs
+                logits (torch.Tensor) of shape [B, K, T, card] corresponding to the provided codes,
+                    i.e. the first item corresponds to logits to predict the first code, meaning that
+                    no additional shifting of codes and logits is required.
+                mask (torch.Tensor) of shape [B, K, T], mask over valid and invalid positions.
+                    Given the specified interleaving strategies, parts of the logits and codes should
+                    not be considered as valid predictions because of invalid context.
+        """
+        B, K, T = codes.shape
+        codes = codes.contiguous()
+        # map codes [B, K, T] into pattern sequence [B, K, S] using special_token_id for masked tokens
+        pattern = self.pattern_provider.get_pattern(T)
+        sequence_codes, sequence_indexes, sequence_mask = pattern.build_pattern_sequence(
+            codes, self.special_token_id, keep_only_valid_steps=True
+        )
+        # apply model on pattern sequence
+        model = self if self._fsdp is None else self._fsdp
+        logits = model(sequence_codes, conditions, condition_tensors)  # [B, K, S, card]
+        # map back the logits on pattern sequence to logits on original codes: [B, K, S, card] -> [B, K, T, card]
+        # and provide the corresponding mask over invalid positions of tokens
+        logits = logits.permute(0, 3, 1, 2)  # [B, card, K, S]
+        # note: we use nans as special token to make it obvious if we feed unexpected logits
+        logits, logits_indexes, logits_mask = pattern.revert_pattern_logits(
+            logits, float('nan'), keep_only_valid_steps=True
+        )
+        logits = logits.permute(0, 2, 3, 1)  # [B, K, T, card]
+        logits_mask = logits_mask[None, :, :].expand(B, -1, -1)  # [K, T] -> [B, K, T]
+        return LMOutput(logits, logits_mask)
+
+    def _sample_next_token(self,
+                           sequence: torch.Tensor,
+                           cfg_conditions: CFGConditions,
+                           unconditional_state: State,
+                           use_sampling: bool = False,
+                           temp: float = 1.0,
+                           top_k: int = 0,
+                           top_p: float = 0.0,
+                           cfg_coef: tp.Optional[float] = None) -> torch.Tensor:
+        """Sample next token from the model given a sequence and a set of conditions. The model supports
+        multiple sampling strategies (greedy sampling, softmax, top-k, top-p...).
+
+        Args:
+            sequence (torch.Tensor): Current sequence of shape [B, K, S]
+                with K corresponding to the number of codebooks and S the number of sequence steps.
+                S = 1 in streaming mode, except for the first step that contains a bigger prompt.
+            condition_tensors (dict[str, ConditionType): Set of conditions. If CFG is used,
+                should be twice the batch size, being the concatenation of the conditions + null conditions.
+            use_sampling (bool): Whether to use a sampling strategy or not.
+            temp (float): Sampling temperature.
+            top_k (int): K for "top-k" sampling.
+            top_p (float): P for "top-p" sampling.
+            cfg_coef (float, optional): classifier free guidance coefficient
+        Returns:
+            next_token (torch.Tensor): Next token tensor of shape [B, K, 1].
+        """
+        B = sequence.shape[0]
+        cfg_coef = self.cfg_coef if cfg_coef is None else cfg_coef
+        model = self if self._fsdp is None else self._fsdp
+        if self.two_step_cfg and cfg_conditions != {}:
+            assert isinstance(cfg_conditions, tuple), type(cfg_conditions)
+            condition_tensors, null_condition_tensors = cfg_conditions
+            cond_logits = model(sequence, conditions=[], condition_tensors=condition_tensors)
+            state = self.get_streaming_state()
+            self.set_streaming_state(unconditional_state)
+            uncond_logits = model(sequence, conditions=[], condition_tensors=null_condition_tensors)
+            unconditional_state.update(self.get_streaming_state())
+            self.set_streaming_state(state)
+            logits = uncond_logits + (cond_logits - uncond_logits) * self.cfg_coef
+        else:
+            assert isinstance(cfg_conditions, dict)
+            condition_tensors = cfg_conditions
+            if condition_tensors:
+                # Preparing for CFG, predicting both conditional and unconditional logits.
+                sequence = torch.cat([sequence, sequence], dim=0)
+            all_logits = model(
+                sequence,
+                conditions=[], condition_tensors=condition_tensors)
+            if condition_tensors:
+                cond_logits, uncond_logits = all_logits.split(B, dim=0)  # [B, K, T, card]
+                logits = uncond_logits + (cond_logits - uncond_logits) * cfg_coef
+            else:
+                logits = all_logits
+
+        logits = logits.permute(0, 1, 3, 2)  # [B, K, card, T]
+        logits = logits[..., -1]  # [B x K x card]
+
+        # Apply softmax for sampling if temp > 0. Else, do greedy sampling to avoid zero division error.
+        if use_sampling and temp > 0.0:
+            probs = torch.softmax(logits / temp, dim=-1)
+            if top_p > 0.0:
+                next_token = utils.sample_top_p(probs, p=top_p)
+            elif top_k > 0:
+                next_token = utils.sample_top_k(probs, k=top_k)
+            else:
+                next_token = utils.multinomial(probs, num_samples=1)
+        else:
+            next_token = torch.argmax(logits, dim=-1, keepdim=True)
+
+        return next_token
+
+    @torch.no_grad()
+    def generate(self,
+                 prompt: tp.Optional[torch.Tensor] = None,
+                 conditions: tp.List[ConditioningAttributes] = [],
+                 condition_tensors: tp.Optional[ConditionTensors] = None,
+                 num_samples: tp.Optional[int] = None,
+                 max_gen_len: int = 256,
+                 use_sampling: bool = True,
+                 temp: float = 1.0,
+                 top_k: int = 250,
+                 top_p: float = 0.0,
+                 cfg_coef: tp.Optional[float] = None,
+                 two_step_cfg: tp.Optional[bool] = None,
+                 remove_prompts: bool = False,
+                 check: bool = False,
+                 callback: tp.Optional[tp.Callable[[int, int], None]] = None) -> torch.Tensor:
+        """Generate tokens sampling from the model given a prompt or unconditionally. Generation can
+        be perform in a greedy fashion or using sampling with top K and top P strategies.
+
+        Args:
+            prompt (torch.Tensor, optional): Prompt tokens of shape [B, K, T].
+            conditions_tensors (list of ConditioningAttributes, optional): List of conditions.
+            num_samples (int, optional): Number of samples to generate when no prompt and no conditions are given.
+            max_gen_len (int): Maximum generation length.
+            use_sampling (bool): Whether to use a sampling strategy or not.
+            temp (float): Sampling temperature.
+            top_k (int): K for "top-k" sampling.
+            top_p (float): P for "top-p" sampling.
+            cfg_coeff (float, optional): Classifier-free guidance coefficient.
+            two_step_cfg (bool, optional): Whether to perform classifier-free guidance with two steps generation.
+            remove_prompts (bool): Whether to remove prompts from generation or not.
+            check (bool): Whether to apply further checks on generated sequence.
+            callback (Callback, optional): Callback function to report generation progress.
+        Returns:
+            torch.Tensor: Generated tokens.
+        """
+        assert not self.training, "generation shouldn't be used in training mode."
+        first_param = next(iter(self.parameters()))
+        device = first_param.device
+
+        # Checking all input shapes are consistent.
+        possible_num_samples = []
+        if num_samples is not None:
+            possible_num_samples.append(num_samples)
+        elif prompt is not None:
+            possible_num_samples.append(prompt.shape[0])
+        elif conditions:
+            possible_num_samples.append(len(conditions))
+        else:
+            possible_num_samples.append(1)
+        assert [x == possible_num_samples[0] for x in possible_num_samples], "Inconsistent inputs shapes"
+        num_samples = possible_num_samples[0]
+
+        # below we create set of conditions: one conditional and one unconditional
+        # to do that we merge the regular condition together with the null condition
+        # we then do 1 forward pass instead of 2.
+        # the reason for that is two-fold:
+        # 1. it is about x2 faster than doing 2 forward passes
+        # 2. avoid the streaming API treating the 2 passes as part of different time steps
+        # We also support doing two different passes, in particular to ensure that
+        # the padding structure is exactly the same between train and test.
+        # With a batch size of 1, this can be slower though.
+        cfg_conditions: CFGConditions
+        two_step_cfg = self.two_step_cfg if two_step_cfg is None else two_step_cfg
+        if condition_tensors is not None:
+            cfg_conditions = condition_tensors
+        elif conditions:
+            null_conditions = ClassifierFreeGuidanceDropout(p=1.0)(conditions)
+            if two_step_cfg:
+                cfg_conditions = (
+                    self.condition_provider(self.condition_provider.tokenize(conditions)),
+                    self.condition_provider(self.condition_provider.tokenize(null_conditions)),
+                )
+            else:
+                conditions = conditions + null_conditions
+                tokenized = self.condition_provider.tokenize(conditions)
+                cfg_conditions = self.condition_provider(tokenized)
+        else:
+            cfg_conditions = {}
+
+        if prompt is None:
+            assert num_samples > 0
+            prompt = torch.zeros((num_samples, self.num_codebooks, 0), dtype=torch.long, device=device)
+
+        B, K, T = prompt.shape
+        start_offset = T
+        assert start_offset < max_gen_len
+
+        pattern = self.pattern_provider.get_pattern(max_gen_len)
+        # this token is used as default value for codes that are not generated yet
+        unknown_token = -1
+
+        # we generate codes up to the max_gen_len that will be mapped to the pattern sequence
+        gen_codes = torch.full((B, K, max_gen_len), unknown_token, dtype=torch.long, device=device)
+        # filling the gen_codes with the prompt if needed
+        gen_codes[..., :start_offset] = prompt
+        # create the gen_sequence with proper interleaving from the pattern: [B, K, S]
+        gen_sequence, indexes, mask = pattern.build_pattern_sequence(gen_codes, self.special_token_id)
+        # retrieve the start_offset in the sequence:
+        # it is the first sequence step that contains the `start_offset` timestep
+        start_offset_sequence = pattern.get_first_step_with_timesteps(start_offset)
+        assert start_offset_sequence is not None
+
+        with self.streaming():
+            unconditional_state = self.get_streaming_state()
+            prev_offset = 0
+            gen_sequence_len = gen_sequence.shape[-1]  # gen_sequence shape is [B, K, S]
+            for offset in range(start_offset_sequence, gen_sequence_len):
+                # get current sequence (note that the streaming API is providing the caching over previous offsets)
+                curr_sequence = gen_sequence[..., prev_offset:offset]
+                curr_mask = mask[None, ..., prev_offset:offset].expand(B, -1, -1)
+                if check:
+                    # check coherence between mask and sequence
+                    assert (curr_sequence == torch.where(curr_mask, curr_sequence, self.special_token_id)).all()
+                    # should never happen as gen_sequence is filled progressively
+                    assert not (curr_sequence == unknown_token).any()
+                # sample next token from the model, next token shape is [B, K, 1]
+                next_token = self._sample_next_token(
+                    curr_sequence, cfg_conditions, unconditional_state, use_sampling, temp, top_k, top_p,
+                    cfg_coef=cfg_coef)
+                # ensure the tokens that should be masked are properly set to special_token_id
+                # as the model never output special_token_id
+                valid_mask = mask[..., offset:offset+1].expand(B, -1, -1)
+                next_token[~valid_mask] = self.special_token_id
+                # ensure we don't overwrite prompt tokens, we only write over unknown tokens
+                # (then mask tokens should be left as is as well, which is correct)
+                gen_sequence[..., offset:offset+1] = torch.where(
+                    gen_sequence[..., offset:offset+1] == unknown_token,
+                    next_token, gen_sequence[..., offset:offset+1]
+                )
+                prev_offset = offset
+                if callback is not None:
+                    callback(1 + offset - start_offset_sequence, gen_sequence_len - start_offset_sequence)
+        unconditional_state.clear()
+
+        # ensure sequence has been entirely filled
+        assert not (gen_sequence == unknown_token).any()
+        # ensure gen_sequence pattern and mask are matching
+        # which means the gen_sequence is valid according to the pattern
+        assert (
+            gen_sequence == torch.where(mask[None, ...].expand(B, -1, -1), gen_sequence, self.special_token_id)
+        ).all()
+        # get back the codes, trimming the prompt if needed and cutting potentially incomplete timesteps
+        out_codes, out_indexes, out_mask = pattern.revert_pattern_sequence(gen_sequence, special_token=unknown_token)
+
+        # sanity checks over the returned codes and corresponding masks
+        assert (out_codes[..., :max_gen_len] != unknown_token).all()
+        assert (out_mask[..., :max_gen_len] == 1).all()
+
+        out_start_offset = start_offset if remove_prompts else 0
+        out_codes = out_codes[..., out_start_offset:max_gen_len]
+
+        # ensure the returned codes are all valid
+        assert (out_codes >= 0).all() and (out_codes <= self.card).all()
+        return out_codes

code/code_adaptations_audiocraft/important_note.md

@@ -0,0 +1,6 @@
+### Important information about this ``audiocraft_lm_generate.py`
+
+As the audiocraft model does not allow to pass precalculated text embeddings into the `generate()` method directly, the 
+code within the audiocraft library was altered. `audiocraft_lm_generate.py` contains an updated version of
+`audiocraft/audiocraft/models/lm.py`. This is also the reason why audiocraft has to be installed from the local folder 
+`/audiocraft`. This file is just a backup to track the changes applied in the `lm.py`.

code/inference/gradio_app/app.py

@@ -0,0 +1,78 @@
+import gradio as gr
+import os
+import sys
+sys.path.insert(1, '..')
+import inference
+import torch
+
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+def generate_background_music(video_path, dataset, use_peft, musicgen_size):
+    print(f"Start generating background music for {video_path} with model \"{'peft' if use_peft else 'audiocraft'}_{dataset}_{musicgen_size}\"")
+
+    new_video_path = inference.generate_background_music(
+        video_path=video_path,
+        dataset=dataset,
+        musicgen_size=musicgen_size,
+        use_stereo=True,
+        use_peft=use_peft,
+        musicgen_temperature=1.0,
+        musicgen_guidance_scale=3.0,
+        top_k_sampling=250,
+        device=device
+    )
+    return gr.Video(new_video_path)
+
+
+interface = gr.Interface(fn=generate_background_music,
+                         inputs=[
+                                 gr.Video(
+                                          label="video input",
+                                          min_length=5,
+                                          max_length=20,
+                                          sources=['upload'],
+                                          show_download_button=True,
+                                          include_audio=True
+                                          ),
+                                 gr.Radio(["nature", "symmv"],
+                                          label="Video Encoder Version",
+                                          value="nature",
+                                          info="Choose one of the available Video Encoders."),
+                                 gr.Radio([False, True],
+                                          label="Use MusicGen Audio Decoder Model trained with PEFT",
+                                          value=False,
+                                          info="If set to 'True' the MusicGen Audio Decoder models trained with LoRA "
+                                               "(Low Rank Adaptation) are used. If set to 'False', the original "
+                                               "MusicGen models are used."),
+                                 gr.Radio(["small", "medium", "large"],
+                                          label="MusicGen Audio Decoder Size",
+                                          value="small",
+                                          info="Choose the size of the MusicGen audio decoder."),
+                                ],
+
+                         outputs=[gr.Video(label="video output")],
+                         examples=[
+                             [os.path.abspath("../../../videos/originals/n_1.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_2.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_3.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_4.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_5.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_6.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_7.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/n_8.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_1.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_2.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_3.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_4.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_5.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_6.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_7.mp4"), "nature", True, "small"],
+                             [os.path.abspath("../../../videos/originals/s_8.mp4"), "nature", True, "small"],
+                           ],
+                         cache_examples=False
+                         )
+
+if __name__ == "__main__":
+    interface.launch(
+        share=False
+    )

code/inference/inference.py

@@ -0,0 +1,117 @@
+from omegaconf import OmegaConf
+from peft import PeftConfig, get_peft_model 
+from audiocraft.models import MusicGen
+from moviepy.editor import AudioFileClip
+from code.inference.inference_utils import *
+import re
+import time
+
+re_file_name = re.compile('([^/]+$)')
+
+
+def generate_background_music(video_path: str,
+                              dataset: str,
+                              musicgen_size: str,
+                              use_stereo: bool,
+                              use_peft: bool,
+                              device: str,
+                              musicgen_temperature: float = 1.0,
+                              musicgen_guidance_scale: float = 3.0,
+                              top_k_sampling: int = 250) -> str:
+    start = time.time()
+    model_path = "../training/"
+    model_path += "models_peft" if use_peft else "models_audiocraft"
+    model_path += f"/{dataset}" + f"_{musicgen_size}"
+
+    conf = OmegaConf.load(model_path + '/configuration.yml')
+    use_sampling = True if top_k_sampling > 0 else False
+    video = mpe.VideoFileClip(video_path)
+
+    musicgen_model_id = "facebook/musicgen-" + "stereo-" if use_stereo else ""
+    musicgen_model_id += musicgen_size
+
+    result_dir = "./results"
+    os.makedirs(result_dir, exist_ok=True)
+
+    encoder_output_dimension = None
+    if "small" in conf.musicgen_model_id:
+        encoder_output_dimension = 1024
+    elif "medium" in conf.musicgen_model_id:
+        encoder_output_dimension = 1536
+    elif "large" in conf.musicgen_model_id:
+        encoder_output_dimension = 2048
+    assert encoder_output_dimension, f"Video Encoder output dimension could not be determined by {conf.musicgen_model_id}"
+
+    musicgen_model = MusicGen.get_pretrained(musicgen_model_id)
+    musicgen_model.lm.to(device)
+    musicgen_model.compression_model.to(device)
+    if use_peft:
+        peft_path = model_path + "/musicgen_peft_final"
+        peft_config = PeftConfig.from_pretrained(peft_path)
+        musicgen_model.lm = get_peft_model(musicgen_model.lm, peft_config)
+        musicgen_model.lm.load_adapter(peft_path, "default")
+
+    print("MusicGen Model loaded.")
+
+    video_to_t5 = VideoToT5(
+        video_extraction_framerate=conf.video_extraction_framerate,
+        encoder_input_dimension=conf.encoder_input_dimension,
+        encoder_output_dimension=encoder_output_dimension,
+        encoder_heads=conf.encoder_heads,
+        encoder_dim_feedforward=conf.encoder_dim_feedforward,
+        encoder_layers=conf.encoder_layers,
+        device=device
+    )
+
+    video_to_t5.load_state_dict(torch.load(model_path + "/lm_final.pt", map_location=device))
+    print("Video Encoder Model loaded.")
+
+    print("Starting Video Feature Extraction.")
+    video_embedding_t5 = video_to_t5(video_paths=[video_path])
+
+    condition_tensors = create_condition_tensors(
+        video_embeddings=video_embedding_t5,
+        batch_size=1,
+        video_extraction_framerate=video_to_t5.video_extraction_framerate,
+        device=device
+    )
+
+    musicgen_model.generation_params = {
+        'max_gen_len': int(video.duration * musicgen_model.frame_rate),
+        'use_sampling': use_sampling,
+        'temp': musicgen_temperature,
+        'cfg_coef': musicgen_guidance_scale,
+        'two_step_cfg': False,
+    }
+    if use_sampling:
+        musicgen_model.generation_params['top_k'] = 250
+
+    print("Starting Audio Generation.")
+    prompt_tokens = None
+    with torch.no_grad():
+        with musicgen_model.autocast:
+            gen_tokens = musicgen_model.lm.generate(prompt_tokens, [], condition_tensors, callback=None,
+                                                    **musicgen_model.generation_params)
+        gen_audio = musicgen_model.compression_model.decode(gen_tokens)
+
+    end = time.time()
+    print("Elapsed time for generation: " + str(end - start))
+
+    _, video_file_name = os.path.split(video_path)
+    video_file_name = video_file_name[:-4]  # remove .mp4
+
+    re_result = re_file_name.search(video_file_name)  # get video file name
+    result_path = f"{'peft' if use_peft else 'audiocraft'}_{dataset}_{musicgen_size}_{re_result.group(1)}"
+    audio_result_path = f"{result_dir}/tmp.wav"
+    video_result_path = f"{result_dir}/{result_path}_video.mp4"
+
+    gen_audio = torch.squeeze(gen_audio.detach().cpu())  # remove mini-batch dimension, move to CPU for saving
+    sample_rate = musicgen_model.sample_rate
+    torchaudio.save(audio_result_path, gen_audio, sample_rate)
+    audio_file_clip = AudioFileClip(audio_result_path)
+    video.audio = audio_file_clip
+
+    print("Rendering Video.")
+    video.write_videofile(video_result_path)
+
+    return video_result_path

code/inference/inference_utils.py

@@ -0,0 +1,283 @@
+from torch.utils.data import Dataset
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import torchaudio
+import os
+import logging
+from torchvision.models import resnet50, ResNet50_Weights, resnet152, resnet18, resnet34, ResNet152_Weights
+from PIL import Image
+from time import strftime
+import math
+import numpy as np
+import moviepy.editor as mpe
+
+
+class VideoDataset(Dataset):
+    def __init__(self, data_dir):
+        self.data_dir = data_dir
+        self.data_map = []
+
+        dir_map = os.listdir(data_dir)
+        for d in dir_map:
+            name, extension = os.path.splitext(d)
+            if extension == ".mp4":
+                self.data_map.append({"video": os.path.join(data_dir, d)})
+
+    def __len__(self):
+        return len(self.data_map)
+
+    def __getitem__(self, idx):
+        return self.data_map[idx]["video"]
+
+
+# input: video_path, output: wav_music
+class VideoToT5(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.video_extraction_framerate = video_extraction_framerate
+        self.video_feature_extractor = VideoFeatureExtractor(video_extraction_framerate=video_extraction_framerate,
+                                                             device=device)
+        self.video_encoder = VideoEncoder(
+            device,
+            encoder_input_dimension,
+            encoder_output_dimension,
+            encoder_heads,
+            encoder_dim_feedforward,
+            encoder_layers
+        )
+
+    def forward(self, video_paths: [str]):
+        image_embeddings = []
+        for video_path in video_paths:
+            video = mpe.VideoFileClip(video_path)
+            video_embedding = self.video_feature_extractor(video)
+            image_embeddings.append(video_embedding)
+        video_embedding = torch.stack(
+            image_embeddings)  # resulting shape: [batch_size, video_extraction_framerate, resnet_output_dimension]
+        # not used, gives worse results!
+        # video_embeddings = torch.mean(video_embeddings, 0, True)  # average out all image embedding to one video embedding
+
+        t5_embeddings = self.video_encoder(video_embedding)  # T5 output: [batch_size, num_tokens,
+        # t5_embedding_size]
+        return t5_embeddings
+
+
+class VideoEncoder(nn.Module):
+    def __init__(self,
+                 device: str,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.device = device
+        self.encoder = (nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=encoder_input_dimension,
+                nhead=encoder_heads,
+                dim_feedforward=encoder_dim_feedforward
+            ),
+            num_layers=encoder_layers,
+        )
+        ).to(device)
+
+        # linear layer to match T5 embedding dimension
+        self.linear = (nn.Linear(
+            in_features=encoder_input_dimension,
+            out_features=encoder_output_dimension)
+                       .to(device))
+
+    def forward(self, x):
+        assert x.dim() == 3
+        x = torch.transpose(x, 0, 1)  # encoder expects [sequence_length, batch_size, embedding_dimension]
+        x = self.encoder(x)  # encoder forward pass
+        x = self.linear(x)  # forward pass through the linear layer
+        x = torch.transpose(x, 0, 1)  # shape: [batch_size, sequence_length, embedding_dimension]
+        return x
+
+
+class VideoFeatureExtractor(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int = 1,
+                 resnet_output_dimension: int = 2048):
+        super().__init__()
+        self.device = device
+
+        # using a ResNet trained on ImageNet
+        self.resnet = resnet50(weights="IMAGENET1K_V2").eval()
+        self.resnet = torch.nn.Sequential(*(list(self.resnet.children())[:-1])).to(device)  # remove ResNet layer
+        self.resnet_preprocessor = ResNet50_Weights.DEFAULT.transforms().to(device)
+        self.video_extraction_framerate = video_extraction_framerate  # setting the fps at which the video is processed
+        self.positional_encoder = PositionalEncoding(resnet_output_dimension).to(device)
+
+    def forward(self, video: mpe.VideoFileClip):
+        embeddings = []
+        for i in range(0, 30 * self.video_extraction_framerate):
+            i = video.get_frame(i)  # get frame as numpy array
+            i = Image.fromarray(i)  # create PIL image from numpy array
+            i = self.resnet_preprocessor(i)  # preprocess image
+            i = i.to(self.device)
+            i = i.unsqueeze(0)  # adding a batch dimension
+            i = self.resnet(i).squeeze()  # ResNet forward pass
+            i = i.squeeze()
+            embeddings.append(i)  # collect embeddings
+
+        embeddings = torch.stack(embeddings)  # concatenate all frame embeddings into one video embedding
+        embeddings = embeddings.unsqueeze(1)
+        embeddings = self.positional_encoder(embeddings)  # apply positional encoding with a sequence length of 30
+        embeddings = embeddings.squeeze()
+        return embeddings
+
+
+# from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, dropout: float = 0.1, max_length: int = 5000):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        position = torch.arange(30).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(30, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)
+
+
+def freeze_model(model: nn.Module):
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+
+def split_dataset_randomly(dataset, validation_split: float, test_split: float, seed: int = None):
+    dataset_size = len(dataset)
+    indices = list(range(dataset_size))
+    datapoints_validation = int(np.floor(validation_split * dataset_size))
+    datapoints_testing = int(np.floor(test_split * dataset_size))
+
+    if seed:
+        np.random.seed(seed)
+
+    np.random.shuffle(indices)  # in-place operation
+    training = indices[datapoints_validation + datapoints_testing:]
+    validation = indices[datapoints_validation:datapoints_testing + datapoints_validation]
+    testing = indices[:datapoints_testing]
+
+    assert len(validation) == datapoints_validation, "Validation set length incorrect"
+    assert len(testing) == datapoints_testing, "Testing set length incorrect"
+    assert len(training) == dataset_size - (datapoints_testing + datapoints_testing), "Training set length incorrect"
+    assert not any([item in training for item in validation]), "Training and Validation overlap"
+    assert not any([item in training for item in testing]), "Training and Testing overlap"
+    assert not any([item in validation for item in testing]), "Validation and Testing overlap"
+
+    return training, validation, testing
+
+
+### private function from audiocraft.solver.musicgen.py => _compute_cross_entropy
+def compute_cross_entropy(logits: torch.Tensor, targets: torch.Tensor, mask: 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 = []
+    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 generate_audio_codes(audio_paths: [str],
+                         audiocraft_compression_model: torch.nn.Module,
+                         device: str) -> torch.Tensor:
+    audio_duration = 30
+    encodec_sample_rate = audiocraft_compression_model.sample_rate
+
+    torch_audios = []
+    for audio_path in audio_paths:
+        wav, original_sample_rate = torchaudio.load(audio_path)  # load audio from file
+        wav = torchaudio.functional.resample(wav, original_sample_rate,
+                                             encodec_sample_rate)  # cast audio to model sample rate
+        wav = wav[:, :encodec_sample_rate * audio_duration]  # enforce an exact audio length of 30 seconds
+
+        assert len(wav.shape) == 2, f"audio data is not of shape [channels, duration]"
+        assert wav.shape[0] == 2, "audio data should be in stereo, but has not 2 channels"
+
+        torch_audios.append(wav)
+
+    torch_audios = torch.stack(torch_audios)
+    torch_audios = torch_audios.to(device)
+
+    with torch.no_grad():
+        gen_audio = audiocraft_compression_model.encode(torch_audios)
+
+    codes, scale = gen_audio
+    assert scale is None
+
+    return codes
+
+
+def create_condition_tensors(
+        video_embeddings: torch.Tensor,
+        batch_size: int,
+        video_extraction_framerate: int,
+        device: str
+):
+    # model T5 mask
+    mask = torch.ones((batch_size, video_extraction_framerate * 30), dtype=torch.int).to(device)
+
+    condition_tensors = {
+        'description': (video_embeddings, mask)
+    }
+    return condition_tensors
+
+
+def get_current_timestamp():
+    return strftime("%Y_%m_%d___%H_%M_%S")
+
+
+def configure_logging(output_dir: str, filename: str, log_level):
+    # create logs folder, if not existing
+    os.makedirs(output_dir, exist_ok=True)
+    level = getattr(logging, log_level)
+    file_path = output_dir + "/" + filename
+    logging.basicConfig(filename=file_path, encoding='utf-8', level=level)
+    logger = logging.getLogger()
+    # only add a StreamHandler if it is not present yet
+    if len(logger.handlers) <= 1:
+        logger.addHandler(logging.StreamHandler())

code/training/training.py

@@ -0,0 +1,239 @@
+from audiocraft.models import MusicGen
+from torch.optim import AdamW
+from torch.utils.data import DataLoader
+import wandb
+from omegaconf import OmegaConf
+from training_utils import *
+import json
+import os
+import sys
+import shutil
+import copy
+from peft import get_peft_model, LoraConfig
+
+
+def get_trainable_parameters(model):
+    """
+    Prints the number of trainable parameters in the model.
+    """
+    trainable_params = 0
+    all_param = 0
+    for _, param in model.named_parameters():
+        all_param += param.numel()
+        if param.requires_grad:
+            trainable_params += param.numel()
+    return f"trainable params: {trainable_params} || all params: {all_param} || trainable%: {100 * trainable_params / all_param}"
+
+
+def save_model(step: int, is_final_save: bool):
+    logging.info(f"Model save at step {int(step)}")
+    torch.save(video_to_t5.state_dict(),
+               f"{model_save_path}/lm_{'final' if is_final_save else int(step)}.pt")
+
+    if conf.use_peft:
+        musicgen_model.lm.save_pretrained(f"{model_save_path}/musicgen_peft_{'final' if is_final_save else int(step)}")
+
+
+conf = OmegaConf.load('training_conf.yml')
+start_timestamp = get_current_timestamp()
+model_save_path = f"./{'models_peft' if conf.use_peft else 'models_audiocraft'}/{conf.output_dir_name}"
+
+if os.path.isdir(model_save_path):
+    i = input(f"Model output directory {model_save_path} already exists, overwrite directory? confirm with [y]\n")
+    if i == "y" or i == "yes":
+        shutil.rmtree(model_save_path)
+    else:
+        print("Aborting.")
+        sys.exit()
+
+configure_logging(model_save_path, f"{start_timestamp}.log", conf.log_level)
+os.makedirs(model_save_path, exist_ok=True)
+
+if conf.use_wandb:
+    wandb = wandb.init(project=conf.wandb_project_name,
+                       config=OmegaConf.to_container(conf))
+    logging.info(f"Wandb project_name: {conf.wandb_project_name}, run_id: {wandb.id}, run_name: {wandb.id}")
+
+logging.info("Start Training")
+musicgen_model = MusicGen.get_pretrained(conf.musicgen_model_id, device=conf.device)
+musicgen_model.compression_model = musicgen_model.compression_model.to(conf.device)
+musicgen_model.lm = musicgen_model.lm.to(conf.device)
+musicgen_model.lm = musicgen_model.lm.train()
+
+encoder_output_dimension = None
+if "small" in conf.musicgen_model_id:
+    encoder_output_dimension = 1024
+elif "medium" in conf.musicgen_model_id:
+    encoder_output_dimension = 1536
+elif "large" in conf.musicgen_model_id:
+    encoder_output_dimension = 2048
+assert encoder_output_dimension, f"Video Encoder output dimension could not be determined by {conf.musicgen_model_id}"
+
+# initialize video-to-text model
+video_to_t5 = VideoToT5(video_extraction_framerate=conf.video_extraction_framerate,
+                        encoder_input_dimension=conf.encoder_input_dimension,
+                        encoder_output_dimension=encoder_output_dimension,
+                        encoder_heads=conf.encoder_heads,
+                        encoder_dim_feedforward=conf.encoder_dim_feedforward,
+                        encoder_layers=conf.encoder_layers,
+                        device=conf.device)
+
+# freeze all model layers that except the video-to-text encoder
+freeze_model(video_to_t5.video_feature_extractor)
+freeze_model(musicgen_model.compression_model)
+if not conf.use_peft:
+    freeze_model(musicgen_model.lm)
+
+logging.info(f"Trainable parameters video_to_t5: {get_trainable_parameters(video_to_t5)}")
+
+if conf.use_peft:
+    lora_config = LoraConfig(
+        r=conf.lora_r,
+        lora_alpha=conf.lora_alpha,
+        target_modules=["out_proj", "linear1", "linear2"],
+        lora_dropout=conf.lora_dropout,
+        bias="none",
+        modules_to_save=["classifier"]
+    )
+
+    logging.info(f"Trainable parameters MusicGen before LoRA: {get_trainable_parameters(musicgen_model.lm)}")
+    musicgen_model.lm = get_peft_model(musicgen_model.lm, lora_config)
+    logging.info(f"Trainable parameters MusicGen with LoRA: {get_trainable_parameters(musicgen_model.lm)}")
+
+logging.info(f"Training on {conf.musicgen_model_id}")
+
+# create dataset train and validation split
+dataset = VideoDataset(conf.dataset_video_folder)
+train_indices, validation_indices, test_indices = split_dataset_randomly(dataset,
+                                                                         conf.dataset_validation_split,
+                                                                         conf.dataset_test_split,
+                                                                         seed=conf.dataset_shuffling_seed)
+
+train_dataset = copy.copy(dataset)
+train_dataset.data_map = [dataset.data_map[i] for i in train_indices]
+validation_dataset = copy.copy(dataset)
+validation_dataset.data_map = [dataset.data_map[i] for i in validation_indices]
+test_dataset = copy.copy(dataset)
+test_dataset.data_map = [dataset.data_map[i] for i in test_indices]
+train_dataloader = DataLoader(train_dataset, batch_size=conf.batch_size)
+validation_dataloader = DataLoader(validation_dataset, batch_size=conf.batch_size)
+test_dataloader = DataLoader(test_dataset, batch_size=1)
+
+with open(f"{model_save_path}/dataset_split.json", 'w') as f:
+    json.dump({
+        "dataset": dataset.data_map,
+        "training": train_dataloader.dataset.data_map,
+        "validation": validation_dataloader.dataset.data_map,
+    }, f)
+
+with open(f"{model_save_path}/configuration.yml", 'w') as f:
+    OmegaConf.save(conf, f)
+
+logging.info(f"Video path: {conf.dataset_video_folder}, "
+             f"Audio path: {conf.dataset_audio_folder} with {len(dataset)} examples, "
+             f"Batch Size: {conf.batch_size}.")
+
+optimizer = AdamW(
+    video_to_t5.video_encoder.parameters(),
+    betas=(conf.beta1, conf.beta2),
+    weight_decay=conf.weight_decay,
+    lr=conf.learning_rate
+)
+
+
+def forward_pass(video_path_list: [str]):
+    optimizer.zero_grad()
+
+    # get corresponding audio for the video data
+    audio_paths = []
+    for video_path in video_paths:
+        # load corresponding audio file
+        _, video_file_name = os.path.split(video_path)
+        video_file_name = video_file_name[:-4]  # remove .mp4
+        if conf.use_demucs_folder_structure:
+            audio_path = f"{conf.dataset_audio_folder}/htdemucs/{video_file_name}/no_vocals.wav"
+        else:
+            audio_path = f"{conf.dataset_audio_folder}/{video_file_name}.wav"
+        audio_paths.append(audio_path)
+
+    # batch encode audio data
+    audio_batches = generate_audio_codes(audio_paths=audio_paths,
+                                         audiocraft_compression_model=musicgen_model.compression_model,
+                                         device=conf.device)
+
+    # batch encode video data
+    video_embedding_batches = video_to_t5(video_path_list)
+
+    condition_tensors = create_condition_tensors(video_embedding_batches,
+                                                 conf.batch_size,
+                                                 video_to_t5.video_extraction_framerate,
+                                                 device=conf.device)
+
+    # forward pass with MusicGen
+    with musicgen_model.autocast:
+        musicgen_output = musicgen_model.lm.compute_predictions(
+            codes=audio_batches,
+            conditions=[],
+            condition_tensors=condition_tensors
+        )
+    loss, _ = compute_cross_entropy(logits=musicgen_output.logits,
+                                    targets=audio_batches,
+                                    mask=musicgen_output.mask)
+    return musicgen_output, loss
+
+
+training_step = 0
+for epoch in range(conf.num_epochs):
+    epoch_training_loss = []
+    epoch_validation_loss = []
+    logging.info("Starting next Epoch.")
+    for batch_idx, video_paths in enumerate(train_dataloader):
+        _, training_loss = forward_pass(video_paths)
+        epoch_training_loss.append(training_loss)
+
+        training_loss.backward()
+        # gradient clipping
+        torch.nn.utils.clip_grad_norm_(musicgen_model.lm.parameters(), conf.gradient_clipping)
+        optimizer.step()
+        training_step += 1
+
+        # update metrics
+        if conf.use_wandb:
+            wandb.log(dict(training_loss=training_loss.item()))
+        logging.info(
+            f"Epoch: {epoch + 1}/{conf.num_epochs}, "
+            f"Batch: {batch_idx}/{len(train_dataloader)}, "
+            f"Loss: {training_loss.item()}"
+        )
+
+    # save model after each epoch
+    save_model(training_step, False)
+
+    # testing
+    logging.info("Start Validation.")
+    with torch.no_grad():
+        for batch_idx, video_paths in enumerate(validation_dataloader):
+            _, validation_loss = forward_pass(video_paths)
+            epoch_validation_loss.append(validation_loss)
+            if conf.use_wandb:
+                wandb.log(dict(validation_loss=validation_loss.item()))
+            logging.info(
+                f"Epoch: {epoch + 1}/{conf.num_epochs}, "
+                f"Batch: {batch_idx}/{len(validation_dataloader)}, "
+                f"Loss: {validation_loss.item()}"
+            )
+    logging.info(
+        f"Epoch results: epoch_training_loss {epoch_training_loss}, epoch_validation_loss {epoch_validation_loss}")
+save_model(training_step, True)
+logging.info(f"Finished Training. Start Testing")
+with torch.no_grad():
+    for batch_idx, video_paths in enumerate(test_dataloader):
+        _, testing_loss = forward_pass(video_paths)
+        if conf.use_wandb:
+            wandb.log(dict(testing_loss=testing_loss.item()))
+        logging.info(
+            f"Epoch: {epoch + 1}/{conf.num_epochs}, "
+            f"Batch: {batch_idx}/{len(test_dataloader)}, "
+            f"Loss: {testing_loss.item()}"
+        )
+logging.info(f"Finished Testing.")

code/training/training_conf.yml

@@ -0,0 +1,47 @@
+# LOGGING
+log_level: "INFO"
+
+# TRAINING LOOP
+num_epochs: 4
+batch_size: 2
+save_step: 800
+device: "cuda"
+
+# WANDB
+use_wandb: true
+wandb_project_name: "pe_30"
+
+# OPTIMIZER
+gradient_clipping: 1.0
+beta1: 0.9  # from audiocraft
+beta2: 0.95  # from audiocraft
+weight_decay: 0.1  # from audiocraft
+learning_rate: 1e-5
+
+# VIDEO ENCODER
+video_extraction_framerate: 1
+encoder_input_dimension: 2048  # determined by ResNet
+encoder_dim_feedforward: 8192
+encoder_heads: 16
+encoder_layers: 6
+
+# MUSICGEN
+musicgen_model_id: "facebook/musicgen-stereo-small"
+
+# PEFT PARAMETERS
+use_peft: false
+lora_r: 16
+lora_alpha: 16
+lora_dropout: 0.1
+
+# OUTPUT FOLDER
+output_dir_name: "nature_small"
+
+# DATASET
+dataset_video_folder: "../../datasets/nature/videos_30sec"
+dataset_audio_folder: "../../datasets/nature/audio_30sec"
+#dataset_audio_folder: "../../datasets/symmv/audio_30sec_separated"
+use_demucs_folder_structure: false
+dataset_validation_split: 0.1
+dataset_test_split: 0.1
+dataset_shuffling_seed: 42

code/training/training_utils.py

@@ -0,0 +1,282 @@
+from torch.utils.data import Dataset
+import torch
+from torch import nn, Tensor
+import torch.nn.functional as F
+import torchaudio
+import os
+import logging
+from torchvision.models import resnet50, ResNet50_Weights, resnet152, resnet18, resnet34, ResNet152_Weights
+from PIL import Image
+from time import strftime
+import math
+import numpy as np
+import moviepy.editor as mpe
+
+
+class VideoDataset(Dataset):
+    def __init__(self, data_dir):
+        self.data_dir = data_dir
+        self.data_map = []
+
+        dir_map = os.listdir(data_dir)
+        for d in dir_map:
+            name, extension = os.path.splitext(d)
+            if extension == ".mp4":
+                self.data_map.append({"video": os.path.join(data_dir, d)})
+
+    def __len__(self):
+        return len(self.data_map)
+
+    def __getitem__(self, idx):
+        return self.data_map[idx]["video"]
+
+
+# input: video_path, output: wav_music
+class VideoToT5(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.video_extraction_framerate = video_extraction_framerate
+        self.video_feature_extractor = VideoFeatureExtractor(video_extraction_framerate=video_extraction_framerate,
+                                                             device=device)
+        self.video_encoder = VideoEncoder(
+            device,
+            encoder_input_dimension,
+            encoder_output_dimension,
+            encoder_heads,
+            encoder_dim_feedforward,
+            encoder_layers
+        )
+
+    def forward(self, video_paths: [str]):
+        image_embeddings = []
+        for video_path in video_paths:
+            video = mpe.VideoFileClip(video_path)
+            video_embedding = self.video_feature_extractor(video)
+            image_embeddings.append(video_embedding)
+        video_embedding = torch.stack(
+            image_embeddings)  # resulting shape: [batch_size, video_extraction_framerate, resnet_output_dimension]
+        # not used, gives worse results!
+        # video_embeddings = torch.mean(video_embeddings, 0, True)  # average out all image embedding to one video embedding
+
+        t5_embeddings = self.video_encoder(video_embedding)  # T5 output: [batch_size, num_tokens,
+        # t5_embedding_size]
+        return t5_embeddings
+
+
+class VideoEncoder(nn.Module):
+    def __init__(self,
+                 device: str,
+                 encoder_input_dimension: int,
+                 encoder_output_dimension: int,
+                 encoder_heads: int,
+                 encoder_dim_feedforward: int,
+                 encoder_layers: int
+                 ):
+        super().__init__()
+        self.device = device
+        self.encoder = (nn.TransformerEncoder(
+            nn.TransformerEncoderLayer(
+                d_model=encoder_input_dimension,
+                nhead=encoder_heads,
+                dim_feedforward=encoder_dim_feedforward
+            ),
+            num_layers=encoder_layers,
+        )
+        ).to(device)
+
+        # linear layer to match T5 embedding dimension
+        self.linear = (nn.Linear(
+            in_features=encoder_input_dimension,
+            out_features=encoder_output_dimension)
+                       .to(device))
+
+    def forward(self, x):
+        assert x.dim() == 3
+        x = torch.transpose(x, 0, 1)  # encoder expects [sequence_length, batch_size, embedding_dimension]
+        x = self.encoder(x)  # encoder forward pass
+        x = self.linear(x)  # forward pass through the linear layer
+        x = torch.transpose(x, 0, 1)  # shape: [batch_size, sequence_length, embedding_dimension]
+        return x
+
+
+class VideoFeatureExtractor(nn.Module):
+    def __init__(self,
+                 device: str,
+                 video_extraction_framerate: int = 1,
+                 resnet_output_dimension: int = 2048):
+        super().__init__()
+        self.device = device
+
+        # using a ResNet trained on ImageNet
+        self.resnet = resnet50(weights="IMAGENET1K_V2").eval()
+        self.resnet = torch.nn.Sequential(*(list(self.resnet.children())[:-1])).to(device)  # remove ResNet layer
+        self.resnet_preprocessor = ResNet50_Weights.DEFAULT.transforms().to(device)
+        self.video_extraction_framerate = video_extraction_framerate  # setting the fps at which the video is processed
+        self.positional_encoder = PositionalEncoding(resnet_output_dimension).to(device)
+
+    def forward(self, video: mpe.VideoFileClip):
+        embeddings = []
+        for i in range(0, 30 * self.video_extraction_framerate):
+            i = video.get_frame(i)  # get frame as numpy array
+            i = Image.fromarray(i)  # create PIL image from numpy array
+            i = self.resnet_preprocessor(i)  # preprocess image
+            i = i.to(self.device)
+            i = i.unsqueeze(0)  # adding a batch dimension
+            i = self.resnet(i).squeeze()  # ResNet forward pass
+            i = i.squeeze()
+            embeddings.append(i)  # collect embeddings
+
+        embeddings = torch.stack(embeddings)  # concatenate all frame embeddings into one video embedding
+        embeddings = embeddings.unsqueeze(1)
+        embeddings = self.positional_encoder(embeddings)  # apply positional encoding with a sequence length of 30
+        embeddings = embeddings.squeeze()
+        return embeddings
+
+
+# from https://pytorch.org/tutorials/beginner/transformer_tutorial.html
+class PositionalEncoding(nn.Module):
+    def __init__(self, d_model: int, dropout: float = 0.1, max_length: int = 30):
+        super().__init__()
+        self.dropout = nn.Dropout(p=dropout)
+        position = torch.arange(30).unsqueeze(1)
+        div_term = torch.exp(torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
+        pe = torch.zeros(30, 1, d_model)
+        pe[:, 0, 0::2] = torch.sin(position * div_term)
+        pe[:, 0, 1::2] = torch.cos(position * div_term)
+        self.register_buffer('pe', pe)
+
+    def forward(self, x: Tensor) -> Tensor:
+        x = x + self.pe[:x.size(0)]
+        return self.dropout(x)
+
+
+def freeze_model(model: nn.Module):
+    for param in model.parameters():
+        param.requires_grad = False
+    model.eval()
+
+
+def split_dataset_randomly(dataset, validation_split: float, test_split: float, seed: int = None):
+    dataset_size = len(dataset)
+    indices = list(range(dataset_size))
+    datapoints_validation = int(np.floor(validation_split * dataset_size))
+    datapoints_testing = int(np.floor(test_split * dataset_size))
+
+    if seed:
+        np.random.seed(seed)
+
+    np.random.shuffle(indices)  # in-place operation
+    training = indices[datapoints_validation + datapoints_testing:]
+    validation = indices[datapoints_validation:datapoints_testing + datapoints_validation]
+    testing = indices[:datapoints_testing]
+
+    assert len(validation) == datapoints_validation, "Validation set length incorrect"
+    assert len(testing) == datapoints_testing, "Testing set length incorrect"
+    assert len(training) == dataset_size - (datapoints_testing + datapoints_testing), "Training set length incorrect"
+    assert not any([item in training for item in validation]), "Training and Validation overlap"
+    assert not any([item in training for item in testing]), "Training and Testing overlap"
+    assert not any([item in validation for item in testing]), "Validation and Testing overlap"
+
+    return training, validation, testing
+
+
+### private function from audiocraft.solver.musicgen.py => _compute_cross_entropy
+def compute_cross_entropy(logits: torch.Tensor, targets: torch.Tensor, mask: 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 = []
+    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 generate_audio_codes(audio_paths: [str],
+                         audiocraft_compression_model: torch.nn.Module,
+                         device: str) -> torch.Tensor:
+    audio_duration = 30
+    encodec_sample_rate = audiocraft_compression_model.sample_rate
+
+    torch_audios = []
+    for audio_path in audio_paths:
+        wav, original_sample_rate = torchaudio.load(audio_path)  # load audio from file
+        wav = torchaudio.functional.resample(wav, original_sample_rate,
+                                             encodec_sample_rate)  # cast audio to model sample rate
+        wav = wav[:, :encodec_sample_rate * audio_duration]  # enforce an exact audio length of 30 seconds
+
+        assert len(wav.shape) == 2, f"audio data is not of shape [channels, duration]"
+        assert wav.shape[0] == 2, "audio data should be in stereo, but has not 2 channels"
+
+        torch_audios.append(wav)
+
+    torch_audios = torch.stack(torch_audios)
+    torch_audios = torch_audios.to(device)
+
+    with torch.no_grad():
+        gen_audio = audiocraft_compression_model.encode(torch_audios)
+
+    codes, scale = gen_audio
+    assert scale is None
+
+    return codes
+
+
+def create_condition_tensors(
+        video_embeddings: torch.Tensor,
+        batch_size: int,
+        video_extraction_framerate: int,
+        device: str
+):
+    mask = torch.ones((batch_size, video_extraction_framerate * 30), dtype=torch.int).to(device)
+
+    condition_tensors = {
+        'description': (video_embeddings, mask)
+    }
+    return condition_tensors
+
+
+def get_current_timestamp():
+    return strftime("%Y_%m_%d___%H_%M_%S")
+
+
+def configure_logging(output_dir: str, filename: str, log_level):
+    # create logs folder, if not existing
+    os.makedirs(output_dir, exist_ok=True)
+    level = getattr(logging, log_level)
+    file_path = output_dir + "/" + filename
+    logging.basicConfig(filename=file_path, encoding='utf-8', level=level)
+    logger = logging.getLogger()
+    # only add a StreamHandler if it is not present yet
+    if len(logger.handlers) <= 1:
+        logger.addHandler(logging.StreamHandler())