stuff and more stuff

postprocessing/mmaudio/model/embeddings.py

@@ -0,0 +1,49 @@
+import torch
+import torch.nn as nn
+
+# https://github.com/facebookresearch/DiT
+
+
+class TimestepEmbedder(nn.Module):
+    """
+    Embeds scalar timesteps into vector representations.
+    """
+
+    def __init__(self, dim, frequency_embedding_size, max_period):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, dim),
+            nn.SiLU(),
+            nn.Linear(dim, dim),
+        )
+        self.dim = dim
+        self.max_period = max_period
+        assert dim % 2 == 0, 'dim must be even.'
+
+        with torch.autocast('cuda', enabled=False):
+            self.freqs = nn.Buffer(
+                1.0 / (10000**(torch.arange(0, frequency_embedding_size, 2, dtype=torch.float32) /
+                               frequency_embedding_size)),
+                persistent=False)
+            freq_scale = 10000 / max_period
+            self.freqs = freq_scale * self.freqs
+
+    def timestep_embedding(self, t):
+        """
+        Create sinusoidal timestep embeddings.
+        :param t: a 1-D Tensor of N indices, one per batch element.
+                          These may be fractional.
+        :param dim: the dimension of the output.
+        :param max_period: controls the minimum frequency of the embeddings.
+        :return: an (N, D) Tensor of positional embeddings.
+        """
+        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+
+        args = t[:, None].float() * self.freqs[None]
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        return embedding
+
+    def forward(self, t):
+        t_freq = self.timestep_embedding(t).to(t.dtype)
+        t_emb = self.mlp(t_freq)
+        return t_emb

postprocessing/mmaudio/model/flow_matching.py

@@ -0,0 +1,71 @@
+import logging
+from typing import Callable, Optional
+
+import torch
+from torchdiffeq import odeint
+
+log = logging.getLogger()
+
+
+# Partially from https://github.com/gle-bellier/flow-matching
+class FlowMatching:
+
+    def __init__(self, min_sigma: float = 0.0, inference_mode='euler', num_steps: int = 25):
+        # inference_mode: 'euler' or 'adaptive'
+        # num_steps: number of steps in the euler inference mode
+        super().__init__()
+        self.min_sigma = min_sigma
+        self.inference_mode = inference_mode
+        self.num_steps = num_steps
+
+        # self.fm = ExactOptimalTransportConditionalFlowMatcher(sigma=min_sigma)
+
+        assert self.inference_mode in ['euler', 'adaptive']
+        if self.inference_mode == 'adaptive' and num_steps > 0:
+            log.info('The number of steps is ignored in adaptive inference mode ')
+
+    def get_conditional_flow(self, x0: torch.Tensor, x1: torch.Tensor,
+                             t: torch.Tensor) -> torch.Tensor:
+        # which is psi_t(x), eq 22 in flow matching for generative models
+        t = t[:, None, None].expand_as(x0)
+        return (1 - (1 - self.min_sigma) * t) * x0 + t * x1
+
+    def loss(self, predicted_v: torch.Tensor, x0: torch.Tensor, x1: torch.Tensor) -> torch.Tensor:
+        # return the mean error without reducing the batch dimension
+        reduce_dim = list(range(1, len(predicted_v.shape)))
+        target_v = x1 - (1 - self.min_sigma) * x0
+        return (predicted_v - target_v).pow(2).mean(dim=reduce_dim)
+
+    def get_x0_xt_c(
+        self,
+        x1: torch.Tensor,
+        t: torch.Tensor,
+        Cs: list[torch.Tensor],
+        generator: Optional[torch.Generator] = None
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+        x0 = torch.empty_like(x1).normal_(generator=generator)
+
+        xt = self.get_conditional_flow(x0, x1, t)
+        return x0, x1, xt, Cs
+
+    def to_prior(self, fn: Callable, x1: torch.Tensor) -> torch.Tensor:
+        return self.run_t0_to_t1(fn, x1, 1, 0)
+
+    def to_data(self, fn: Callable, x0: torch.Tensor) -> torch.Tensor:
+        return self.run_t0_to_t1(fn, x0, 0, 1)
+
+    def run_t0_to_t1(self, fn: Callable, x0: torch.Tensor, t0: float, t1: float) -> torch.Tensor:
+        # fn: a function that takes (t, x) and returns the direction x0->x1
+
+        if self.inference_mode == 'adaptive':
+            return odeint(fn, x0, torch.tensor([t0, t1], device=x0.device, dtype=x0.dtype))
+        elif self.inference_mode == 'euler':
+            x = x0
+            steps = torch.linspace(t0, t1 - self.min_sigma, self.num_steps + 1)
+            for ti, t in enumerate(steps[:-1]):
+                flow = fn(t, x)
+                next_t = steps[ti + 1]
+                dt = next_t - t
+                x = x + dt * flow
+
+        return x

postprocessing/mmaudio/model/low_level.py

@@ -0,0 +1,95 @@
+import torch
+from torch import nn
+from torch.nn import functional as F
+
+
+class ChannelLastConv1d(nn.Conv1d):
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = x.permute(0, 2, 1)
+        x = super().forward(x)
+        x = x.permute(0, 2, 1)
+        return x
+
+
+# https://github.com/Stability-AI/sd3-ref
+class MLP(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int = 256,
+    ):
+        """
+        Initialize the FeedForward module.
+
+        Args:
+            dim (int): Input dimension.
+            hidden_dim (int): Hidden dimension of the feedforward layer.
+            multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
+
+        Attributes:
+            w1 (ColumnParallelLinear): Linear transformation for the first layer.
+            w2 (RowParallelLinear): Linear transformation for the second layer.
+            w3 (ColumnParallelLinear): Linear transformation for the third layer.
+
+        """
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(dim, hidden_dim, bias=False)
+        self.w2 = nn.Linear(hidden_dim, dim, bias=False)
+        self.w3 = nn.Linear(dim, hidden_dim, bias=False)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))
+
+
+class ConvMLP(nn.Module):
+
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: int,
+        multiple_of: int = 256,
+        kernel_size: int = 3,
+        padding: int = 1,
+    ):
+        """
+        Initialize the FeedForward module.
+
+        Args:
+            dim (int): Input dimension.
+            hidden_dim (int): Hidden dimension of the feedforward layer.
+            multiple_of (int): Value to ensure hidden dimension is a multiple of this value.
+
+        Attributes:
+            w1 (ColumnParallelLinear): Linear transformation for the first layer.
+            w2 (RowParallelLinear): Linear transformation for the second layer.
+            w3 (ColumnParallelLinear): Linear transformation for the third layer.
+
+        """
+        super().__init__()
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = ChannelLastConv1d(dim,
+                                    hidden_dim,
+                                    bias=False,
+                                    kernel_size=kernel_size,
+                                    padding=padding)
+        self.w2 = ChannelLastConv1d(hidden_dim,
+                                    dim,
+                                    bias=False,
+                                    kernel_size=kernel_size,
+                                    padding=padding)
+        self.w3 = ChannelLastConv1d(dim,
+                                    hidden_dim,
+                                    bias=False,
+                                    kernel_size=kernel_size,
+                                    padding=padding)
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)) * self.w3(x))

