minicpmo-utils 0.1.0__py3-none-any.whl

stepaudio2/cosyvoice2/flow/flow.py

@@ -0,0 +1,230 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import torch
+import torch.nn as nn
+from torch.nn import functional as F
+
+from stepaudio2.cosyvoice2.utils.mask import make_pad_mask
+from stepaudio2.cosyvoice2.flow.flow_matching import CausalConditionalCFM
+from stepaudio2.cosyvoice2.transformer.upsample_encoder_v2 import UpsampleConformerEncoderV2
+
+
+class CausalMaskedDiffWithXvec(torch.nn.Module):
+    def __init__(self,
+                 input_size: int = 512,
+                 output_size: int = 80,
+                 spk_embed_dim: int = 192,
+                 output_type: str = "mel",
+                 vocab_size: int = 5121,
+                 encoder: UpsampleConformerEncoderV2 = None,
+                 decoder: CausalConditionalCFM = None,
+                 input_embedding: torch.nn.Module = None,
+                 ):
+        super().__init__()
+        self.input_size = input_size
+        self.output_size = output_size
+        self.vocab_size = vocab_size
+        self.output_type = output_type
+        self.pre_lookahead_len = int(encoder.pre_lookahead_layer.pre_lookahead_len)
+        self.up_rate = int(encoder.up_layer.stride)
+        if input_embedding is None:
+            self.input_embedding = nn.Embedding(vocab_size, input_size)
+        else:
+            self.input_embedding = input_embedding
+        self.spk_embed_affine_layer = torch.nn.Linear(spk_embed_dim, output_size)
+        self.encoder = encoder
+        self.encoder_proj = torch.nn.Linear(self.encoder.output_size(), output_size)
+        self.decoder = decoder
+
+        # xvec projection with CUDA Graph optimization
+        # 初始化 CUDA Graph 相关变量
+        self.enable_cuda_graph = False
+        self.static_embedding = None
+        self.static_output = None
+        self.graph = None
+        self.embedding_shape = None
+
+    def scatter_cuda_graph(self, enable_cuda_graph: bool):
+        self.enable_cuda_graph = enable_cuda_graph
+        if self.enable_cuda_graph:
+            # self.encoder.scatter_cuda_graph(enable_cuda_graph)
+            self.decoder.scatter_cuda_graph(enable_cuda_graph)
+
+    @torch.inference_mode()
+    def inference(self,
+                  token,
+                  token_len,
+                  prompt_token,
+                  prompt_token_len,
+                  prompt_feat,
+                  prompt_feat_len,
+                  embedding,
+                  n_timesteps: int = 10,
+                  ):
+        assert token.shape[0] == 1
+
+        # xvec projection
+        embedding = F.normalize(embedding, dim=1)
+        embedding = self.spk_embed_affine_layer(embedding)
+    
+        # concat text and prompt_text
+        token_len = prompt_token_len + token_len
+        token = torch.concat([prompt_token, token], dim=1)
+        
+        mask = (~make_pad_mask(token_len)).unsqueeze(-1).to(embedding)
+        token = self.input_embedding(torch.clamp(token, min=0)) * mask
+
+        # token encode
+        h, _ = self.encoder.forward(token, token_len)
+        h = self.encoder_proj(h)
+
+        # condition
+        mel_len1 = prompt_feat.shape[1]
+        mel_len2 = h.shape[1] - prompt_feat.shape[1]
+
+        conds = torch.zeros_like(h)
+        conds[:, :mel_len1] = prompt_feat
+        conds = conds.transpose(1, 2).contiguous()
+
+        mask = (~make_pad_mask(torch.tensor([mel_len1 + mel_len2]))).to(h)
+
+        feat = self.decoder.forward(
+            mu=h.transpose(1, 2).contiguous(),
+            mask=mask.unsqueeze(1),
+            spks=embedding,
+            cond=conds,
+            n_timesteps=n_timesteps,
+        )
+
+        feat = feat[:, :, mel_len1:]
+        assert feat.shape[2] == mel_len2
+        return feat
+
+    @torch.inference_mode()
+    def setup_cache(self, 
+                    token: torch.Tensor, 
+                    mel: torch.Tensor, 
+                    spk: torch.Tensor, 
+                    n_timesteps: int = 10,
+                    ):
+        """
+        Args:
+            token: shape (b, t), with look ahead tokens
+            mel: shape (b, t, c), groundtruth mel
+            spk: shape (b, 192), speaker embedding
+        Returns:
+            cache: dict {
+                'conformer': {'cnn_cache': xxx, 'att_cache': xxx}, 
+                'estimator': {'cnn_cache': xxx, 'att_cache': xxx}
+            }
+        """
+        # check if look ahead token included
+        assert (token.shape[1] - self.pre_lookahead_len) * self.up_rate == mel.shape[1], (token.shape, mel.shape)
+
+        # xvec projection
+        spk = F.normalize(spk, dim=1)
+        spk = self.spk_embed_affine_layer(spk)
+
+        token = self.input_embedding(token)
+        # NOTE encoder.forward_chunk will strip the look ahead part
+        h, conformer_cnn_cache, conformer_att_cache = self.encoder.forward_chunk(
+            xs = token,
+            last_chunk = False,
+            cnn_cache = None,
+            att_cache = None,
+        )
+        h = self.encoder_proj(h)
+
+        feat, estimator_cnn_cache, estimator_att_cache = self.decoder.forward_chunk(
+            mu = h.transpose(1, 2).contiguous(),
+            spks = spk,
+            cond = mel.transpose(1, 2).contiguous(),
+            n_timesteps = n_timesteps,
+            temperature = 1.0,
+            cnn_cache = None,
+            att_cache = None,
+        )
+
+        cache = {
+            'conformer_cnn_cache': conformer_cnn_cache,
+            'conformer_att_cache': conformer_att_cache,
+            'estimator_cnn_cache': estimator_cnn_cache,
+            'estimator_att_cache': estimator_att_cache,
+        }
+        
+        # print("examining flow cache")
+        # from IPython import embed; embed()
+        
+        
+        return cache
+
+    @torch.inference_mode()
+    def inference_chunk(self, 
+                        token: torch.Tensor,
+                        spk: torch.Tensor, 
+                        cache: dict,
+                        last_chunk: bool = False,
+                        n_timesteps: int = 10,
+                        ):
+        """
+        Args:
+            token: shape (b, t), with look ahead tokens
+            spk: shape (b, 192), speaker embedding
+            cache: dict {
+                'conformer_cnn_cache': xxx,
+                ...
+            }
+        """
+        # unpack cache
+        conformer_cnn_cache = cache['conformer_cnn_cache']
+        conformer_att_cache = cache['conformer_att_cache']
+        estimator_cnn_cache = cache['estimator_cnn_cache']
+        estimator_att_cache = cache['estimator_att_cache']
+
+        # xvec projection
+        spk = F.normalize(spk, dim=1)
+        spk = self.spk_embed_affine_layer(spk)
+
+        token = self.input_embedding(token)
+        # if not the last chunk, h is shorter than xs for a length of lookahead_length * stride (6)
+        h, conformer_cnn_cache, conformer_att_cache = self.encoder.forward_chunk(
+            xs = token,
+            last_chunk = last_chunk,
+            cnn_cache = conformer_cnn_cache,
+            att_cache = conformer_att_cache,
+        )
+        h = self.encoder_proj(h)
+
+        cond = torch.zeros_like(h)
+        # forward estimator
+        feat, estimator_cnn_cache, estimator_att_cache = self.decoder.forward_chunk(
+            mu = h.transpose(1, 2).contiguous(),
+            spks = spk,
+            cond = cond.transpose(1, 2).contiguous(),
+            n_timesteps = n_timesteps,
+            temperature = 1.0,
+            cnn_cache = estimator_cnn_cache,
+            att_cache = estimator_att_cache,
+        )
+
+
+        new_cache = {
+            'conformer_cnn_cache': conformer_cnn_cache,
+            'conformer_att_cache': conformer_att_cache,
+            'estimator_cnn_cache': estimator_cnn_cache,
+            'estimator_att_cache': estimator_att_cache,
+        }
+
+        return feat, new_cache
+

