xinference 1.4.1__py3-none-any.whl → 1.5.0.post1__py3-none-any.whl

xinference/thirdparty/megatts3/tts/modules/llm_dit/transformer.py

@@ -0,0 +1,230 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 Any, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+from torch import nn
+
+
+def precompute_freqs_cis(dim: int, end: int, theta: float = 10000.0):
+    freqs = 1.0 / (theta ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+    t = torch.arange(end, device=freqs.device)  # type: ignore
+    freqs = torch.outer(t, freqs).float()  # type: ignore
+    freqs_cis = torch.polar(torch.ones_like(freqs), freqs)  # complex64
+    return freqs_cis
+
+
+def reshape_for_broadcast(freqs_cis: torch.Tensor, x: torch.Tensor):
+    ndim = x.ndim
+    assert 0 <= 1 < ndim
+    assert freqs_cis.shape == (x.shape[1], x.shape[-1])
+    shape = [d if i == 1 or i == ndim - 1 else 1 for i, d in enumerate(x.shape)]
+    return freqs_cis.view(*shape)
+
+
+def apply_rotary_emb(
+        xq: torch.Tensor,
+        xk: torch.Tensor,
+        freqs_cis: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-1], -1, 2))
+    xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-1], -1, 2))
+    freqs_cis = reshape_for_broadcast(freqs_cis, xq_)
+    xq_out = torch.view_as_real(xq_ * freqs_cis).flatten(3)
+    xk_out = torch.view_as_real(xk_ * freqs_cis).flatten(3)
+    return xq_out.type_as(xq), xk_out.type_as(xk)
+
+
+class AdaLNZero(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 6)
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb=None):
+        emb = self.linear(self.silu(emb))
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
+        x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+class AdaLNZero_Out(nn.Module):
+    def __init__(self, dim):
+        super().__init__()
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(dim, dim * 2)
+        self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(self, x, emb):
+        emb = self.linear(self.silu(emb))
+        scale, shift = torch.chunk(emb, 2, dim=1)
+        x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :]
+        return x
+
+
+class Attention(nn.Module):
+    def __init__(self, encoder_dim, encoder_n_heads, max_seq_len):
+        super().__init__()
+        self.encoder_n_kv_heads = encoder_n_heads
+        model_parallel_size = 1
+        self.n_local_heads = encoder_n_heads // model_parallel_size
+        self.n_local_kv_heads = self.encoder_n_kv_heads // model_parallel_size
+        self.n_rep = self.n_local_heads // self.n_local_kv_heads
+        self.head_dim = encoder_dim // encoder_n_heads
+
+        self.wq = nn.Linear(
+            encoder_dim,
+            encoder_n_heads * self.head_dim,
+        )
+        self.wk = nn.Linear(
+            encoder_dim,
+            self.encoder_n_kv_heads * self.head_dim,
+        )
+        self.wv = nn.Linear(
+            encoder_dim,
+            self.encoder_n_kv_heads * self.head_dim,
+        )
+        self.wo = nn.Linear(
+            encoder_n_heads * self.head_dim,
+            encoder_dim,
+        )
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            start_pos: int,
+            freqs_cis: torch.Tensor,
+            mask: Optional[torch.Tensor],
+    ):
+        bsz, seqlen, _ = x.shape
+        xq, xk, xv = self.wq(x), self.wk(x), self.wv(x)
+        xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim)
+        xk = xk.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+        xv = xv.view(bsz, seqlen, self.n_local_kv_heads, self.head_dim)
+
+        xq, xk = apply_rotary_emb(xq, xk, freqs_cis=freqs_cis)
+        xq = xq.transpose(1, 2)  # (bs, n_local_heads, seqlen, head_dim)
+        keys = xk.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
+        values = xv.transpose(1, 2)  # (bs, n_local_heads, cache_len + seqlen, head_dim)
+
+        output = F.scaled_dot_product_attention(xq, keys, values, mask[:, None, None, :], is_causal=False)
+        output = output.transpose(1, 2).contiguous().view(bsz, seqlen, -1)
+        return self.wo(output)
+
+
+class FeedForward(nn.Module):
+    def __init__(
+            self,
+            dim: int,
+            hidden_dim: int,
+            multiple_of: int,
+            ffn_dim_multiplier: Optional[float],
+    ):
+        super().__init__()
+        if ffn_dim_multiplier is not None:
+            hidden_dim = int(ffn_dim_multiplier * hidden_dim)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+
+        self.w1 = nn.Linear(
+            dim, hidden_dim
+        )
+        self.w2 = nn.Linear(
+            hidden_dim, dim
+        )
+
+    def forward(self, x):
+        return self.w2(F.silu(self.w1(x)))
+
+
+class TransformerBlock(nn.Module):
+    def __init__(self, encoder_dim, encoder_n_heads, max_seq_len):
+        super().__init__()
+        self.encoder_n_heads = encoder_n_heads
+        self.encoder_dim = encoder_dim
+        self.head_dim = encoder_dim // encoder_n_heads
+        self.attention = Attention(encoder_dim, encoder_n_heads, max_seq_len)
+        self.feed_forward = FeedForward(
+            dim=encoder_dim,
+            hidden_dim=2 * encoder_dim,
+            multiple_of=256,
+            ffn_dim_multiplier=None,
+        )
+        self.attention_norm = AdaLNZero(encoder_dim)
+        self.ffn_norm = nn.LayerNorm(encoder_dim, elementwise_affine=False, eps=1e-6)
+
+    def forward(
+            self,
+            x: torch.Tensor,
+            t: torch.Tensor,
+            start_pos: int,
+            freqs_cis: torch.Tensor,
+            mask: Optional[torch.Tensor],
+    ):
+        """
+        Perform a forward pass through the TransformerBlock.
+
+        Args:
+            x (torch.Tensor): Input tensor.
+            start_pos (int): Starting position for attention caching.
+            freqs_cis (torch.Tensor): Precomputed cosine and sine frequencies.
+            mask (torch.Tensor, optional): Masking tensor for attention. Defaults to None.
+
+        Returns:
+            torch.Tensor: Output tensor after applying attention and feedforward layers.
+
+        """
+        # pre-norm & modulation for attention input
+        norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attention_norm(x, emb=t)
+
+        # attention
+        attn_output = self.attention(norm, start_pos, freqs_cis, mask=mask)
+
+        # process attention output for input x
+        h = x + gate_msa.unsqueeze(1) * attn_output
+
+        norm = self.ffn_norm(h) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
+        ff_output = self.feed_forward(norm)
+        out = h + gate_mlp.unsqueeze(1) * ff_output
+
+        return out
+
+
+class Transformer(nn.Module):
+    def __init__(self, encoder_n_layers, encoder_dim, encoder_n_heads, max_seq_len):
+        super().__init__()
+        # Decoder
+        self.layers = torch.nn.ModuleList()
+        for _ in range(encoder_n_layers):
+            self.layers.append(TransformerBlock(encoder_dim, encoder_n_heads, max_seq_len))
+
+        self.norm = AdaLNZero_Out(encoder_dim)
+        self.out_proj = nn.Linear(encoder_dim, encoder_dim)
+
+        # Rope embedding
+        freqs_cis = precompute_freqs_cis(
+            encoder_dim // encoder_n_heads, max_seq_len
+        )
+        self.register_buffer("freqs_cis", torch.view_as_real(freqs_cis), persistent=False)
+    
+    def forward(self, x, t, attn_mask, start_pos=0):
+        freqs_cis = torch.view_as_complex(self.freqs_cis.float())[start_pos: start_pos + x.size(1)]
+        for i, layer in enumerate(self.layers):
+            x = layer(x, t, start_pos, freqs_cis, attn_mask)
+        x = self.norm(x, t)
+        x = self.out_proj(x)
+        return x