postprocessing/mmaudio/model/networks.py

@@ -0,0 +1,477 @@
+import logging
+from dataclasses import dataclass
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from ..ext.rotary_embeddings import compute_rope_rotations
+from .embeddings import TimestepEmbedder
+from .low_level import MLP, ChannelLastConv1d, ConvMLP
+from .transformer_layers import (FinalBlock, JointBlock, MMDitSingleBlock)
+
+log = logging.getLogger()
+
+
+@dataclass
+class PreprocessedConditions:
+    clip_f: torch.Tensor
+    sync_f: torch.Tensor
+    text_f: torch.Tensor
+    clip_f_c: torch.Tensor
+    text_f_c: torch.Tensor
+
+
+# Partially from https://github.com/facebookresearch/DiT
+class MMAudio(nn.Module):
+
+    def __init__(self,
+                 *,
+                 latent_dim: int,
+                 clip_dim: int,
+                 sync_dim: int,
+                 text_dim: int,
+                 hidden_dim: int,
+                 depth: int,
+                 fused_depth: int,
+                 num_heads: int,
+                 mlp_ratio: float = 4.0,
+                 latent_seq_len: int,
+                 clip_seq_len: int,
+                 sync_seq_len: int,
+                 text_seq_len: int = 77,
+                 latent_mean: Optional[torch.Tensor] = None,
+                 latent_std: Optional[torch.Tensor] = None,
+                 empty_string_feat: Optional[torch.Tensor] = None,
+                 v2: bool = False) -> None:
+        super().__init__()
+
+        self.v2 = v2
+        self.latent_dim = latent_dim
+        self._latent_seq_len = latent_seq_len
+        self._clip_seq_len = clip_seq_len
+        self._sync_seq_len = sync_seq_len
+        self._text_seq_len = text_seq_len
+        self.hidden_dim = hidden_dim
+        self.num_heads = num_heads
+
+        if v2:
+            self.audio_input_proj = nn.Sequential(
+                ChannelLastConv1d(latent_dim, hidden_dim, kernel_size=7, padding=3),
+                nn.SiLU(),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=7, padding=3),
+            )
+
+            self.clip_input_proj = nn.Sequential(
+                nn.Linear(clip_dim, hidden_dim),
+                nn.SiLU(),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=3, padding=1),
+            )
+
+            self.sync_input_proj = nn.Sequential(
+                ChannelLastConv1d(sync_dim, hidden_dim, kernel_size=7, padding=3),
+                nn.SiLU(),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=3, padding=1),
+            )
+
+            self.text_input_proj = nn.Sequential(
+                nn.Linear(text_dim, hidden_dim),
+                nn.SiLU(),
+                MLP(hidden_dim, hidden_dim * 4),
+            )
+        else:
+            self.audio_input_proj = nn.Sequential(
+                ChannelLastConv1d(latent_dim, hidden_dim, kernel_size=7, padding=3),
+                nn.SELU(),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=7, padding=3),
+            )
+
+            self.clip_input_proj = nn.Sequential(
+                nn.Linear(clip_dim, hidden_dim),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=3, padding=1),
+            )
+
+            self.sync_input_proj = nn.Sequential(
+                ChannelLastConv1d(sync_dim, hidden_dim, kernel_size=7, padding=3),
+                nn.SELU(),
+                ConvMLP(hidden_dim, hidden_dim * 4, kernel_size=3, padding=1),
+            )
+
+            self.text_input_proj = nn.Sequential(
+                nn.Linear(text_dim, hidden_dim),
+                MLP(hidden_dim, hidden_dim * 4),
+            )
+
+        self.clip_cond_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.text_cond_proj = nn.Linear(hidden_dim, hidden_dim)
+        self.global_cond_mlp = MLP(hidden_dim, hidden_dim * 4)
+        # each synchformer output segment has 8 feature frames
+        self.sync_pos_emb = nn.Parameter(torch.zeros((1, 1, 8, sync_dim)))
+
+        self.final_layer = FinalBlock(hidden_dim, latent_dim)
+
+        if v2:
+            self.t_embed = TimestepEmbedder(hidden_dim,
+                                            frequency_embedding_size=hidden_dim,
+                                            max_period=1)
+        else:
+            self.t_embed = TimestepEmbedder(hidden_dim,
+                                            frequency_embedding_size=256,
+                                            max_period=10000)
+        self.joint_blocks = nn.ModuleList([
+            JointBlock(hidden_dim,
+                       num_heads,
+                       mlp_ratio=mlp_ratio,
+                       pre_only=(i == depth - fused_depth - 1)) for i in range(depth - fused_depth)
+        ])
+
+        self.fused_blocks = nn.ModuleList([
+            MMDitSingleBlock(hidden_dim, num_heads, mlp_ratio=mlp_ratio, kernel_size=3, padding=1)
+            for i in range(fused_depth)
+        ])
+
+        if latent_mean is None:
+            # these values are not meant to be used
+            # if you don't provide mean/std here, we should load them later from a checkpoint
+            assert latent_std is None
+            latent_mean = torch.ones(latent_dim).view(1, 1, -1).fill_(float('nan'))
+            latent_std = torch.ones(latent_dim).view(1, 1, -1).fill_(float('nan'))
+        else:
+            assert latent_std is not None
+            assert latent_mean.numel() == latent_dim, f'{latent_mean.numel()=} != {latent_dim=}'
+        if empty_string_feat is None:
+            empty_string_feat = torch.zeros((text_seq_len, text_dim))
+        self.latent_mean = nn.Parameter(latent_mean.view(1, 1, -1), requires_grad=False)
+        self.latent_std = nn.Parameter(latent_std.view(1, 1, -1), requires_grad=False)
+
+        self.empty_string_feat = nn.Parameter(empty_string_feat, requires_grad=False)
+        self.empty_clip_feat = nn.Parameter(torch.zeros(1, clip_dim), requires_grad=True)
+        self.empty_sync_feat = nn.Parameter(torch.zeros(1, sync_dim), requires_grad=True)
+
+        self.initialize_weights()
+        self.initialize_rotations()
+
+    def initialize_rotations(self):
+        base_freq = 1.0
+        latent_rot = compute_rope_rotations(self._latent_seq_len,
+                                            self.hidden_dim // self.num_heads,
+                                            10000,
+                                            freq_scaling=base_freq,
+                                            device=self.device)
+        clip_rot = compute_rope_rotations(self._clip_seq_len,
+                                          self.hidden_dim // self.num_heads,
+                                          10000,
+                                          freq_scaling=base_freq * self._latent_seq_len /
+                                          self._clip_seq_len,
+                                          device=self.device)
+
+        self.latent_rot = latent_rot #, persistent=False)
+        self.clip_rot = clip_rot #, persistent=False)
+
+    def update_seq_lengths(self, latent_seq_len: int, clip_seq_len: int, sync_seq_len: int) -> None:
+        self._latent_seq_len = latent_seq_len
+        self._clip_seq_len = clip_seq_len
+        self._sync_seq_len = sync_seq_len
+        self.initialize_rotations()
+
+    def initialize_weights(self):
+
+        def _basic_init(module):
+            if isinstance(module, nn.Linear):
+                torch.nn.init.xavier_uniform_(module.weight)
+                if module.bias is not None:
+                    nn.init.constant_(module.bias, 0)
+
+        self.apply(_basic_init)
+
+        # Initialize timestep embedding MLP:
+        nn.init.normal_(self.t_embed.mlp[0].weight, std=0.02)
+        nn.init.normal_(self.t_embed.mlp[2].weight, std=0.02)
+
+        # Zero-out adaLN modulation layers in DiT blocks:
+        for block in self.joint_blocks:
+            nn.init.constant_(block.latent_block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.latent_block.adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(block.clip_block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.clip_block.adaLN_modulation[-1].bias, 0)
+            nn.init.constant_(block.text_block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.text_block.adaLN_modulation[-1].bias, 0)
+        for block in self.fused_blocks:
+            nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
+            nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
+
+        # Zero-out output layers:
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].weight, 0)
+        nn.init.constant_(self.final_layer.adaLN_modulation[-1].bias, 0)
+        nn.init.constant_(self.final_layer.conv.weight, 0)
+        nn.init.constant_(self.final_layer.conv.bias, 0)
+
+        # empty string feat shall be initialized by a CLIP encoder
+        nn.init.constant_(self.sync_pos_emb, 0)
+        nn.init.constant_(self.empty_clip_feat, 0)
+        nn.init.constant_(self.empty_sync_feat, 0)
+
+    def normalize(self, x: torch.Tensor) -> torch.Tensor:
+        # return (x - self.latent_mean) / self.latent_std
+        return x.sub_(self.latent_mean).div_(self.latent_std)
+
+    def unnormalize(self, x: torch.Tensor) -> torch.Tensor:
+        # return x * self.latent_std + self.latent_mean
+        return x.mul_(self.latent_std).add_(self.latent_mean)
+
+    def preprocess_conditions(self, clip_f: torch.Tensor, sync_f: torch.Tensor,
+                              text_f: torch.Tensor) -> PreprocessedConditions:
+        """
+        cache computations that do not depend on the latent/time step
+        i.e., the features are reused over steps during inference
+        """
+        assert clip_f.shape[1] == self._clip_seq_len, f'{clip_f.shape=} {self._clip_seq_len=}'
+        assert sync_f.shape[1] == self._sync_seq_len, f'{sync_f.shape=} {self._sync_seq_len=}'
+        assert text_f.shape[1] == self._text_seq_len, f'{text_f.shape=} {self._text_seq_len=}'
+
+        bs = clip_f.shape[0]
+
+        # B * num_segments (24) * 8 * 768
+        num_sync_segments = self._sync_seq_len // 8
+        sync_f = sync_f.view(bs, num_sync_segments, 8, -1) + self.sync_pos_emb
+        sync_f = sync_f.flatten(1, 2)  # (B, VN, D)
+
+        # extend vf to match x
+        clip_f = self.clip_input_proj(clip_f)  # (B, VN, D)
+        sync_f = self.sync_input_proj(sync_f)  # (B, VN, D)
+        text_f = self.text_input_proj(text_f)  # (B, VN, D)
+
+        # upsample the sync features to match the audio
+        sync_f = sync_f.transpose(1, 2)  # (B, D, VN)
+        sync_f = F.interpolate(sync_f, size=self._latent_seq_len, mode='nearest-exact')
+        sync_f = sync_f.transpose(1, 2)  # (B, N, D)
+
+        # get conditional features from the clip side
+        clip_f_c = self.clip_cond_proj(clip_f.mean(dim=1))  # (B, D)
+        text_f_c = self.text_cond_proj(text_f.mean(dim=1))  # (B, D)
+
+        return PreprocessedConditions(clip_f=clip_f,
+                                      sync_f=sync_f,
+                                      text_f=text_f,
+                                      clip_f_c=clip_f_c,
+                                      text_f_c=text_f_c)
+
+    def predict_flow(self, latent: torch.Tensor, t: torch.Tensor,
+                     conditions: PreprocessedConditions) -> torch.Tensor:
+        """
+        for non-cacheable computations
+        """
+        assert latent.shape[1] == self._latent_seq_len, f'{latent.shape=} {self._latent_seq_len=}'
+
+        clip_f = conditions.clip_f
+        sync_f = conditions.sync_f
+        text_f = conditions.text_f
+        clip_f_c = conditions.clip_f_c
+        text_f_c = conditions.text_f_c
+
+        latent = self.audio_input_proj(latent)  # (B, N, D)
+        global_c = self.global_cond_mlp(clip_f_c + text_f_c)  # (B, D)
+
+        global_c = self.t_embed(t).unsqueeze(1) + global_c.unsqueeze(1)  # (B, D)
+        extended_c = global_c + sync_f
+
+
+
+        self.latent_rot = self.latent_rot.to("cuda")
+        self.clip_rot = self.clip_rot.to("cuda")
+        for block in self.joint_blocks:
+            latent, clip_f, text_f = block(latent, clip_f, text_f, global_c, extended_c,
+                                           self.latent_rot, self.clip_rot)  # (B, N, D)
+
+        for block in self.fused_blocks:
+            latent = block(latent, extended_c, self.latent_rot)
+        self.latent_rot = self.latent_rot.to("cpu")
+        self.clip_rot = self.clip_rot.to("cpu")
+
+        # should be extended_c; this is a minor implementation error #55
+        flow = self.final_layer(latent, global_c)  # (B, N, out_dim), remove t
+        return flow
+
+    def forward(self, latent: torch.Tensor, clip_f: torch.Tensor, sync_f: torch.Tensor,
+                text_f: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """
+        latent: (B, N, C) 
+        vf: (B, T, C_V)
+        t: (B,)
+        """
+        conditions = self.preprocess_conditions(clip_f, sync_f, text_f)
+        flow = self.predict_flow(latent, t, conditions)
+        return flow
+
+    def get_empty_string_sequence(self, bs: int) -> torch.Tensor:
+        return self.empty_string_feat.unsqueeze(0).expand(bs, -1, -1)
+
+    def get_empty_clip_sequence(self, bs: int) -> torch.Tensor:
+        return self.empty_clip_feat.unsqueeze(0).expand(bs, self._clip_seq_len, -1)
+
+    def get_empty_sync_sequence(self, bs: int) -> torch.Tensor:
+        return self.empty_sync_feat.unsqueeze(0).expand(bs, self._sync_seq_len, -1)
+
+    def get_empty_conditions(
+            self,
+            bs: int,
+            *,
+            negative_text_features: Optional[torch.Tensor] = None) -> PreprocessedConditions:
+        if negative_text_features is not None:
+            empty_text = negative_text_features
+        else:
+            empty_text = self.get_empty_string_sequence(1)
+
+        empty_clip = self.get_empty_clip_sequence(1)
+        empty_sync = self.get_empty_sync_sequence(1)
+        conditions = self.preprocess_conditions(empty_clip, empty_sync, empty_text)
+        conditions.clip_f = conditions.clip_f.expand(bs, -1, -1)
+        conditions.sync_f = conditions.sync_f.expand(bs, -1, -1)
+        conditions.clip_f_c = conditions.clip_f_c.expand(bs, -1)
+        if negative_text_features is None:
+            conditions.text_f = conditions.text_f.expand(bs, -1, -1)
+            conditions.text_f_c = conditions.text_f_c.expand(bs, -1)
+
+        return conditions
+
+    def ode_wrapper(self, t: torch.Tensor, latent: torch.Tensor, conditions: PreprocessedConditions,
+                    empty_conditions: PreprocessedConditions, cfg_strength: float) -> torch.Tensor:
+        t = t * torch.ones(len(latent), device=latent.device, dtype=latent.dtype)
+
+        if cfg_strength < 1.0:
+            return self.predict_flow(latent, t, conditions)
+        else:
+            return (cfg_strength * self.predict_flow(latent, t, conditions) +
+                    (1 - cfg_strength) * self.predict_flow(latent, t, empty_conditions))
+
+    def load_weights(self, src_dict) -> None:
+        if 't_embed.freqs' in src_dict:
+            del src_dict['t_embed.freqs']
+        if 'latent_rot' in src_dict:
+            del src_dict['latent_rot']
+        if 'clip_rot' in src_dict:
+            del src_dict['clip_rot']
+
+        a,b = self.load_state_dict(src_dict, strict=True, assign= True)
+        pass
+
+    @property
+    def device(self) -> torch.device:
+        return self.latent_mean.device
+
+    @property
+    def latent_seq_len(self) -> int:
+        return self._latent_seq_len
+
+    @property
+    def clip_seq_len(self) -> int:
+        return self._clip_seq_len
+
+    @property
+    def sync_seq_len(self) -> int:
+        return self._sync_seq_len
+
+
+def small_16k(**kwargs) -> MMAudio:
+    num_heads = 7
+    return MMAudio(latent_dim=20,
+                   clip_dim=1024,
+                   sync_dim=768,
+                   text_dim=1024,
+                   hidden_dim=64 * num_heads,
+                   depth=12,
+                   fused_depth=8,
+                   num_heads=num_heads,
+                   latent_seq_len=250,
+                   clip_seq_len=64,
+                   sync_seq_len=192,
+                   **kwargs)
+
+
+def small_44k(**kwargs) -> MMAudio:
+    num_heads = 7
+    return MMAudio(latent_dim=40,
+                   clip_dim=1024,
+                   sync_dim=768,
+                   text_dim=1024,
+                   hidden_dim=64 * num_heads,
+                   depth=12,
+                   fused_depth=8,
+                   num_heads=num_heads,
+                   latent_seq_len=345,
+                   clip_seq_len=64,
+                   sync_seq_len=192,
+                   **kwargs)
+
+
+def medium_44k(**kwargs) -> MMAudio:
+    num_heads = 14
+    return MMAudio(latent_dim=40,
+                   clip_dim=1024,
+                   sync_dim=768,
+                   text_dim=1024,
+                   hidden_dim=64 * num_heads,
+                   depth=12,
+                   fused_depth=8,
+                   num_heads=num_heads,
+                   latent_seq_len=345,
+                   clip_seq_len=64,
+                   sync_seq_len=192,
+                   **kwargs)
+
+
+def large_44k(**kwargs) -> MMAudio:
+    num_heads = 14
+    return MMAudio(latent_dim=40,
+                   clip_dim=1024,
+                   sync_dim=768,
+                   text_dim=1024,
+                   hidden_dim=64 * num_heads,
+                   depth=21,
+                   fused_depth=14,
+                   num_heads=num_heads,
+                   latent_seq_len=345,
+                   clip_seq_len=64,
+                   sync_seq_len=192,
+                   **kwargs)
+
+
+def large_44k_v2(**kwargs) -> MMAudio:
+    num_heads = 14
+    return MMAudio(latent_dim=40,
+                   clip_dim=1024,
+                   sync_dim=768,
+                   text_dim=1024,
+                   hidden_dim=64 * num_heads,
+                   depth=21,
+                   fused_depth=14,
+                   num_heads=num_heads,
+                   latent_seq_len=345,
+                   clip_seq_len=64,
+                   sync_seq_len=192,
+                   v2=True,
+                   **kwargs)
+
+
+def get_my_mmaudio(name: str, **kwargs) -> MMAudio:
+    if name == 'small_16k':
+        return small_16k(**kwargs)
+    if name == 'small_44k':
+        return small_44k(**kwargs)
+    if name == 'medium_44k':
+        return medium_44k(**kwargs)
+    if name == 'large_44k':
+        return large_44k(**kwargs)
+    if name == 'large_44k_v2':
+        return large_44k_v2(**kwargs)
+
+    raise ValueError(f'Unknown model name: {name}')
+
+
+if __name__ == '__main__':
+    network = get_my_mmaudio('small_16k')
+
+    # print the number of parameters in terms of millions
+    num_params = sum(p.numel() for p in network.parameters()) / 1e6
+    print(f'Number of parameters: {num_params:.2f}M')

postprocessing/mmaudio/model/sequence_config.py

@@ -0,0 +1,58 @@
+import dataclasses
+import math
+
+
+@dataclasses.dataclass
+class SequenceConfig:
+    # general
+    duration: float
+
+    # audio
+    sampling_rate: int
+    spectrogram_frame_rate: int
+    latent_downsample_rate: int = 2
+
+    # visual
+    clip_frame_rate: int = 8
+    sync_frame_rate: int = 25
+    sync_num_frames_per_segment: int = 16
+    sync_step_size: int = 8
+    sync_downsample_rate: int = 2
+
+    @property
+    def num_audio_frames(self) -> int:
+        # we need an integer number of latents
+        return self.latent_seq_len * self.spectrogram_frame_rate * self.latent_downsample_rate
+
+    @property
+    def latent_seq_len(self) -> int:
+        return int(
+            math.ceil(self.duration * self.sampling_rate / self.spectrogram_frame_rate /
+                      self.latent_downsample_rate))
+
+    @property
+    def clip_seq_len(self) -> int:
+        return int(self.duration * self.clip_frame_rate)
+
+    @property
+    def sync_seq_len(self) -> int:
+        num_frames = self.duration * self.sync_frame_rate
+        num_segments = (num_frames - self.sync_num_frames_per_segment) // self.sync_step_size + 1
+        return int(num_segments * self.sync_num_frames_per_segment / self.sync_downsample_rate)
+
+
+CONFIG_16K = SequenceConfig(duration=8.0, sampling_rate=16000, spectrogram_frame_rate=256)
+CONFIG_44K = SequenceConfig(duration=8.0, sampling_rate=44100, spectrogram_frame_rate=512)
+
+if __name__ == '__main__':
+    assert CONFIG_16K.latent_seq_len == 250
+    assert CONFIG_16K.clip_seq_len == 64
+    assert CONFIG_16K.sync_seq_len == 192
+    assert CONFIG_16K.num_audio_frames == 128000
+
+    assert CONFIG_44K.latent_seq_len == 345
+    assert CONFIG_44K.clip_seq_len == 64
+    assert CONFIG_44K.sync_seq_len == 192
+    assert CONFIG_44K.num_audio_frames == 353280
+
+    print('Passed')

postprocessing/mmaudio/model/transformer_layers.py

@@ -0,0 +1,202 @@
+from typing import Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from einops.layers.torch import Rearrange
+
+from ..ext.rotary_embeddings import apply_rope
+from ..model.low_level import MLP, ChannelLastConv1d, ConvMLP
+
+
+def modulate(x: torch.Tensor, shift: torch.Tensor, scale: torch.Tensor):
+    return x * (1 + scale) + shift
+
+
+def attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor):
+    # training will crash without these contiguous calls and the CUDNN limitation
+    # I believe this is related to https://github.com/pytorch/pytorch/issues/133974
+    # unresolved at the time of writing
+    q = q.contiguous()
+    k = k.contiguous()
+    v = v.contiguous()
+    out = F.scaled_dot_product_attention(q, k, v)
+    out = rearrange(out, 'b h n d -> b n (h d)').contiguous()
+    return out
+
+
+class SelfAttention(nn.Module):
+
+    def __init__(self, dim: int, nheads: int):
+        super().__init__()
+        self.dim = dim
+        self.nheads = nheads
+
+        self.qkv = nn.Linear(dim, dim * 3, bias=True)
+        self.q_norm = nn.RMSNorm(dim // nheads)
+        self.k_norm = nn.RMSNorm(dim // nheads)
+
+        self.split_into_heads = Rearrange('b n (h d j) -> b h n d j',
+                                          h=nheads,
+                                          d=dim // nheads,
+                                          j=3)
+
+    def pre_attention(
+            self, x: torch.Tensor,
+            rot: Optional[torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        # x: batch_size * n_tokens * n_channels
+        qkv = self.qkv(x)
+        q, k, v = self.split_into_heads(qkv).chunk(3, dim=-1)
+        q = q.squeeze(-1)
+        k = k.squeeze(-1)
+        v = v.squeeze(-1)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if rot is not None:
+            q = apply_rope(q, rot)
+            k = apply_rope(k, rot)
+
+        return q, k, v
+
+    def forward(
+            self,
+            x: torch.Tensor,  # batch_size * n_tokens * n_channels
+    ) -> torch.Tensor:
+        q, v, k = self.pre_attention(x)
+        out = attention(q, k, v)
+        return out
+
+
+class MMDitSingleBlock(nn.Module):
+
+    def __init__(self,
+                 dim: int,
+                 nhead: int,
+                 mlp_ratio: float = 4.0,
+                 pre_only: bool = False,
+                 kernel_size: int = 7,
+                 padding: int = 3):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, elementwise_affine=False)
+        self.attn = SelfAttention(dim, nhead)
+
+        self.pre_only = pre_only
+        if pre_only:
+            self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim, bias=True))
+        else:
+            if kernel_size == 1:
+                self.linear1 = nn.Linear(dim, dim)
+            else:
+                self.linear1 = ChannelLastConv1d(dim, dim, kernel_size=kernel_size, padding=padding)
+            self.norm2 = nn.LayerNorm(dim, elementwise_affine=False)
+
+            if kernel_size == 1:
+                self.ffn = MLP(dim, int(dim * mlp_ratio))
+            else:
+                self.ffn = ConvMLP(dim,
+                                   int(dim * mlp_ratio),
+                                   kernel_size=kernel_size,
+                                   padding=padding)
+
+            self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 6 * dim, bias=True))
+
+    def pre_attention(self, x: torch.Tensor, c: torch.Tensor, rot: Optional[torch.Tensor]):
+        # x: BS * N * D
+        # cond: BS * D
+        modulation = self.adaLN_modulation(c)
+        if self.pre_only:
+            (shift_msa, scale_msa) = modulation.chunk(2, dim=-1)
+            gate_msa = shift_mlp = scale_mlp = gate_mlp = None
+        else:
+            (shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp,
+             gate_mlp) = modulation.chunk(6, dim=-1)
+
+        x = modulate(self.norm1(x), shift_msa, scale_msa)
+        q, k, v = self.attn.pre_attention(x, rot)
+        return (q, k, v), (gate_msa, shift_mlp, scale_mlp, gate_mlp)
+
+    def post_attention(self, x: torch.Tensor, attn_out: torch.Tensor, c: tuple[torch.Tensor]):
+        if self.pre_only:
+            return x
+
+        (gate_msa, shift_mlp, scale_mlp, gate_mlp) = c
+        x = x + self.linear1(attn_out) * gate_msa
+        r = modulate(self.norm2(x), shift_mlp, scale_mlp)
+        x = x + self.ffn(r) * gate_mlp
+
+        return x
+
+    def forward(self, x: torch.Tensor, cond: torch.Tensor,
+                rot: Optional[torch.Tensor]) -> torch.Tensor:
+        # x: BS * N * D
+        # cond: BS * D
+        x_qkv, x_conditions = self.pre_attention(x, cond, rot)
+        attn_out = attention(*x_qkv)
+        x = self.post_attention(x, attn_out, x_conditions)
+
+        return x
+
+
+class JointBlock(nn.Module):
+
+    def __init__(self, dim: int, nhead: int, mlp_ratio: float = 4.0, pre_only: bool = False):
+        super().__init__()
+        self.pre_only = pre_only
+        self.latent_block = MMDitSingleBlock(dim,
+                                             nhead,
+                                             mlp_ratio,
+                                             pre_only=False,
+                                             kernel_size=3,
+                                             padding=1)
+        self.clip_block = MMDitSingleBlock(dim,
+                                           nhead,
+                                           mlp_ratio,
+                                           pre_only=pre_only,
+                                           kernel_size=3,
+                                           padding=1)
+        self.text_block = MMDitSingleBlock(dim, nhead, mlp_ratio, pre_only=pre_only, kernel_size=1)
+
+    def forward(self, latent: torch.Tensor, clip_f: torch.Tensor, text_f: torch.Tensor,
+                global_c: torch.Tensor, extended_c: torch.Tensor, latent_rot: torch.Tensor,
+                clip_rot: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
+        # latent: BS * N1 * D
+        # clip_f: BS * N2 * D
+        # c: BS * (1/N) * D
+        x_qkv, x_mod = self.latent_block.pre_attention(latent, extended_c, latent_rot)
+        c_qkv, c_mod = self.clip_block.pre_attention(clip_f, global_c, clip_rot)
+        t_qkv, t_mod = self.text_block.pre_attention(text_f, global_c, rot=None)
+
+        latent_len = latent.shape[1]
+        clip_len = clip_f.shape[1]
+        text_len = text_f.shape[1]
+
+        joint_qkv = [torch.cat([x_qkv[i], c_qkv[i], t_qkv[i]], dim=2) for i in range(3)]
+
+        attn_out = attention(*joint_qkv)
+        x_attn_out = attn_out[:, :latent_len]
+        c_attn_out = attn_out[:, latent_len:latent_len + clip_len]
+        t_attn_out = attn_out[:, latent_len + clip_len:]
+
+        latent = self.latent_block.post_attention(latent, x_attn_out, x_mod)
+        if not self.pre_only:
+            clip_f = self.clip_block.post_attention(clip_f, c_attn_out, c_mod)
+            text_f = self.text_block.post_attention(text_f, t_attn_out, t_mod)
+
+        return latent, clip_f, text_f
+
+
+class FinalBlock(nn.Module):
+
+    def __init__(self, dim, out_dim):
+        super().__init__()
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(dim, 2 * dim, bias=True))
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False)
+        self.conv = ChannelLastConv1d(dim, out_dim, kernel_size=7, padding=3)
+
+    def forward(self, latent, c):
+        shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
+        latent = modulate(self.norm(latent), shift, scale)
+        latent = self.conv(latent)
+        return latent