stepaudio2/cosyvoice2/flow/flow_matching.py

@@ -0,0 +1,205 @@
+# Copyright (c) 2024 Alibaba Inc (authors: Xiang Lyu, Zhihao Du)
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from typing import List
+import onnxruntime
+import torch
+import torch.nn.functional as F
+
+from stepaudio2.cosyvoice2.flow.decoder_dit import DiT
+from stepaudio2.cosyvoice2.utils.mask import make_pad_mask
+
+
+"""
+Inference wrapper
+"""
+class CausalConditionalCFM(torch.nn.Module):
+    def __init__(self, estimator: DiT, inference_cfg_rate:float=0.7):
+        super().__init__()
+        self.estimator = estimator
+        self.inference_cfg_rate = inference_cfg_rate
+        self.out_channels = estimator.out_channels
+         # a maximum of 600s
+        self.register_buffer('rand_noise', torch.randn([1, self.out_channels, 50 * 600]), persistent=False)
+
+        self.register_buffer('cnn_cache_buffer', torch.zeros(16, 16, 2, 1024, 2), persistent=False)
+        self.register_buffer('att_cache_buffer', torch.zeros(16, 16, 2, 8, 1000, 128), persistent=False)
+
+    def scatter_cuda_graph(self, enable_cuda_graph: bool):
+        if enable_cuda_graph:
+            self.estimator._init_cuda_graph_all()
+
+    def solve_euler(self, x, t_span, mu, mask, spks, cond):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+        """
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+        t = t.unsqueeze(dim=0)
+        assert self.inference_cfg_rate > 0, 'inference_cfg_rate better > 0'
+
+        # constant during denoising
+        mask_in = torch.cat([mask, mask], dim=0)
+        mu_in = torch.cat([mu, torch.zeros_like(mu)], dim=0)
+        spks_in = torch.cat([spks, torch.zeros_like(spks)], dim=0)
+        cond_in = torch.cat([cond, torch.zeros_like(cond)], dim=0)
+        
+        for step in range(1, len(t_span)):
+
+            x_in = torch.cat([x, x], dim=0)
+            t_in = torch.cat([t, t], dim=0)
+
+            dphi_dt = self.estimator.forward(
+                x_in,
+                mask_in,
+                mu_in,
+                t_in,
+                spks_in,
+                cond_in,
+            )
+            dphi_dt, cfg_dphi_dt = torch.split(dphi_dt, [x.size(0), x.size(0)], dim=0)
+            dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt)
+            x = x + dt * dphi_dt
+            t = t + dt
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        return x
+
+    @torch.inference_mode()
+    def forward(self, mu, mask, spks, cond, n_timesteps=10, temperature=1.0):
+        z = self.rand_noise[:, :, :mu.size(2)] * temperature
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
+        # cosine scheduling
+        t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        return self.solve_euler(z, t_span, mu, mask, spks, cond)
+
+    def solve_euler_chunk(self, 
+                          x:torch.Tensor, 
+                          t_span:torch.Tensor, 
+                          mu:torch.Tensor, 
+                          spks:torch.Tensor, 
+                          cond:torch.Tensor, 
+                          cnn_cache:torch.Tensor=None,
+                          att_cache:torch.Tensor=None,
+                          ):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+            mask (torch.Tensor): output_mask
+                shape: (batch_size, 1, mel_timesteps)
+            spks (torch.Tensor, optional): speaker ids. Defaults to None.
+                shape: (batch_size, spk_emb_dim)
+            cond: Not used but kept for future purposes
+            cnn_cache: shape (n_time, depth, b, c1+c2, 2)
+            att_cache: shape (n_time, depth, b, nh, t, c * 2)
+        """
+        assert self.inference_cfg_rate > 0, 'cfg rate should be > 0'
+        
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+        t = t.unsqueeze(dim=0)  # (b,)
+
+        # setup initial cache
+        if cnn_cache is None:
+            cnn_cache = [None for _ in range(len(t_span)-1)]
+        if att_cache is None:
+            att_cache = [None for _ in range(len(t_span)-1)]
+        # next chunk's cache at each timestep
+
+        if att_cache[0] is not None:
+            last_att_len = att_cache.shape[4]
+        else:
+            last_att_len = 0
+
+        # constant during denoising
+        mu_in = torch.cat([mu, torch.zeros_like(mu)], dim=0)
+        spks_in = torch.cat([spks, torch.zeros_like(spks)], dim=0)
+        cond_in = torch.cat([cond, torch.zeros_like(cond)], dim=0)
+        for step in range(1, len(t_span)):
+            # torch.cuda.memory._record_memory_history(max_entries=100000)
+            # torch.cuda.memory._record_memory_history(max_entries=100000)
+            this_att_cache = att_cache[step-1]
+            this_cnn_cache = cnn_cache[step-1]
+
+            dphi_dt, this_new_cnn_cache, this_new_att_cache = self.estimator.forward_chunk(
+                x = x.repeat(2, 1, 1),
+                mu = mu_in,
+                t = t.repeat(2),
+                spks = spks_in,
+                cond = cond_in,
+                cnn_cache = this_cnn_cache,
+                att_cache = this_att_cache,
+            )
+            dphi_dt, cfg_dphi_dt = dphi_dt.chunk(2, dim=0)
+            dphi_dt = ((1.0 + self.inference_cfg_rate) * dphi_dt - self.inference_cfg_rate * cfg_dphi_dt)
+            x = x + dt * dphi_dt
+            t = t + dt
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+            self.cnn_cache_buffer[step-1] = this_new_cnn_cache
+            self.att_cache_buffer[step-1][:, :, :, :x.shape[2]+last_att_len, :] = this_new_att_cache
+        
+        cnn_cache = self.cnn_cache_buffer
+        att_cache = self.att_cache_buffer[:, :, :, :, :x.shape[2]+last_att_len, :]
+        return x, cnn_cache, att_cache
+    
+    @torch.inference_mode()
+    def forward_chunk(self, 
+                      mu:torch.Tensor, 
+                      spks:torch.Tensor, 
+                      cond:torch.Tensor,
+                      n_timesteps:int=10, 
+                      temperature:float=1.0, 
+                      cnn_cache:torch.Tensor=None,
+                      att_cache:torch.Tensor=None,
+                      ):
+        """
+        Args:
+            mu(torch.Tensor): shape (b, c, t)
+            spks(torch.Tensor): shape (b, 192)
+            cond(torch.Tensor): shape (b, c, t)
+            cnn_cache: shape (n_time, depth, b, c1+c2, 2)
+            att_cache: shape (n_time, depth, b, nh, t, c * 2)
+        """
+        # get offset from att_cache
+        offset = att_cache.shape[4] if att_cache is not None else 0
+        z = self.rand_noise[:, :, offset:offset+mu.size(2)] * temperature
+        t_span = torch.linspace(0, 1, n_timesteps + 1, device=mu.device, dtype=mu.dtype)
+        # cosine scheduling
+        t_span = 1 - torch.cos(t_span * 0.5 * torch.pi)
+        x, new_cnn_cache, new_att_cache = self.solve_euler_chunk(
+            x=z,
+            t_span=t_span,
+            mu=mu,
+            spks=spks,
+            cond=cond,
+            att_cache=att_cache,
+            cnn_cache=cnn_cache,
+        )
+        return x, new_cnn_cache, new_att_cache