Initial commit

zipvoice/models/modules/solver.py

@@ -0,0 +1,277 @@
+#!/usr/bin/env python3
+# Copyright         2024  Xiaomi Corp.        (authors: Han Zhu)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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 Optional, Union
+
+import torch
+
+
+class DiffusionModel(torch.nn.Module):
+    """A wrapper of diffusion models for inference.
+    Args:
+        model: The diffusion model.
+        func_name: The function name to call.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        func_name: str = "forward_fm_decoder",
+    ):
+        super().__init__()
+        self.model = model
+        self.func_name = func_name
+        self.model_func = getattr(self.model, func_name)
+
+    def forward(
+        self,
+        t: torch.Tensor,
+        x: torch.Tensor,
+        text_condition: torch.Tensor,
+        speech_condition: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        guidance_scale: Union[float, torch.Tensor] = 0.0,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward function that Handles the classifier-free guidance.
+        Args:
+            t: The current timestep, a tensor of a tensor of a single float.
+            x: The initial value, with the shape (batch, seq_len, emb_dim).
+            text_condition: The text_condition of the diffision model, with
+                the shape (batch, seq_len, emb_dim).
+            speech_condition: The speech_condition of the diffision model, with the
+                shape (batch, seq_len, emb_dim).
+            padding_mask: The mask for padding; True means masked position, with the
+                shape (batch, seq_len).
+            guidance_scale: The scale of classifier-free guidance, a float or a tensor
+                of shape (batch, 1, 1).
+        Retrun:
+            The prediction with the shape (batch, seq_len, emb_dim).
+        """
+        if not torch.is_tensor(guidance_scale):
+            guidance_scale = torch.tensor(
+                guidance_scale, dtype=t.dtype, device=t.device
+            )
+
+        if (guidance_scale == 0.0).all():
+            return self.model_func(
+                t=t,
+                xt=x,
+                text_condition=text_condition,
+                speech_condition=speech_condition,
+                padding_mask=padding_mask,
+                **kwargs
+            )
+        else:
+            assert t.dim() == 0
+
+            x = torch.cat([x] * 2, dim=0)
+            padding_mask = torch.cat([padding_mask] * 2, dim=0)
+
+            text_condition = torch.cat(
+                [torch.zeros_like(text_condition), text_condition], dim=0
+            )
+
+            if t > 0.5:
+                speech_condition = torch.cat(
+                    [torch.zeros_like(speech_condition), speech_condition], dim=0
+                )
+            else:
+                guidance_scale = guidance_scale * 2
+                speech_condition = torch.cat(
+                    [speech_condition, speech_condition], dim=0
+                )
+
+            data_uncond, data_cond = self.model_func(
+                t=t,
+                xt=x,
+                text_condition=text_condition,
+                speech_condition=speech_condition,
+                padding_mask=padding_mask,
+                **kwargs
+            ).chunk(2, dim=0)
+
+            res = (1 + guidance_scale) * data_cond - guidance_scale * data_uncond
+            return res
+
+
+class DistillDiffusionModel(DiffusionModel):
+    """A wrapper of distilled diffusion models for inference.
+    Args:
+        model: The distilled diffusion model.
+        func_name: The function name to call.
+    """
+
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        func_name: str = "forward_fm_decoder",
+    ):
+        super().__init__(model=model, func_name=func_name)
+
+    def forward(
+        self,
+        t: torch.Tensor,
+        x: torch.Tensor,
+        text_condition: torch.Tensor,
+        speech_condition: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        guidance_scale: Union[float, torch.Tensor] = 0.0,
+        **kwargs
+    ) -> torch.Tensor:
+        """
+        Forward function that Handles the classifier-free guidance.
+        Args:
+            t: The current timestep, a tensor of a single float.
+            x: The initial value, with the shape (batch, seq_len, emb_dim).
+            text_condition: The text_condition of the diffision model, with
+                the shape (batch, seq_len, emb_dim).
+            speech_condition: The speech_condition of the diffision model, with the
+                shape (batch, seq_len, emb_dim).
+            padding_mask: The mask for padding; True means masked position, with the
+                shape (batch, seq_len).
+            guidance_scale: The scale of classifier-free guidance, a float or a tensor
+                of shape (batch, 1, 1).
+        Retrun:
+            The prediction with the shape (batch, seq_len, emb_dim).
+        """
+        if not torch.is_tensor(guidance_scale):
+            guidance_scale = torch.tensor(
+                guidance_scale, dtype=t.dtype, device=t.device
+            )
+        return self.model_func(
+            t=t,
+            xt=x,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+            guidance_scale=guidance_scale,
+            **kwargs
+        )
+
+
+class EulerSolver:
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        func_name: str = "forward_fm_decoder",
+    ):
+        """Construct a Euler Solver
+        Args:
+            model: The diffusion model.
+            func_name: The function name to call.
+        """
+        self.model = DiffusionModel(model, func_name=func_name)
+
+    def sample(
+        self,
+        x: torch.Tensor,
+        text_condition: torch.Tensor,
+        speech_condition: torch.Tensor,
+        padding_mask: torch.Tensor,
+        num_step: int = 10,
+        guidance_scale: Union[float, torch.Tensor] = 0.0,
+        t_start: float = 0.0,
+        t_end: float = 1.0,
+        t_shift: float = 1.0,
+        **kwargs
+    ) -> torch.Tensor:
+        device = x.device
+        assert isinstance(t_start, float) and isinstance(t_end, float)
+
+        # Generate the schedule of timesteps
+        timesteps = get_time_steps(
+            t_start=t_start,
+            t_end=t_end,
+            num_step=num_step,
+            t_shift=t_shift,
+            device=device,
+        )
+
+        for step in range(num_step):
+            t_cur = timesteps[step]
+            t_next = timesteps[step + 1]
+
+            # Predict velocity (v)
+            v = self.model(
+                t=t_cur,
+                x=x,
+                text_condition=text_condition,
+                speech_condition=speech_condition,
+                padding_mask=padding_mask,
+                guidance_scale=guidance_scale,
+                **kwargs
+            )
+
+            # 1. Predict the clean 'data' (x_1) and 'noise' (x_0)
+            # Flow matching formulation: x_t = (1 - t) * x_0 + t * x_1
+            # Therefore: v = x_1 - x_0
+            x_1_pred = x + (1.0 - t_cur) * v
+            x_0_pred = x - t_cur * v
+
+            if step < num_step - 1:
+                # 2. Probability Flow ODE update (Anchor-based)
+                # This 'anchors' the next point along the predicted line, 
+                # making it more robust than simple Euler integration.
+                x = (1.0 - t_next) * x_0_pred + t_next * x_1_pred
+            else:
+                # Final step: Snap directly to the predicted clean data
+                x = x_1_pred
+
+        return x
+
+
+class DistillEulerSolver(EulerSolver):
+    def __init__(
+        self,
+        model: torch.nn.Module,
+        func_name: str = "forward_fm_decoder",
+    ):
+        """Construct a Euler Solver for distilled diffusion models.
+        Args:
+            model: The diffusion model.
+        """
+        self.model = DistillDiffusionModel(model, func_name=func_name)
+
+
+def get_time_steps(
+    t_start: float = 0.0,
+    t_end: float = 1.0,
+    num_step: int = 10,
+    t_shift: float = 1.0,
+    device: torch.device = torch.device("cpu"),
+) -> torch.Tensor:
+    """Compute the intermediate time steps for sampling.
+
+    Args:
+        t_start: The starting time of the sampling (default is 0).
+        t_end: The starting time of the sampling (default is 1).
+        num_step: The number of sampling.
+        t_shift: shift the t toward smaller numbers so that the sampling
+            will emphasize low SNR region. Should be in the range of (0, 1].
+            The shifting will be more significant when the number is smaller.
+        device: A torch device.
+    Returns:
+        The time step with the shape (num_step + 1,).
+    """
+
+    timesteps = torch.linspace(t_start, t_end, num_step + 1).to(device)
+
+    timesteps = t_shift * timesteps / (1 + (t_shift - 1) * timesteps)
+
+    return timesteps