postprocessing/mmaudio/model/utils/distributions.py

@@ -0,0 +1,46 @@
+from typing import Optional
+
+import numpy as np
+import torch
+
+
+class DiagonalGaussianDistribution:
+
+    def __init__(self, parameters, deterministic=False):
+        self.parameters = parameters
+        self.mean, self.logvar = torch.chunk(parameters, 2, dim=1)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if self.deterministic:
+            self.var = self.std = torch.zeros_like(self.mean).to(device=self.parameters.device)
+
+    def sample(self, rng: Optional[torch.Generator] = None):
+        # x = self.mean + self.std * torch.randn(self.mean.shape).to(device=self.parameters.device)
+
+        r = torch.empty_like(self.mean).normal_(generator=rng)
+        x = self.mean + self.std * r
+
+        return x
+
+    def kl(self, other=None):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        else:
+            if other is None:
+
+                return 0.5 * torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar
+            else:
+                return 0.5 * (torch.pow(self.mean - other.mean, 2) / other.var +
+                              self.var / other.var - 1.0 - self.logvar + other.logvar)
+
+    def nll(self, sample, dims=[1, 2, 3]):
+        if self.deterministic:
+            return torch.Tensor([0.])
+        logtwopi = np.log(2.0 * np.pi)
+        return 0.5 * torch.sum(logtwopi + self.logvar + torch.pow(sample - self.mean, 2) / self.var,
+                               dim=dims)
+
+    def mode(self):
+        return self.mean

