Update src/sbv2/synthesizer_trn.py

src/sbv2/synthesizer_trn.py

@@ -1,109 +1,70 @@
-# 完全新規構成 synthesize_trn.py
-# (config.jsonにstrict対応)
-
-import torch
-import torch.nn as nn
-
-from src.sbv2.generator import Generator
-from src.sbv2.posterior_encoder import PosteriorEncoder
-from src.sbv2.flow import ResidualCouplingBlock
-from src.sbv2.flow import Flip
-from src.sbv2.duration_predictor import DurationPredictor
-from src.sbv2.stochastic_duration_predictor import StochasticDurationPredictor
-
-class SynthesizerTrn(nn.Module):
-    def __init__(self,
-                 n_vocab,
-                 spec_channels,
-                 inter_channels,
-                 hidden_channels,
-                 filter_channels,
-                 n_heads,
-                 n_layers,
-                 kernel_size,
-                 p_dropout,
-                 resblock,
-                 resblock_kernel_sizes,
-                 resblock_dilation_sizes,
-                 upsample_rates,
-                 upsample_initial_channel,
-                 upsample_kernel_sizes,
-                 segment_size,
-                 gin_channels,
-                 out_channels,
-                 dec_kernel_size,
-                 enc_channels,
-                 enc_out_channels,
-                 enc_kernel_size,
-                 enc_dilation_rate,
-                 enc_n_layers,
-                 flow_hidden_channels,
-                 flow_kernel_size,
-                 flow_n_layers,
-                 flow_n_flows,
-                 sdp_hidden_channels,
-                 sdp_kernel_size,
-                 sdp_n_layers,
-                 sdp_dropout,
-                 sampling_rate,
-                 filter_length,
-                 hop_length,
-                 win_length):
-        super().__init__()
-        
-        self.n_vocab = n_vocab
-        self.spec_channels = spec_channels
-        self.inter_channels = inter_channels
-        self.hidden_channels = hidden_channels
-        self.filter_channels = filter_channels
-        self.n_heads = n_heads
-        self.n_layers = n_layers
-        self.kernel_size = kernel_size
-        self.p_dropout = p_dropout
-        self.resblock = resblock
-        self.resblock_kernel_sizes = resblock_kernel_sizes
-        self.resblock_dilation_sizes = resblock_dilation_sizes
-        self.upsample_rates = upsample_rates
-        self.upsample_initial_channel = upsample_initial_channel
-        self.upsample_kernel_sizes = upsample_kernel_sizes
-        self.segment_size = segment_size
-        self.gin_channels = gin_channels
-        self.out_channels = out_channels
-        self.dec_kernel_size = dec_kernel_size
-        self.enc_channels = enc_channels
-        self.enc_out_channels = enc_out_channels
-        self.enc_kernel_size = enc_kernel_size
-        self.enc_dilation_rate = enc_dilation_rate
-        self.enc_n_layers = enc_n_layers
-        self.flow_hidden_channels = flow_hidden_channels
-        self.flow_kernel_size = flow_kernel_size
-        self.flow_n_layers = flow_n_layers
-        self.flow_n_flows = flow_n_flows
-        self.sdp_hidden_channels = sdp_hidden_channels
-        self.sdp_kernel_size = sdp_kernel_size
-        self.sdp_n_layers = sdp_n_layers
-        self.sdp_dropout = sdp_dropout
-        self.sampling_rate = sampling_rate
-        self.filter_length = filter_length
-        self.hop_length = hop_length
-        self.win_length = win_length
-
-        # ネットワークモジュール
-        self.enc_p = PosteriorEncoder(
-            spec_channels, inter_channels, hidden_channels,
-            kernel_size, enc_dilation_rate, int(enc_n_layers))
-        self.decoder = Generator(
-            upsample_rates, upsample_initial_channel)
-        self.flow = ResidualCouplingBlock(
-            inter_channels, flow_hidden_channels, flow_kernel_size, flow_n_layers)
-        self.flow_post = Flip()
-        self.dp = DurationPredictor(
-            inter_channels, filter_channels, kernel_size, p_dropout)
-        self.sdp = StochasticDurationPredictor(
-            inter_channels, filter_channels, kernel_size, p_dropout)
-
-    def forward(self, *args, **kwargs):
-        raise NotImplementedError("Training用 forwardは未実装です")
-
-    def infer(self, *args, **kwargs):
-        raise NotImplementedError("推論用 inferは未実装です")
+import torch.nn as nn
+from src.sbv2.generator import Generator
+from src.sbv2.posterior_encoder import PosteriorEncoder
+
+class SynthesizerTrn(nn.Module):
+    def __init__(
+        self,
+        n_vocab,
+        p_dropout,
+        segment_size,
+        inter_channels,
+        out_channels,
+        hidden_channels,
+        filter_channels,
+        dec_kernel_size,
+        enc_channels,
+        enc_out_channels,
+        enc_kernel_size,
+        enc_dilation_rate,
+        enc_n_layers,
+        flow_hidden_channels,
+        flow_kernel_size,
+        flow_n_layers,
+        flow_n_flows,
+        sdp_hidden_channels,
+        sdp_kernel_size,
+        sdp_n_layers,
+        sdp_dropout,
+        sampling_rate,
+        filter_length,
+        hop_length,
+        win_length,
+        resblock,
+        resblock_kernel_sizes,
+        resblock_dilation_sizes,
+        upsample_rates,
+        upsample_initial_channel,
+        upsample_kernel_sizes,
+        gin_channels
+    ):
+        super().__init__()
+
+        # PosteriorEncoder
+        self.enc_p = PosteriorEncoder(
+            in_channels=enc_channels,
+            out_channels=enc_out_channels,
+            hidden_channels=hidden_channels,
+            kernel_size=enc_kernel_size,
+            dilation_rate=enc_dilation_rate,
+            n_layers=enc_n_layers
+        )
+
+        # Generator (Decoder)
+        self.decoder = Generator(
+            upsample_rates=upsample_rates,
+            upsample_initial_channel=upsample_initial_channel,
+            resblock_kernel_sizes=resblock_kernel_sizes,
+            resblock_dilation_sizes=resblock_dilation_sizes,
+            resblock=resblock,
+            upsample_kernel_sizes=upsample_kernel_sizes,
+            inter_channels=inter_channels,
+            out_channels=out_channels,
+            sampling_rate=sampling_rate
+        )
+
+        # その他層（ここでは省略）
+
+    def infer(self, x, noise_scale=0.667, noise_scale_w=0.8, length_scale=1.0):
+        # 仮の推論実装（必要に応じて調整）
+        return self.decoder(x), None