xinference/thirdparty/megatts3/tts/modules/wavvae/decoder/diag_gaussian.py

@@ -0,0 +1,67 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 numpy as np
+
+class DiagonalGaussianDistribution(object):
+    def __init__(self, parameters: torch.Tensor, deterministic: bool = 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, device=self.parameters.device, dtype=self.parameters.dtype
+            )
+
+    def sample(self, generator=None) -> torch.Tensor:
+        # make sure sample is on the same device as the parameters and has same dtype
+        sample = torch.randn(
+            self.mean.shape,
+            generator=generator,
+            device=self.parameters.device,
+            dtype=self.parameters.dtype,
+        )
+        x = self.mean + self.std * sample
+        return x
+
+    def kl(self, other: "DiagonalGaussianDistribution" = None) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.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) -> torch.Tensor:
+        if self.deterministic:
+            return torch.Tensor([0.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) -> torch.Tensor:
+        return self.mean

xinference/thirdparty/megatts3/tts/modules/wavvae/decoder/hifigan_modules.py

@@ -0,0 +1,283 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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.nn as nn
+import torch.nn.functional as F
+import torch
+import torch.utils.data
+from librosa.filters import mel as librosa_mel_fn
+from torch.nn.utils import weight_norm, remove_weight_norm
+from torch.nn import Conv1d
+import numpy as np
+
+
+def init_weights(m, mean=0.0, std=0.01):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(mean, std)
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size*dilation - dilation)/2)
+
+
+class Upsample(nn.Module):
+    def __init__(self, mult, r):
+        super(Upsample, self).__init__()
+        self.r = r
+        self.upsample = nn.Sequential(nn.Upsample(mode="nearest", scale_factor=r),
+                                      nn.LeakyReLU(0.2),
+                                      nn.ReflectionPad1d(3),
+                                      nn.utils.weight_norm(nn.Conv1d(mult, mult // 2, kernel_size=7, stride=1))
+                                      )
+        r_kernel = r if r >= 5 else 5
+        self.trans_upsample = nn.Sequential(nn.LeakyReLU(0.2),
+                                            nn.utils.weight_norm(nn.ConvTranspose1d(mult, mult // 2,
+                                                                                    kernel_size=r_kernel * 2, stride=r,
+                                                                                    padding=r_kernel - r // 2,
+                                                                                    output_padding=r % 2)
+                                                                 ))
+
+    def forward(self, x):
+        x = torch.sin(x) + x
+        out1 = self.upsample(x)
+        out2 = self.trans_upsample(x)
+        return out1 + out2
+
+
+class Downsample(nn.Module):
+    def __init__(self, mult, r):
+        super(Downsample, self).__init__()
+        self.r = r
+        r_kernel = r if r >= 5 else 5
+        self.trans_downsample = nn.Sequential(nn.LeakyReLU(0.2),
+                                              nn.utils.weight_norm(nn.Conv1d(mult, mult * 2,
+                                                                             kernel_size=r_kernel * 2, stride=r,
+                                                                             padding=r_kernel - r // 2)
+                                                                   ))
+
+    def forward(self, x):
+        out = self.trans_downsample(x)
+        return out
+
+
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find("BatchNorm2d") != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+def weights_zero_init(m):
+    classname = m.__class__.__name__
+    if classname.find("Conv") != -1:
+        m.weight.data.fill_(0.0)
+        m.bias.data.fill_(0.0)
+
+
+def WNConv1d(*args, **kwargs):
+    return weight_norm(nn.Conv1d(*args, **kwargs))
+
+
+def WNConvTranspose1d(*args, **kwargs):
+    return weight_norm(nn.ConvTranspose1d(*args, **kwargs))
+
+
+class Audio2Mel(nn.Module):
+    def __init__(
+            self,
+            hop_length=300,
+            sampling_rate=24000,
+            n_mel_channels=80,
+            mel_fmin=0.,
+            mel_fmax=None,
+            frame_size=0.05,
+            device='cpu'
+    ):
+        super().__init__()
+        ##############################################
+        # FFT Parameters                              #
+        ##############################################
+
+        self.n_fft = int(np.power(2., np.ceil(np.log(sampling_rate * frame_size) / np.log(2))))
+        window = torch.hann_window(int(sampling_rate * frame_size)).float()
+        mel_basis = librosa_mel_fn(
+            sampling_rate, self.n_fft, n_mel_channels, mel_fmin, mel_fmax
+        )  # Mel filter (by librosa)
+        mel_basis = torch.from_numpy(mel_basis).float()
+        self.register_buffer("mel_basis", mel_basis)
+        self.register_buffer("window", window)
+
+        self.hop_length = hop_length
+        self.win_length = int(sampling_rate * frame_size)
+        self.sampling_rate = sampling_rate
+        self.n_mel_channels = n_mel_channels
+
+    def forward(self, audio):
+        fft = torch.stft(
+            audio.squeeze(1),
+            n_fft=self.n_fft,
+            hop_length=self.hop_length,
+            win_length=self.win_length,
+            window=self.window,
+            center=True,
+        )
+        real_part, imag_part = fft.unbind(-1)
+        magnitude = torch.sqrt(torch.clamp(real_part ** 2 + imag_part ** 2, min=1e-5))
+        mel_output = torch.matmul(self.mel_basis, magnitude)
+
+        log_mel_spec = 20 * torch.log10(torch.clamp(mel_output, min=1e-5)) - 20
+        norm_mel = (log_mel_spec + 115.) / 115.
+        mel_comp = torch.clamp(norm_mel * 8. - 4., -4., 4.)
+
+        return mel_comp
+
+
+class ResnetBlock(nn.Module):
+    def __init__(self, dim, dilation=1, dim_in=None):
+        super().__init__()
+        if dim_in is None:
+            dim_in = dim
+
+        self.block = nn.Sequential(
+            nn.LeakyReLU(0.2),
+            nn.ReflectionPad1d(dilation),
+            WNConv1d(dim_in, dim, kernel_size=3, dilation=dilation),
+            nn.LeakyReLU(0.2),
+            WNConv1d(dim, dim, kernel_size=1),
+        )
+        self.shortcut = WNConv1d(dim_in, dim, kernel_size=1)
+
+    def forward(self, x):
+        return self.shortcut(x) + self.block(x)
+
+
+'''
+参照hifigan（https://arxiv.org/pdf/2010.05646.pdf）v2结构
+多尺度主要是kernel_size不同，3组并行卷积模块，每个卷积模块内部采用不同的串行dilation size，且中间交叉正常无dilation卷积层
+'''
+
+
+class ResBlockMRFV2(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3, dilation=(1, 3, 5)):
+        super(ResBlockMRFV2, self).__init__()
+        self.convs1 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[0],
+                               padding=get_padding(kernel_size, dilation[0]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[1],
+                               padding=get_padding(kernel_size, dilation[1]))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=dilation[2],
+                               padding=get_padding(kernel_size, dilation[2])))
+        ])
+        self.convs1.apply(init_weights)
+
+        self.convs2 = nn.ModuleList([
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1))),
+            weight_norm(Conv1d(channels, channels, kernel_size, 1, dilation=1,
+                               padding=get_padding(kernel_size, 1)))
+        ])
+        self.convs2.apply(init_weights)
+
+    def forward(self, x):
+        for c1, c2 in zip(self.convs1, self.convs2):
+            xt = F.leaky_relu(x, 0.2)
+            xt = c1(xt)
+            xt = F.leaky_relu(xt, 0.2)
+            xt = c2(xt)
+            x = xt + x
+        return x
+
+    def remove_weight_norm(self):
+        for l in self.convs1:
+            remove_weight_norm(l)
+        for l in self.convs2:
+            remove_weight_norm(l)
+
+
+class ResBlockMRFV2Inter(torch.nn.Module):
+    def __init__(self, channels, kernel_size=3):
+        super(ResBlockMRFV2Inter, self).__init__()
+        self.block1 = ResBlockMRFV2(channels)
+        self.block2 = ResBlockMRFV2(channels, 7)
+        self.block3 = ResBlockMRFV2(channels, 11)
+
+    def forward(self, x):
+        xs = self.block1(x)
+        xs += self.block2(x)
+        xs += self.block3(x)
+        x = xs / 3
+        return x
+
+
+class Generator(nn.Module):
+    def __init__(self, input_size_, ngf, n_residual_layers, num_band, args, ratios=[5, 5, 4, 3], onnx_export=False,
+                 device='cpu'):
+        super().__init__()
+        self.hop_length = args.frame_shift
+        self.args = args
+        self.onnx_export = onnx_export
+
+        # ------------- Define upsample layers ----------------
+        mult = int(2 ** len(ratios))
+        model_up = []
+        input_size = input_size_
+        model_up += [
+            nn.ReflectionPad1d(3),
+            WNConv1d(input_size, mult * ngf, kernel_size=7, padding=0),
+        ]
+
+        # Upsample to raw audio scale
+        for i, r in enumerate(ratios):
+            model_up += [Upsample(mult * ngf, r)]
+            model_up += [ResBlockMRFV2Inter(mult * ngf // 2)]
+            mult //= 2
+
+        model_up += [
+            nn.LeakyReLU(0.2),
+            nn.ReflectionPad1d(3),
+            WNConv1d(ngf, num_band, kernel_size=7, padding=0),
+            nn.Tanh(),
+        ]
+        if not args.use_tanh:
+            model_up[-1] = nn.Conv1d(num_band, num_band, 1)
+        model_up[-2].apply(weights_zero_init)
+
+        self.model_up = nn.Sequential(*model_up)
+
+        self.apply(weights_init)
+
+    def forward(self, mel, step=None):
+        # mel input: (batch_size, seq_num, 80)
+        if self.onnx_export:
+            mel = mel.transpose(1, 2)
+            # on onnx, for engineering, mel input: (batch_size, 80, seq_num)
+
+        # Between Down and up
+        x = mel
+
+        # Upsample pipline
+        cnt_after_upsample = 0
+
+        for i, m in enumerate(self.model_up):
+            x = m(x)
+
+            if type(m) == Upsample:
+                cnt_after_upsample += 1
+
+        return x

xinference/thirdparty/megatts3/tts/modules/wavvae/decoder/seanet_encoder.py

@@ -0,0 +1,38 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 torch import nn
+from tts.modules.wavvae.encoder.common_modules.seanet import SEANetEncoder
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        dowmsamples: List[int] = [6, 5, 5, 4, 2],
+    ):
+        super().__init__()
+
+        # breakpoint()
+        self.frame_rate = 25  # not use
+        self.encoder = SEANetEncoder(causal=False, n_residual_layers=1, norm='weight_norm', pad_mode='reflect', lstm=2,
+                                dimension=512, channels=1, n_filters=32, ratios=dowmsamples, activation='ELU',
+                                kernel_size=7, residual_kernel_size=3, last_kernel_size=7, dilation_base=2,
+                                true_skip=False, compress=2)
+
+    def forward(self, audio: torch.Tensor):
+        audio = audio.unsqueeze(1)                  # audio(16,24000)
+        emb = self.encoder(audio)
+        return emb

xinference/thirdparty/megatts3/tts/modules/wavvae/decoder/wavvae_v3.py

@@ -0,0 +1,60 @@
+# Copyright 2025 ByteDance and/or its affiliates.
+#
+# 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 argparse
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from tts.modules.wavvae.decoder.seanet_encoder import Encoder
+from tts.modules.wavvae.decoder.diag_gaussian import DiagonalGaussianDistribution
+from tts.modules.wavvae.decoder.hifigan_modules import Generator, Upsample
+
+
+class WavVAE_V3(nn.Module):
+    def __init__(self, hparams=None):
+        super().__init__()
+        self.encoder = Encoder(dowmsamples=[6, 5, 4, 4, 2])
+        self.proj_to_z = nn.Linear(512, 64)
+        self.proj_to_decoder = nn.Linear(32, 320)
+
+        config_path = hparams['melgan_config']
+        args = argparse.Namespace()
+        args.__dict__.update(config_path)
+        self.latent_upsampler = Upsample(320, 4)
+        self.decoder = Generator(
+            input_size_=160, ngf=128, n_residual_layers=4,
+            num_band=1, args=args, ratios=[5,4,4,3])
+
+    ''' encode waveform into 25 hz latent representation '''
+    def encode_latent(self, audio):
+        posterior = self.encode(audio)
+        latent = posterior.sample().permute(0, 2, 1)  # (b,t,latent_channel)
+        return latent
+
+    def encode(self, audio):
+        x = self.encoder(audio).permute(0, 2, 1)
+        x = self.proj_to_z(x).permute(0, 2, 1)
+        poseterior = DiagonalGaussianDistribution(x)
+        return poseterior
+
+    def decode(self, latent):
+        latent = self.proj_to_decoder(latent).permute(0, 2, 1)
+        return self.decoder(self.latent_upsampler(latent))
+
+    def forward(self, audio):
+        posterior = self.encode(audio)
+        latent = posterior.sample().permute(0, 2, 1)  # (b, t, latent_channel)
+        recon_wav = self.decode(latent)
+        return recon_wav, posterior