postprocessing/mmaudio/model/utils/features_utils.py

@@ -0,0 +1,174 @@
+from typing import Literal, Optional
+import json
+import open_clip
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange
+from open_clip import create_model_from_pretrained, create_model
+from torchvision.transforms import Normalize
+
+from ...ext.autoencoder import AutoEncoderModule
+from ...ext.mel_converter import get_mel_converter
+from ...ext.synchformer.synchformer import Synchformer
+from ...model.utils.distributions import DiagonalGaussianDistribution
+
+
+def patch_clip(clip_model):
+    # a hack to make it output last hidden states
+    # https://github.com/mlfoundations/open_clip/blob/fc5a37b72d705f760ebbc7915b84729816ed471f/src/open_clip/model.py#L269
+    def new_encode_text(self, text, normalize: bool = False):
+        cast_dtype = self.transformer.get_cast_dtype()
+
+        x = self.token_embedding(text).to(cast_dtype)  # [batch_size, n_ctx, d_model]
+
+        x = x + self.positional_embedding.to(cast_dtype)
+        x = self.transformer(x, attn_mask=self.attn_mask)
+        x = self.ln_final(x)  # [batch_size, n_ctx, transformer.width]
+        return F.normalize(x, dim=-1) if normalize else x
+
+    clip_model.encode_text = new_encode_text.__get__(clip_model)
+    return clip_model
+
+def get_model_config(model_name):
+    with open("ckpts/DFN5B-CLIP-ViT-H-14-378/open_clip_config.json", 'r', encoding='utf-8') as f:
+        return json.load(f)["model_cfg"]
+
+class FeaturesUtils(nn.Module):
+
+    def __init__(
+        self,
+        *,
+        tod_vae_ckpt: Optional[str] = None,
+        bigvgan_vocoder_ckpt: Optional[str] = None,
+        synchformer_ckpt: Optional[str] = None,
+        enable_conditions: bool = True,
+        mode=Literal['16k', '44k'],
+        need_vae_encoder: bool = True,
+    ):
+        super().__init__()
+        self.device ="cuda"
+        if enable_conditions:
+            old_get_model_config = open_clip.factory.get_model_config
+            open_clip.factory.get_model_config = get_model_config
+            with open("ckpts/DFN5B-CLIP-ViT-H-14-378/open_clip_config.json", 'r', encoding='utf-8') as f:
+                override_preprocess = json.load(f)["preprocess_cfg"]
+
+            self.clip_model = create_model('DFN5B-CLIP-ViT-H-14-378', pretrained='ckpts/DFN5B-CLIP-ViT-H-14-378/open_clip_pytorch_model.bin',  force_preprocess_cfg= override_preprocess)
+            open_clip.factory.get_model_config = old_get_model_config
+
+            # self.clip_model = create_model_from_pretrained('hf-hub:apple/DFN5B-CLIP-ViT-H-14-384', return_transform=False)
+            self.clip_preprocess = Normalize(mean=[0.48145466, 0.4578275, 0.40821073],
+                                             std=[0.26862954, 0.26130258, 0.27577711])
+            self.clip_model = patch_clip(self.clip_model)
+
+            self.synchformer = Synchformer()
+            self.synchformer.load_state_dict(
+                torch.load(synchformer_ckpt, weights_only=True, map_location='cpu'))
+
+            self.tokenizer = open_clip.get_tokenizer('ViT-H-14-378-quickgelu')  # same as 'ViT-H-14'
+        else:
+            self.clip_model = None
+            self.synchformer = None
+            self.tokenizer = None
+
+        if tod_vae_ckpt is not None:
+            self.mel_converter = get_mel_converter(mode)
+            self.tod = AutoEncoderModule(vae_ckpt_path=tod_vae_ckpt,
+                                         vocoder_ckpt_path=bigvgan_vocoder_ckpt,
+                                         mode=mode,
+                                         need_vae_encoder=need_vae_encoder)
+        else:
+            self.tod = None
+
+    def compile(self):
+        if self.clip_model is not None:
+            self.clip_model.encode_image = torch.compile(self.clip_model.encode_image)
+            self.clip_model.encode_text = torch.compile(self.clip_model.encode_text)
+        if self.synchformer is not None:
+            self.synchformer = torch.compile(self.synchformer)
+        self.decode = torch.compile(self.decode)
+        self.vocode = torch.compile(self.vocode)
+
+    def train(self, mode: bool) -> None:
+        return super().train(False)
+
+    @torch.inference_mode()
+    def encode_video_with_clip(self, x: torch.Tensor, batch_size: int = -1) -> torch.Tensor:
+        assert self.clip_model is not None, 'CLIP is not loaded'
+        # x: (B, T, C, H, W) H/W: 384
+        b, t, c, h, w = x.shape
+        assert c == 3 and h == 384 and w == 384
+        x = self.clip_preprocess(x)
+        x = rearrange(x, 'b t c h w -> (b t) c h w')
+        outputs = []
+        if batch_size < 0:
+            batch_size = b * t
+        for i in range(0, b * t, batch_size):
+            outputs.append(self.clip_model.encode_image(x[i:i + batch_size], normalize=True))
+        x = torch.cat(outputs, dim=0)
+        # x = self.clip_model.encode_image(x, normalize=True)
+        x = rearrange(x, '(b t) d -> b t d', b=b)
+        return x
+
+    @torch.inference_mode()
+    def encode_video_with_sync(self, x: torch.Tensor, batch_size: int = -1) -> torch.Tensor:
+        assert self.synchformer is not None, 'Synchformer is not loaded'
+        # x: (B, T, C, H, W) H/W: 384
+
+        b, t, c, h, w = x.shape
+        assert c == 3 and h == 224 and w == 224
+
+        # partition the video
+        segment_size = 16
+        step_size = 8
+        num_segments = (t - segment_size) // step_size + 1
+        segments = []
+        for i in range(num_segments):
+            segments.append(x[:, i * step_size:i * step_size + segment_size])
+        x = torch.stack(segments, dim=1)  # (B, S, T, C, H, W)
+
+        outputs = []
+        if batch_size < 0:
+            batch_size = b
+        x = rearrange(x, 'b s t c h w -> (b s) 1 t c h w')
+        for i in range(0, b * num_segments, batch_size):
+            outputs.append(self.synchformer(x[i:i + batch_size]))
+        x = torch.cat(outputs, dim=0)
+        x = rearrange(x, '(b s) 1 t d -> b (s t) d', b=b)
+        return x
+
+    @torch.inference_mode()
+    def encode_text(self, text: list[str]) -> torch.Tensor:
+        assert self.clip_model is not None, 'CLIP is not loaded'
+        assert self.tokenizer is not None, 'Tokenizer is not loaded'
+        # x: (B, L)
+        tokens = self.tokenizer(text).to(self.device)
+        return self.clip_model.encode_text(tokens, normalize=True)
+
+    @torch.inference_mode()
+    def encode_audio(self, x) -> DiagonalGaussianDistribution:
+        assert self.tod is not None, 'VAE is not loaded'
+        # x: (B * L)
+        mel = self.mel_converter(x)
+        dist = self.tod.encode(mel)
+
+        return dist
+
+    @torch.inference_mode()
+    def vocode(self, mel: torch.Tensor) -> torch.Tensor:
+        assert self.tod is not None, 'VAE is not loaded'
+        return self.tod.vocode(mel)
+
+    @torch.inference_mode()
+    def decode(self, z: torch.Tensor) -> torch.Tensor:
+        assert self.tod is not None, 'VAE is not loaded'
+        return self.tod.decode(z.transpose(1, 2))
+
+    # @property
+    # def device(self):
+    #     return next(self.parameters()).device
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype

postprocessing/mmaudio/model/utils/parameter_groups.py

@@ -0,0 +1,72 @@
+import logging
+
+log = logging.getLogger()
+
+
+def get_parameter_groups(model, cfg, print_log=False):
+    """
+    Assign different weight decays and learning rates to different parameters.
+    Returns a parameter group which can be passed to the optimizer.
+    """
+    weight_decay = cfg.weight_decay
+    # embed_weight_decay = cfg.embed_weight_decay
+    # backbone_lr_ratio = cfg.backbone_lr_ratio
+    base_lr = cfg.learning_rate
+
+    backbone_params = []
+    embed_params = []
+    other_params = []
+
+    # embedding_names = ['summary_pos', 'query_init', 'query_emb', 'obj_pe']
+    # embedding_names = [e + '.weight' for e in embedding_names]
+
+    # inspired by detectron2
+    memo = set()
+    for name, param in model.named_parameters():
+        if not param.requires_grad:
+            continue
+        # Avoid duplicating parameters
+        if param in memo:
+            continue
+        memo.add(param)
+
+        if name.startswith('module'):
+            name = name[7:]
+
+        inserted = False
+        # if name.startswith('pixel_encoder.'):
+        #     backbone_params.append(param)
+        #     inserted = True
+        #     if print_log:
+        #         log.info(f'{name} counted as a backbone parameter.')
+        # else:
+        #     for e in embedding_names:
+        #         if name.endswith(e):
+        #             embed_params.append(param)
+        #             inserted = True
+        #             if print_log:
+        #                 log.info(f'{name} counted as an embedding parameter.')
+        #             break
+
+        # if not inserted:
+        other_params.append(param)
+
+    parameter_groups = [
+        # {
+        #     'params': backbone_params,
+        #     'lr': base_lr * backbone_lr_ratio,
+        #     'weight_decay': weight_decay
+        # },
+        # {
+        #     'params': embed_params,
+        #     'lr': base_lr,
+        #     'weight_decay': embed_weight_decay
+        # },
+        {
+            'params': other_params,
+            'lr': base_lr,
+            'weight_decay': weight_decay
+        },
+    ]
+
+    return parameter_groups

postprocessing/mmaudio/model/utils/sample_utils.py

@@ -0,0 +1,12 @@
+from typing import Optional
+
+import torch
+
+
+def log_normal_sample(x: torch.Tensor,
+                      generator: Optional[torch.Generator] = None,
+                      m: float = 0.0,
+                      s: float = 1.0) -> torch.Tensor:
+    bs = x.shape[0]
+    s = torch.randn(bs, device=x.device, generator=generator) * s + m
+    return torch.sigmoid(s)