zipvoice/models/modules/zipformer_two_stream.py

@@ -0,0 +1,264 @@
+#!/usr/bin/env python3
+# Copyright    2025  Xiaomi Corp.        (authors: Han Zhu)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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 math
+from typing import Optional, Tuple, Union
+
+import torch
+from torch import Tensor, nn
+
+from zipvoice.models.modules.scaling import FloatLike, ScheduledFloat, SwooshR
+from zipvoice.models.modules.zipformer import (
+    DownsampledZipformer2Encoder,
+    TTSZipformer,
+    Zipformer2Encoder,
+    Zipformer2EncoderLayer,
+)
+
+
+def timestep_embedding(timesteps, dim, max_period=10000):
+    """Create sinusoidal timestep embeddings.
+
+    :param timesteps: shape of (N) or (N, T)
+    :param dim: the dimension of the output.
+    :param max_period: controls the minimum frequency of the embeddings.
+    :return: an Tensor of positional embeddings. shape of (N, dim) or (T, N, dim)
+    """
+    half = dim // 2
+    freqs = torch.exp(
+        -math.log(max_period)
+        * torch.arange(start=0, end=half, dtype=torch.float32, device=timesteps.device)
+        / half
+    )
+
+    if timesteps.dim() == 2:
+        timesteps = timesteps.transpose(0, 1)  # (N, T) -> (T, N)
+
+    args = timesteps[..., None].float() * freqs[None]
+    embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+    if dim % 2:
+        embedding = torch.cat([embedding, torch.zeros_like(embedding[..., :1])], dim=-1)
+    return embedding
+
+
+class TTSZipformerTwoStream(TTSZipformer):
+    """
+    Args:
+
+    Note: all "int or Tuple[int]" arguments below will be treated as lists of the same
+    length as downsampling_factor if they are single ints or one-element tuples.
+    The length of downsampling_factor defines the number of stacks.
+
+        downsampling_factor (Tuple[int]): downsampling factor for each encoder stack.
+           Note: this is in addition to the downsampling factor of 2 that is applied in
+           the frontend (self.encoder_embed).
+        encoder_dim (Tuple[int]): embedding dimension of each of the encoder stacks,
+            one per encoder stack.
+        num_encoder_layers (int or Tuple[int])): number of encoder layers for each stack
+        query_head_dim (int or Tuple[int]): dimension of query and key per attention
+           head: per stack, if a tuple..
+        pos_head_dim (int or Tuple[int]): dimension of positional-encoding projection
+            per attention head
+        value_head_dim (int or Tuple[int]): dimension of value in each attention head
+        num_heads: (int or Tuple[int]): number of heads in the self-attention mechanism.
+              Must be at least 4.
+        feedforward_dim (int or Tuple[int]): hidden dimension in feedforward modules
+        cnn_module_kernel (int or Tuple[int])): Kernel size of convolution module
+
+        pos_dim (int): the dimension of each positional-encoding vector prior to
+            projection, e.g. 128.
+
+        dropout (float): dropout rate
+        warmup_batches (float): number of batches to warm up over; this controls
+          dropout of encoder layers.
+        use_time_embed: (bool): if True, do not take time embedding as additional input.
+        time_embed_dim: (int): the dimension of the time embedding.
+    """
+
+    def __init__(
+        self,
+        in_dim: Tuple[int],
+        out_dim: Tuple[int],
+        downsampling_factor: Tuple[int] = (2, 4),
+        num_encoder_layers: Union[int, Tuple[int]] = 4,
+        cnn_module_kernel: Union[int, Tuple[int]] = 31,
+        encoder_dim: int = 384,
+        query_head_dim: int = 24,
+        pos_head_dim: int = 4,
+        value_head_dim: int = 12,
+        num_heads: int = 8,
+        feedforward_dim: int = 1536,
+        pos_dim: int = 192,
+        dropout: FloatLike = None,  # see code below for default
+        warmup_batches: float = 4000.0,
+        use_time_embed: bool = True,
+        time_embed_dim: int = 192,
+        use_conv: bool = True,
+    ) -> None:
+        nn.Module.__init__(self)
+
+        if dropout is None:
+            dropout = ScheduledFloat((0.0, 0.3), (20000.0, 0.1))
+        if isinstance(downsampling_factor, int):
+            downsampling_factor = (downsampling_factor,)
+
+        def _to_tuple(x):
+            """Converts a single int or a 1-tuple of an int to a tuple with the same
+            length as downsampling_factor"""
+            if isinstance(x, int):
+                x = (x,)
+            if len(x) == 1:
+                x = x * len(downsampling_factor)
+            else:
+                assert len(x) == len(downsampling_factor) and isinstance(x[0], int)
+            return x
+
+        def _assert_downsampling_factor(factors):
+            """assert downsampling_factor follows u-net style"""
+            assert factors[0] == 1 and factors[-1] == 1
+
+            for i in range(1, len(factors) // 2 + 1):
+                assert factors[i] == factors[i - 1] * 2
+
+            for i in range(len(factors) // 2 + 1, len(factors)):
+                assert factors[i] * 2 == factors[i - 1]
+
+        _assert_downsampling_factor(downsampling_factor)
+        self.downsampling_factor = downsampling_factor  # tuple
+        num_encoder_layers = _to_tuple(num_encoder_layers)
+        self.cnn_module_kernel = cnn_module_kernel = _to_tuple(cnn_module_kernel)
+        self.encoder_dim = encoder_dim
+        self.num_encoder_layers = num_encoder_layers
+        self.query_head_dim = query_head_dim
+        self.value_head_dim = value_head_dim
+        self.num_heads = num_heads
+
+        self.use_time_embed = use_time_embed
+
+        self.time_embed_dim = time_embed_dim
+        if self.use_time_embed:
+            assert time_embed_dim != -1
+        else:
+            time_embed_dim = -1
+
+        assert len(in_dim) == len(out_dim) == 2
+
+        self.in_dim = in_dim
+        self.in_proj = nn.ModuleList(
+            [nn.Linear(in_dim[0], encoder_dim), nn.Linear(in_dim[1], encoder_dim)]
+        )
+        self.out_dim = out_dim
+        self.out_proj = nn.ModuleList(
+            [nn.Linear(encoder_dim, out_dim[0]), nn.Linear(encoder_dim, out_dim[1])]
+        )
+
+        # each one will be Zipformer2Encoder or DownsampledZipformer2Encoder
+        encoders = []
+
+        num_encoders = len(downsampling_factor)
+        for i in range(num_encoders):
+            encoder_layer = Zipformer2EncoderLayer(
+                embed_dim=encoder_dim,
+                pos_dim=pos_dim,
+                num_heads=num_heads,
+                query_head_dim=query_head_dim,
+                pos_head_dim=pos_head_dim,
+                value_head_dim=value_head_dim,
+                feedforward_dim=feedforward_dim,
+                use_conv=use_conv,
+                cnn_module_kernel=cnn_module_kernel[i],
+                dropout=dropout,
+            )
+
+            # For the segment of the warmup period, we let the Conv2dSubsampling
+            # layer learn something.  Then we start to warm up the other encoders.
+            encoder = Zipformer2Encoder(
+                encoder_layer,
+                num_encoder_layers[i],
+                embed_dim=encoder_dim,
+                time_embed_dim=time_embed_dim,
+                pos_dim=pos_dim,
+                warmup_begin=warmup_batches * (i + 1) / (num_encoders + 1),
+                warmup_end=warmup_batches * (i + 2) / (num_encoders + 1),
+                final_layerdrop_rate=0.035 * (downsampling_factor[i] ** 0.5),
+            )
+
+            if downsampling_factor[i] != 1:
+                encoder = DownsampledZipformer2Encoder(
+                    encoder,
+                    dim=encoder_dim,
+                    downsample=downsampling_factor[i],
+                )
+
+            encoders.append(encoder)
+
+        self.encoders = nn.ModuleList(encoders)
+        if self.use_time_embed:
+            self.time_embed = nn.Sequential(
+                nn.Linear(time_embed_dim, time_embed_dim * 2),
+                SwooshR(),
+                nn.Linear(time_embed_dim * 2, time_embed_dim),
+            )
+        else:
+            self.time_embed = None
+
+    def forward(
+        self,
+        x: Tensor,
+        t: Optional[Tensor] = None,
+        padding_mask: Optional[Tensor] = None,
+    ) -> Tuple[Tensor, Tensor]:
+        """
+        Args:
+          x:
+            The input tensor. Its shape is (batch_size, seq_len, feature_dim).
+          t:
+            A t tensor of shape (batch_size,) or (batch_size, seq_len)
+          padding_mask:
+            The mask for padding, of shape (batch_size, seq_len); True means
+            masked position. May be None.
+        Returns:
+          Return the output embeddings. its shape is
+            (batch_size, output_seq_len, encoder_dim)
+        """
+        assert x.size(2) in self.in_dim, f"{x.size(2)} in {self.in_dim}"
+        if x.size(2) == self.in_dim[0]:
+            index = 0
+        else:
+            index = 1
+        x = x.permute(1, 0, 2)
+        x = self.in_proj[index](x)
+
+        if t is not None:
+            assert t.dim() == 1 or t.dim() == 2, t.shape
+            time_emb = timestep_embedding(t, self.time_embed_dim)
+            time_emb = self.time_embed(time_emb)
+        else:
+            time_emb = None
+
+        attn_mask = None
+
+        for i, module in enumerate(self.encoders):
+            x = module(
+                x,
+                time_emb=time_emb,
+                src_key_padding_mask=padding_mask,
+                attn_mask=attn_mask,
+            )
+        x = self.out_proj[index](x)
+        x = x.permute(1, 0, 2)
+        return x

zipvoice/models/zipvoice.py

@@ -0,0 +1,534 @@
+# Copyright    2024    Xiaomi Corp.        (authors:  Wei Kang
+#                                                     Han Zhu)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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, Optional
+
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from zipvoice.models.modules.solver import EulerSolver
+from zipvoice.models.modules.zipformer import TTSZipformer
+from zipvoice.utils.common import (
+    condition_time_mask,
+    get_tokens_index,
+    make_pad_mask,
+    pad_labels,
+    prepare_avg_tokens_durations,
+)
+
+
+class ZipVoice(nn.Module):
+    """The ZipVoice model."""
+
+    def __init__(
+        self,
+        fm_decoder_downsampling_factor: List[int] = [1, 2, 4, 2, 1],
+        fm_decoder_num_layers: List[int] = [2, 2, 4, 4, 4],
+        fm_decoder_cnn_module_kernel: List[int] = [31, 15, 7, 15, 31],
+        fm_decoder_feedforward_dim: int = 1536,
+        fm_decoder_num_heads: int = 4,
+        fm_decoder_dim: int = 512,
+        text_encoder_num_layers: int = 4,
+        text_encoder_feedforward_dim: int = 512,
+        text_encoder_cnn_module_kernel: int = 9,
+        text_encoder_num_heads: int = 4,
+        text_encoder_dim: int = 192,
+        time_embed_dim: int = 192,
+        text_embed_dim: int = 192,
+        query_head_dim: int = 32,
+        value_head_dim: int = 12,
+        pos_head_dim: int = 4,
+        pos_dim: int = 48,
+        feat_dim: int = 100,
+        vocab_size: int = 26,
+        pad_id: int = 0,
+    ):
+        """
+        Initialize the model with specified configuration parameters.
+
+        Args:
+            fm_decoder_downsampling_factor: List of downsampling factors for each layer
+                in the flow-matching decoder.
+            fm_decoder_num_layers: List of the number of layers for each block in the
+                flow-matching decoder.
+            fm_decoder_cnn_module_kernel: List of kernel sizes for CNN modules in the
+                flow-matching decoder.
+            fm_decoder_feedforward_dim: Dimension of the feedforward network in the
+                flow-matching decoder.
+            fm_decoder_num_heads: Number of attention heads in the flow-matching
+                decoder.
+            fm_decoder_dim: Hidden dimension of the flow-matching decoder.
+            text_encoder_num_layers: Number of layers in the text encoder.
+            text_encoder_feedforward_dim: Dimension of the feedforward network in the
+                text encoder.
+            text_encoder_cnn_module_kernel: Kernel size for the CNN module in the
+                text encoder.
+            text_encoder_num_heads: Number of attention heads in the text encoder.
+            text_encoder_dim: Hidden dimension of the text encoder.
+            time_embed_dim: Dimension of the time embedding.
+            text_embed_dim: Dimension of the text embedding.
+            query_head_dim: Dimension of the query attention head.
+            value_head_dim: Dimension of the value attention head.
+            pos_head_dim: Dimension of the position attention head.
+            pos_dim: Dimension of the positional encoding.
+            feat_dim: Dimension of the acoustic features.
+            vocab_size: Size of the vocabulary.
+            pad_id: ID used for padding tokens.
+        """
+        super().__init__()
+
+        self.fm_decoder = TTSZipformer(
+            in_dim=feat_dim * 3,
+            out_dim=feat_dim,
+            downsampling_factor=fm_decoder_downsampling_factor,
+            num_encoder_layers=fm_decoder_num_layers,
+            cnn_module_kernel=fm_decoder_cnn_module_kernel,
+            encoder_dim=fm_decoder_dim,
+            feedforward_dim=fm_decoder_feedforward_dim,
+            num_heads=fm_decoder_num_heads,
+            query_head_dim=query_head_dim,
+            pos_head_dim=pos_head_dim,
+            value_head_dim=value_head_dim,
+            pos_dim=pos_dim,
+            use_time_embed=True,
+            time_embed_dim=time_embed_dim,
+        )
+
+        self.text_encoder = TTSZipformer(
+            in_dim=text_embed_dim,
+            out_dim=feat_dim,
+            downsampling_factor=1,
+            num_encoder_layers=text_encoder_num_layers,
+            cnn_module_kernel=text_encoder_cnn_module_kernel,
+            encoder_dim=text_encoder_dim,
+            feedforward_dim=text_encoder_feedforward_dim,
+            num_heads=text_encoder_num_heads,
+            query_head_dim=query_head_dim,
+            pos_head_dim=pos_head_dim,
+            value_head_dim=value_head_dim,
+            pos_dim=pos_dim,
+            use_time_embed=False,
+        )
+
+        self.feat_dim = feat_dim
+        self.text_embed_dim = text_embed_dim
+        self.pad_id = pad_id
+
+        self.embed = nn.Embedding(vocab_size, text_embed_dim)
+        self.solver = EulerSolver(self, func_name="forward_fm_decoder")
+
+    def forward_fm_decoder(
+        self,
+        t: torch.Tensor,
+        xt: torch.Tensor,
+        text_condition: torch.Tensor,
+        speech_condition: torch.Tensor,
+        padding_mask: Optional[torch.Tensor] = None,
+        guidance_scale: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Compute velocity.
+        Args:
+            t:  A tensor of shape (N, 1, 1) or a tensor of a float,
+                in the range of (0, 1).
+            xt: the input of the current timestep, including condition
+                embeddings and noisy acoustic features.
+            text_condition: the text condition embeddings, with the
+                shape (batch, seq_len, emb_dim).
+            speech_condition: the speech condition embeddings, with the
+                shape (batch, seq_len, emb_dim).
+            padding_mask: The mask for padding, True means masked
+                position, with the shape (N, T).
+            guidance_scale: The guidance scale in classifier-free guidance,
+                which is a tensor of shape (N, 1, 1) or a tensor of a float.
+
+        Returns:
+            predicted velocity, with the shape (batch, seq_len, emb_dim).
+        """
+
+        xt = torch.cat([xt, text_condition, speech_condition], dim=2)
+
+        assert t.dim() in (0, 3)
+        # Handle t with the shape (N, 1, 1):
+        # squeeze the last dimension if it's size is 1.
+        while t.dim() > 1 and t.size(-1) == 1:
+            t = t.squeeze(-1)
+        # Handle t with a single value: expand to the size of batch size.
+        if t.dim() == 0:
+            t = t.repeat(xt.shape[0])
+
+        if guidance_scale is not None:
+            while guidance_scale.dim() > 1 and guidance_scale.size(-1) == 1:
+                guidance_scale = guidance_scale.squeeze(-1)
+            if guidance_scale.dim() == 0:
+                guidance_scale = guidance_scale.repeat(xt.shape[0])
+
+            vt = self.fm_decoder(
+                x=xt, t=t, padding_mask=padding_mask, guidance_scale=guidance_scale
+            )
+        else:
+            vt = self.fm_decoder(x=xt, t=t, padding_mask=padding_mask)
+        return vt
+
+    def forward_text_embed(
+        self,
+        tokens: List[List[int]],
+    ):
+        """
+        Get the text embeddings.
+        Args:
+            tokens: a list of list of token ids.
+        Returns:
+            embed: the text embeddings, shape (batch, seq_len, emb_dim).
+            tokens_lens: the length of each token sequence, shape (batch,).
+        """
+        device = (
+            self.device if isinstance(self, DDP) else next(self.parameters()).device
+        )
+        tokens_padded = pad_labels(tokens, pad_id=self.pad_id, device=device)  # (B, S)
+        embed = self.embed(tokens_padded)  # (B, S, C)
+        tokens_lens = torch.tensor(
+            [len(token) for token in tokens], dtype=torch.int64, device=device
+        )
+        tokens_padding_mask = make_pad_mask(tokens_lens, embed.shape[1])  # (B, S)
+
+        embed = self.text_encoder(
+            x=embed, t=None, padding_mask=tokens_padding_mask
+        )  # (B, S, C)
+        return embed, tokens_lens
+
+    def forward_text_condition(
+        self,
+        embed: torch.Tensor,
+        tokens_lens: torch.Tensor,
+        features_lens: torch.Tensor,
+    ):
+        """
+        Get the text condition with the same length of the acoustic feature.
+        Args:
+            embed: the text embeddings, shape (batch, token_seq_len, emb_dim).
+            tokens_lens: the length of each token sequence, shape (batch,).
+            features_lens: the length of each acoustic feature sequence,
+                shape (batch,).
+        Returns:
+            text_condition: the text condition, shape
+                (batch, feature_seq_len, emb_dim).
+            padding_mask: the padding mask of text condition, shape
+                (batch, feature_seq_len).
+        """
+
+        num_frames = int(features_lens.max())
+
+        padding_mask = make_pad_mask(features_lens, max_len=num_frames)  # (B, T)
+
+        tokens_durations = prepare_avg_tokens_durations(features_lens, tokens_lens)
+
+        tokens_index = get_tokens_index(tokens_durations, num_frames).to(
+            embed.device
+        )  # (B, T)
+
+        text_condition = torch.gather(
+            embed,
+            dim=1,
+            index=tokens_index.unsqueeze(-1).expand(
+                embed.size(0), num_frames, embed.size(-1)
+            ),
+        )  # (B, T, F)
+        return text_condition, padding_mask
+
+    def forward_text_train(
+        self,
+        tokens: List[List[int]],
+        features_lens: torch.Tensor,
+    ):
+        """
+        Process text for training, given text tokens and real feature lengths.
+        """
+        embed, tokens_lens = self.forward_text_embed(tokens)
+        text_condition, padding_mask = self.forward_text_condition(
+            embed, tokens_lens, features_lens
+        )
+        return (
+            text_condition,
+            padding_mask,
+        )
+
+    def forward_text_inference_gt_duration(
+        self,
+        tokens: List[List[int]],
+        features_lens: torch.Tensor,
+        prompt_tokens: List[List[int]],
+        prompt_features_lens: torch.Tensor,
+    ):
+        """
+        Process text for inference, given text tokens, real feature lengths and prompts.
+        """
+        tokens = [
+            prompt_token + token for prompt_token, token in zip(prompt_tokens, tokens)
+        ]
+        features_lens = prompt_features_lens + features_lens
+        embed, tokens_lens = self.forward_text_embed(tokens)
+        text_condition, padding_mask = self.forward_text_condition(
+            embed, tokens_lens, features_lens
+        )
+        return text_condition, padding_mask
+
+    def forward_text_inference_ratio_duration(
+        self,
+        tokens: List[List[int]],
+        prompt_tokens: List[List[int]],
+        prompt_features_lens: torch.Tensor,
+        speed: float,
+    ):
+        """
+        Process text for inference, given text tokens and prompts,
+        feature lengths are predicted with the ratio of token numbers.
+        """
+        device = (
+            self.device if isinstance(self, DDP) else next(self.parameters()).device
+        )
+
+        cat_tokens = [
+            prompt_token + token for prompt_token, token in zip(prompt_tokens, tokens)
+        ]
+
+        prompt_tokens_lens = torch.tensor(
+            [len(token) for token in prompt_tokens],
+            dtype=torch.int64,
+            device=device,
+        )
+
+        tokens_lens = torch.tensor(
+            [len(token) for token in tokens],
+            dtype=torch.int64,
+            device=device,
+        )
+
+        cat_embed, cat_tokens_lens = self.forward_text_embed(cat_tokens)
+
+        features_lens = prompt_features_lens + torch.ceil(
+            (prompt_features_lens / prompt_tokens_lens * tokens_lens / speed)
+        ).to(dtype=torch.int64)
+
+        text_condition, padding_mask = self.forward_text_condition(
+            cat_embed, cat_tokens_lens, features_lens
+        )
+        return text_condition, padding_mask
+
+    def forward(
+        self,
+        tokens: List[List[int]],
+        features: torch.Tensor,
+        features_lens: torch.Tensor,
+        noise: torch.Tensor,
+        t: torch.Tensor,
+        condition_drop_ratio: float = 0.0,
+    ) -> torch.Tensor:
+        """Forward pass of the model for training.
+        Args:
+            tokens: a list of list of token ids.
+            features: the acoustic features, with the shape (batch, seq_len, feat_dim).
+            features_lens: the length of each acoustic feature sequence, shape (batch,).
+            noise: the intitial noise, with the shape (batch, seq_len, feat_dim).
+            t: the time step, with the shape (batch, 1, 1).
+            condition_drop_ratio: the ratio of dropped text condition.
+        Returns:
+            fm_loss: the flow-matching loss.
+        """
+
+        (text_condition, padding_mask,) = self.forward_text_train(
+            tokens=tokens,
+            features_lens=features_lens,
+        )
+
+        speech_condition_mask = condition_time_mask(
+            features_lens=features_lens,
+            mask_percent=(0.7, 1.0),
+            max_len=features.size(1),
+        )
+        speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features)
+
+        if condition_drop_ratio > 0.0:
+            drop_mask = (
+                torch.rand(text_condition.size(0), 1, 1).to(text_condition.device)
+                > condition_drop_ratio
+            )
+            text_condition = text_condition * drop_mask
+
+        xt = features * t + noise * (1 - t)
+        ut = features - noise  # (B, T, F)
+
+        vt = self.forward_fm_decoder(
+            t=t,
+            xt=xt,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+        )
+
+        loss_mask = speech_condition_mask & (~padding_mask)
+        fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2)
+
+        return fm_loss
+
+    def sample(
+        self,
+        tokens: List[List[int]],
+        prompt_tokens: List[List[int]],
+        prompt_features: torch.Tensor,
+        prompt_features_lens: torch.Tensor,
+        features_lens: Optional[torch.Tensor] = None,
+        speed: float = 1.0,
+        t_shift: float = 1.0,
+        duration: str = "predict",
+        num_step: int = 5,
+        guidance_scale: float = 0.5,
+    ) -> torch.Tensor:
+        """
+        Generate acoustic features, given text tokens, prompts feature
+            and prompt transcription's text tokens.
+        Args:
+            tokens: a list of list of text tokens.
+            prompt_tokens: a list of list of prompt tokens.
+            prompt_features: the prompt feature with the shape
+                (batch_size, seq_len, feat_dim).
+            prompt_features_lens: the length of each prompt feature,
+                with the shape (batch_size,).
+            features_lens: the length of the predicted eature, with the
+                shape (batch_size,). It is used only when duration is "real".
+            duration: "real" or "predict". If "real", the predicted
+                feature length is given by features_lens.
+            num_step: the number of steps to use in the ODE solver.
+            guidance_scale: the guidance scale for classifier-free guidance.
+        """
+
+        assert duration in ["real", "predict"]
+
+        if duration == "predict":
+            (
+                text_condition,
+                padding_mask,
+            ) = self.forward_text_inference_ratio_duration(
+                tokens=tokens,
+                prompt_tokens=prompt_tokens,
+                prompt_features_lens=prompt_features_lens,
+                speed=speed,
+            )
+        else:
+            assert features_lens is not None
+            text_condition, padding_mask = self.forward_text_inference_gt_duration(
+                tokens=tokens,
+                features_lens=features_lens,
+                prompt_tokens=prompt_tokens,
+                prompt_features_lens=prompt_features_lens,
+            )
+        batch_size, num_frames, _ = text_condition.shape
+
+        speech_condition = torch.nn.functional.pad(
+            prompt_features, (0, 0, 0, num_frames - prompt_features.size(1))
+        )  # (B, T, F)
+
+        # False means speech condition positions.
+        speech_condition_mask = make_pad_mask(prompt_features_lens, num_frames)
+        speech_condition = torch.where(
+            speech_condition_mask.unsqueeze(-1),
+            torch.zeros_like(speech_condition),
+            speech_condition,
+        )
+
+        x0 = torch.randn(
+            batch_size,
+            num_frames,
+            prompt_features.size(-1),
+            device=text_condition.device,
+        )
+
+        x1 = self.solver.sample(
+            x=x0,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+            num_step=num_step,
+            guidance_scale=guidance_scale,
+            t_shift=t_shift,
+        )
+        x1_wo_prompt_lens = (~padding_mask).sum(-1) - prompt_features_lens
+        x1_prompt = torch.zeros(
+            x1.size(0), prompt_features_lens.max(), x1.size(2), device=x1.device
+        )
+        x1_wo_prompt = torch.zeros(
+            x1.size(0), x1_wo_prompt_lens.max(), x1.size(2), device=x1.device
+        )
+        for i in range(x1.size(0)):
+            x1_wo_prompt[i, : x1_wo_prompt_lens[i], :] = x1[
+                i,
+                prompt_features_lens[i] : prompt_features_lens[i]
+                + x1_wo_prompt_lens[i],
+            ]
+            x1_prompt[i, : prompt_features_lens[i], :] = x1[
+                i, : prompt_features_lens[i]
+            ]
+
+        return x1_wo_prompt, x1_wo_prompt_lens, x1_prompt, prompt_features_lens
+
+    def sample_intermediate(
+        self,
+        tokens: List[List[int]],
+        features: torch.Tensor,
+        features_lens: torch.Tensor,
+        noise: torch.Tensor,
+        speech_condition_mask: torch.Tensor,
+        t_start: float,
+        t_end: float,
+        num_step: int = 1,
+        guidance_scale: torch.Tensor = None,
+    ) -> torch.Tensor:
+        """
+        Generate acoustic features in intermediate timesteps.
+        Args:
+            tokens: List of list of token ids.
+            features: The acoustic features, with the shape (batch, seq_len, feat_dim).
+            features_lens: The length of each acoustic feature sequence,
+                with the shape (batch,).
+            noise: The initial noise, with the shape (batch, seq_len, feat_dim).
+            speech_condition_mask: The mask for speech condition, True means
+                non-condition positions, with the shape (batch, seq_len).
+            t_start: The start timestep.
+            t_end: The end timestep.
+            num_step: The number of steps for sampling.
+            guidance_scale: The scale for classifier-free guidance inference,
+                with the shape (batch, 1, 1).
+        """
+        (text_condition, padding_mask,) = self.forward_text_train(
+            tokens=tokens,
+            features_lens=features_lens,
+        )
+
+        speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features)
+
+        x_t_end = self.solver.sample(
+            x=noise,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+            num_step=num_step,
+            guidance_scale=guidance_scale,
+            t_start=t_start,
+            t_end=t_end,
+        )
+        x_t_end_lens = (~padding_mask).sum(-1)
+        return x_t_end, x_t_end_lens

zipvoice/models/zipvoice_dialog.py

@@ -0,0 +1,358 @@
+# Copyright    2025    Xiaomi Corp.        (authors:  Han Zhu)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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 torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from zipvoice.models.modules.zipformer_two_stream import TTSZipformerTwoStream
+from zipvoice.models.zipvoice import ZipVoice
+from zipvoice.utils.common import condition_time_mask_suffix, make_pad_mask, pad_labels
+
+
+class ZipVoiceDialog(ZipVoice):
+    """The ZipVoice-Dialog model."""
+
+    def __init__(
+        self,
+        fm_decoder_downsampling_factor: List[int] = [1, 2, 4, 2, 1],
+        fm_decoder_num_layers: List[int] = [2, 2, 4, 4, 4],
+        fm_decoder_cnn_module_kernel: List[int] = [31, 15, 7, 15, 31],
+        fm_decoder_feedforward_dim: int = 1536,
+        fm_decoder_num_heads: int = 4,
+        fm_decoder_dim: int = 512,
+        text_encoder_num_layers: int = 4,
+        text_encoder_feedforward_dim: int = 512,
+        text_encoder_cnn_module_kernel: int = 9,
+        text_encoder_num_heads: int = 4,
+        text_encoder_dim: int = 192,
+        time_embed_dim: int = 192,
+        text_embed_dim: int = 192,
+        query_head_dim: int = 32,
+        value_head_dim: int = 12,
+        pos_head_dim: int = 4,
+        pos_dim: int = 48,
+        feat_dim: int = 100,
+        vocab_size: int = 26,
+        pad_id: int = 0,
+        spk_a_id: int = 360,
+        spk_b_id: int = 361,
+    ):
+        """
+        Initialize the model with specified configuration parameters.
+
+        Args:
+            fm_decoder_downsampling_factor: List of downsampling factors for each layer
+                in the flow-matching decoder.
+            fm_decoder_num_layers: List of the number of layers for each block in the
+                flow-matching decoder.
+            fm_decoder_cnn_module_kernel: List of kernel sizes for CNN modules in the
+                flow-matching decoder.
+            fm_decoder_feedforward_dim: Dimension of the feedforward network in the
+                flow-matching decoder.
+            fm_decoder_num_heads: Number of attention heads in the flow-matching
+                decoder.
+            fm_decoder_dim: Hidden dimension of the flow-matching decoder.
+            text_encoder_num_layers: Number of layers in the text encoder.
+            text_encoder_feedforward_dim: Dimension of the feedforward network in the
+                text encoder.
+            text_encoder_cnn_module_kernel: Kernel size for the CNN module in the
+                text encoder.
+            text_encoder_num_heads: Number of attention heads in the text encoder.
+            text_encoder_dim: Hidden dimension of the text encoder.
+            time_embed_dim: Dimension of the time embedding.
+            text_embed_dim: Dimension of the text embedding.
+            query_head_dim: Dimension of the query attention head.
+            value_head_dim: Dimension of the value attention head.
+            pos_head_dim: Dimension of the position attention head.
+            pos_dim: Dimension of the positional encoding.
+            feat_dim: Dimension of the acoustic features.
+            vocab_size: Size of the vocabulary.
+            pad_id: ID used for padding tokens.
+            spk_a_id: ID of speaker A / [S1].
+            spk_b_id: ID of speaker B / [S2].
+        """
+        super().__init__(
+            fm_decoder_downsampling_factor=fm_decoder_downsampling_factor,
+            fm_decoder_num_layers=fm_decoder_num_layers,
+            fm_decoder_cnn_module_kernel=fm_decoder_cnn_module_kernel,
+            fm_decoder_feedforward_dim=fm_decoder_feedforward_dim,
+            fm_decoder_num_heads=fm_decoder_num_heads,
+            fm_decoder_dim=fm_decoder_dim,
+            text_encoder_num_layers=text_encoder_num_layers,
+            text_encoder_feedforward_dim=text_encoder_feedforward_dim,
+            text_encoder_cnn_module_kernel=text_encoder_cnn_module_kernel,
+            text_encoder_num_heads=text_encoder_num_heads,
+            text_encoder_dim=text_encoder_dim,
+            time_embed_dim=time_embed_dim,
+            text_embed_dim=text_embed_dim,
+            query_head_dim=query_head_dim,
+            value_head_dim=value_head_dim,
+            pos_head_dim=pos_head_dim,
+            pos_dim=pos_dim,
+            feat_dim=feat_dim,
+            vocab_size=vocab_size,
+            pad_id=pad_id,
+        )
+
+        self.spk_a_id = spk_a_id
+        self.spk_b_id = spk_b_id
+        self.spk_embed = nn.Embedding(2, feat_dim)
+        torch.nn.init.normal_(self.spk_embed.weight, mean=0, std=0.1)
+
+    def extract_spk_indices(self, tensor):
+        turn_mask = ((tensor == self.spk_a_id) | (tensor == self.spk_b_id)).long()
+        turn_counts = turn_mask.cumsum(dim=1)
+        spk_mask = turn_counts % 2
+        spk_mask = torch.where(tensor == self.pad_id, -1, spk_mask)
+        spk_a_indices = torch.where(spk_mask == 0)
+        spk_b_indices = torch.where(spk_mask == 1)
+        return spk_a_indices, spk_b_indices
+
+    def forward_text_embed(
+        self,
+        tokens: List[List[int]],
+    ):
+        """
+        Get the text embeddings.
+        Args:
+            tokens: a list of list of token ids.
+        Returns:
+            embed: the text embeddings, shape (batch, seq_len, emb_dim).
+            tokens_lens: the length of each token sequence, shape (batch,).
+        """
+        device = (
+            self.device if isinstance(self, DDP) else next(self.parameters()).device
+        )
+        tokens_padded = pad_labels(tokens, pad_id=self.pad_id, device=device)  # (B, S)
+        embed = self.embed(tokens_padded)  # (B, S, C)
+        spk_a_indices, spk_b_indices = self.extract_spk_indices(tokens_padded)
+        tokens_lens = torch.tensor(
+            [len(token) for token in tokens], dtype=torch.int64, device=device
+        )
+        tokens_padding_mask = make_pad_mask(tokens_lens, embed.shape[1])  # (B, S)
+
+        embed = self.text_encoder(
+            x=embed, t=None, padding_mask=tokens_padding_mask
+        )  # (B, S, C)
+        embed[spk_a_indices] += self.spk_embed(torch.tensor(0, device=device)).to(
+            embed.dtype
+        )
+        embed[spk_b_indices] += self.spk_embed(torch.tensor(1, device=device)).to(
+            embed.dtype
+        )
+        return embed, tokens_lens
+
+    def forward(
+        self,
+        tokens: List[List[int]],
+        features: torch.Tensor,
+        features_lens: torch.Tensor,
+        noise: torch.Tensor,
+        t: torch.Tensor,
+        condition_drop_ratio: float = 0.0,
+    ) -> torch.Tensor:
+        """Forward pass of the model for training.
+        Args:
+            tokens: a list of list of token ids.
+            features: the acoustic features, with the shape (batch, seq_len, feat_dim).
+            features_lens: the length of each acoustic feature sequence, shape (batch,).
+            noise: the intitial noise, with the shape (batch, seq_len, feat_dim).
+            t: the time step, with the shape (batch, 1, 1).
+            condition_drop_ratio: the ratio of dropped text condition.
+        Returns:
+            fm_loss: the flow-matching loss.
+        """
+
+        (text_condition, padding_mask,) = self.forward_text_train(
+            tokens=tokens,
+            features_lens=features_lens,
+        )
+
+        speech_condition_mask = condition_time_mask_suffix(
+            features_lens=features_lens,
+            mask_percent=(0.5, 1.0),
+            max_len=features.size(1),
+        )
+        speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features)
+
+        if condition_drop_ratio > 0.0:
+            drop_mask = (
+                torch.rand(text_condition.size(0), 1, 1).to(text_condition.device)
+                > condition_drop_ratio
+            )
+            text_condition = text_condition * drop_mask
+
+        xt = features * t + noise * (1 - t)
+        ut = features - noise  # (B, T, F)
+
+        vt = self.forward_fm_decoder(
+            t=t,
+            xt=xt,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+        )
+
+        loss_mask = speech_condition_mask & (~padding_mask)
+        fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2)
+
+        return fm_loss
+
+
+class ZipVoiceDialogStereo(ZipVoiceDialog):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        required_params = {
+            "feat_dim",
+            "fm_decoder_downsampling_factor",
+            "fm_decoder_num_layers",
+            "fm_decoder_cnn_module_kernel",
+            "fm_decoder_dim",
+            "fm_decoder_feedforward_dim",
+            "fm_decoder_num_heads",
+            "query_head_dim",
+            "pos_head_dim",
+            "value_head_dim",
+            "pos_dim",
+            "time_embed_dim",
+        }
+
+        missing = [p for p in required_params if p not in kwargs]
+        if missing:
+            raise ValueError(f"Missing required parameters: {', '.join(missing)}")
+
+        self.fm_decoder = TTSZipformerTwoStream(
+            in_dim=(kwargs["feat_dim"] * 5, kwargs["feat_dim"] * 3),
+            out_dim=(kwargs["feat_dim"] * 2, kwargs["feat_dim"]),
+            downsampling_factor=kwargs["fm_decoder_downsampling_factor"],
+            num_encoder_layers=kwargs["fm_decoder_num_layers"],
+            cnn_module_kernel=kwargs["fm_decoder_cnn_module_kernel"],
+            encoder_dim=kwargs["fm_decoder_dim"],
+            feedforward_dim=kwargs["fm_decoder_feedforward_dim"],
+            num_heads=kwargs["fm_decoder_num_heads"],
+            query_head_dim=kwargs["query_head_dim"],
+            pos_head_dim=kwargs["pos_head_dim"],
+            value_head_dim=kwargs["value_head_dim"],
+            pos_dim=kwargs["pos_dim"],
+            use_time_embed=True,
+            time_embed_dim=kwargs["time_embed_dim"],
+        )
+
+    def forward(
+        self,
+        tokens: List[List[int]],
+        features: torch.Tensor,
+        features_lens: torch.Tensor,
+        noise: torch.Tensor,
+        t: torch.Tensor,
+        condition_drop_ratio: float = 0.0,
+        se_weight: float = 1.0,
+    ) -> torch.Tensor:
+        """Forward pass of the model for training.
+        Args:
+            tokens: a list of list of token ids.
+            features: the acoustic features, with the shape (batch, seq_len, feat_dim).
+            features_lens: the length of each acoustic feature sequence, shape (batch,).
+            noise: the intitial noise, with the shape (batch, seq_len, feat_dim).
+            t: the time step, with the shape (batch, 1, 1).
+            condition_drop_ratio: the ratio of dropped text condition.
+            se_weight: the weight of the speaker exclusive loss.
+        Returns:
+            fm_loss: the flow-matching loss.
+        """
+
+        (text_condition, padding_mask,) = self.forward_text_train(
+            tokens=tokens,
+            features_lens=features_lens,
+        )
+
+        speech_condition_mask = condition_time_mask_suffix(
+            features_lens=features_lens,
+            mask_percent=(0.5, 1.0),
+            max_len=features.size(1),
+        )
+        speech_condition = torch.where(speech_condition_mask.unsqueeze(-1), 0, features)
+
+        if condition_drop_ratio > 0.0:
+            drop_mask = (
+                torch.rand(text_condition.size(0), 1, 1).to(text_condition.device)
+                > condition_drop_ratio
+            )
+            text_condition = text_condition * drop_mask
+
+        xt = features * t + noise * (1 - t)
+        ut = features - noise  # (B, T, F)
+
+        vt = self.forward_fm_decoder(
+            t=t,
+            xt=xt,
+            text_condition=text_condition,
+            speech_condition=speech_condition,
+            padding_mask=padding_mask,
+        )
+
+        loss_mask = speech_condition_mask & (~padding_mask)
+        fm_loss = torch.mean((vt[loss_mask] - ut[loss_mask]) ** 2)
+
+        if se_weight > 0:
+            target = xt + vt * (1 - t)
+            fbank_1 = target[:, :, : self.feat_dim]
+            fbank_2 = target[:, :, self.feat_dim :]
+            energy_loss = torch.mean(
+                self.energy_based_loss(fbank_1, fbank_2, features)[loss_mask]
+            )
+            loss = fm_loss + energy_loss * se_weight
+        else:
+            loss = fm_loss
+
+        return loss
+
+    def energy_based_loss(self, fbank1, fbank2, gt_fbank):
+        energy1 = self.energy(fbank1)
+        energy2 = self.energy(fbank2)
+
+        energy_thresholds = self.adaptive_threshold_from_gt(
+            torch.cat(
+                [
+                    gt_fbank[:, :, : self.feat_dim],
+                    gt_fbank[:, :, self.feat_dim :],
+                ],
+                dim=1,
+            )
+        )
+
+        both_speaking = (
+            (energy1 > energy_thresholds) & (energy2 > energy_thresholds)
+        ).float()
+
+        penalty = (
+            both_speaking
+            * (energy1 - energy_thresholds)
+            * (energy2 - energy_thresholds)
+        )
+        return penalty
+
+    def energy(self, fbank):
+        return torch.mean(fbank, dim=-1)
+
+    def adaptive_threshold_from_gt(self, gt_fbank, percentile=50):
+        frame_energies = self.energy(gt_fbank)
+        thresholds = torch.quantile(frame_energies, q=percentile / 100, dim=1)
+        return thresholds.unsqueeze(1)

zipvoice/models/zipvoice_distill.py

@@ -0,0 +1,94 @@
+# Copyright    2024    Xiaomi Corp.        (authors:  Wei Kang
+#                                                     Han Zhu)
+#
+# See ../../../../LICENSE for clarification regarding multiple authors
+#
+# 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 torch
+
+from zipvoice.models.modules.solver import DistillEulerSolver
+from zipvoice.models.modules.zipformer import TTSZipformer
+from zipvoice.models.zipvoice import ZipVoice
+
+
+class ZipVoiceDistill(ZipVoice):
+    """ZipVoice-Distill model."""
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+        required_params = {
+            "feat_dim",
+            "fm_decoder_downsampling_factor",
+            "fm_decoder_num_layers",
+            "fm_decoder_cnn_module_kernel",
+            "fm_decoder_dim",
+            "fm_decoder_feedforward_dim",
+            "fm_decoder_num_heads",
+            "query_head_dim",
+            "pos_head_dim",
+            "value_head_dim",
+            "pos_dim",
+            "time_embed_dim",
+        }
+
+        missing = [p for p in required_params if p not in kwargs]
+        if missing:
+            raise ValueError(f"Missing required parameters: {', '.join(missing)}")
+
+        self.fm_decoder = TTSZipformer(
+            in_dim=kwargs["feat_dim"] * 3,
+            out_dim=kwargs["feat_dim"],
+            downsampling_factor=kwargs["fm_decoder_downsampling_factor"],
+            num_encoder_layers=kwargs["fm_decoder_num_layers"],
+            cnn_module_kernel=kwargs["fm_decoder_cnn_module_kernel"],
+            encoder_dim=kwargs["fm_decoder_dim"],
+            feedforward_dim=kwargs["fm_decoder_feedforward_dim"],
+            num_heads=kwargs["fm_decoder_num_heads"],
+            query_head_dim=kwargs["query_head_dim"],
+            pos_head_dim=kwargs["pos_head_dim"],
+            value_head_dim=kwargs["value_head_dim"],
+            pos_dim=kwargs["pos_dim"],
+            use_time_embed=True,
+            time_embed_dim=kwargs["time_embed_dim"],
+            use_guidance_scale_embed=True,
+        )
+        self.solver = DistillEulerSolver(self, func_name="forward_fm_decoder")
+
+    def forward(
+        self,
+        tokens: List[List[int]],
+        features: torch.Tensor,
+        features_lens: torch.Tensor,
+        noise: torch.Tensor,
+        speech_condition_mask: torch.Tensor,
+        t_start: float,
+        t_end: float,
+        num_step: int = 1,
+        guidance_scale: torch.Tensor = None,
+    ) -> torch.Tensor:
+
+        return self.sample_intermediate(
+            tokens=tokens,
+            features=features,
+            features_lens=features_lens,
+            noise=noise,
+            speech_condition_mask=speech_condition_mask,
+            t_start=t_start,
+            t_end=t_end,
+            num_step=num_step,
+            guidance_scale=guidance_scale,
+        )