Update modeling_bert_vits2.py

modeling_bert_vits2.py

@@ -33,16 +33,10 @@ from transformers.modeling_outputs import (
 from transformers.models.bert.modeling_bert import BertModel
 from transformers.modeling_utils import PreTrainedModel
 from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings
-from configuration_bert_vits2 import BertVits2Config
-
 
 logger = logging.get_logger(__name__)
 
 
-# General docstring
-_CONFIG_FOR_DOC = "BertVits2Config"
-
-
 @dataclass
 class BertVits2ModelOutput(ModelOutput):
     """
@@ -328,7 +322,7 @@ def _rational_quadratic_spline(
 
 
 class BertVits2WaveNet(torch.nn.Module):
-    def __init__(self, config: BertVits2Config, num_layers: int):
+    def __init__(self, config, num_layers: int):
         super().__init__()
         self.hidden_size = config.hidden_size
         self.num_layers = num_layers
@@ -408,7 +402,7 @@ class BertVits2WaveNet(torch.nn.Module):
 
 
 class BertVits2PosteriorEncoder(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.out_channels = config.flow_size
 
@@ -485,7 +479,7 @@ class HifiGanResidualBlock(nn.Module):
 
 
 class BertVits2HifiGan(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.num_kernels = len(config.resblock_kernel_sizes)
@@ -571,7 +565,7 @@ class BertVits2HifiGan(nn.Module):
 
 
 class BertVits2ResidualCouplingLayer(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.half_channels = config.flow_size // 2
 
@@ -593,7 +587,7 @@ class BertVits2ResidualCouplingLayer(nn.Module):
 
 
 class BertVits2ResidualCouplingBlock(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.flows = nn.ModuleList()
         for _ in range(config.prior_encoder_num_flows):
@@ -608,7 +602,7 @@ class BertVits2ResidualCouplingBlock(nn.Module):
 
 
 class BertVits2TransformerCouplingLayer(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.half_channels = config.flow_size // 2
 
@@ -653,7 +647,7 @@ class BertVits2TransformerCouplingLayer(nn.Module):
 
 
 class BertVits2TransformerCouplingBlock(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.flows = nn.ModuleList([
             BertVits2TransformerCouplingLayer(config) for _ in range(config.prior_encoder_num_flows)
@@ -672,7 +666,7 @@ class BertVits2TransformerCouplingBlock(nn.Module):
 
 
 class BertVits2DilatedDepthSeparableConv(nn.Module):
-    def __init__(self, config: BertVits2Config, dropout_rate=0.0):
+    def __init__(self, config, dropout_rate=0.0):
         super().__init__()
         kernel_size = config.duration_predictor_kernel_size
         channels = config.hidden_size
@@ -718,7 +712,7 @@ class BertVits2DilatedDepthSeparableConv(nn.Module):
 
 
 class BertVits2ConvFlow(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.filter_channels = config.hidden_size
         self.half_channels = config.depth_separable_channels // 2
@@ -761,7 +755,7 @@ class BertVits2ConvFlow(nn.Module):
 
 
 class BertVits2ElementwiseAffine(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.channels = config.depth_separable_channels
         self.translate = nn.Parameter(torch.zeros(self.channels, 1))
@@ -918,7 +912,7 @@ class BertVits2DurationPredictor(nn.Module):
 class BertVits2Attention(nn.Module):
     """Multi-headed attention with relative positional representation."""
 
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.embed_dim = config.hidden_size
         self.num_heads = config.num_attention_heads
@@ -1130,7 +1124,7 @@ class BertVits2FeedForward(nn.Module):
 
 
 class BertVits2EncoderLayer(nn.Module):
-    def __init__(self, config: BertVits2Config, kernel_size=None):
+    def __init__(self, config, kernel_size=None):
         super().__init__()
         self.attention = BertVits2Attention(config)
         self.dropout = nn.Dropout(config.hidden_dropout)
@@ -1169,7 +1163,7 @@ class BertVits2EncoderLayer(nn.Module):
 
 
 class BertVits2Encoder(nn.Module):
-    def __init__(self, config: BertVits2Config, kernel_size=None, n_layers=None):
+    def __init__(self, config, kernel_size=None, n_layers=None):
         super().__init__()
         self.config = config
         if n_layers is None:
@@ -1260,7 +1254,7 @@ class BertVits2TextEncoder(nn.Module):
     Transformer encoder that uses relative positional representation instead of absolute positional encoding.
     """
 
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, config.pad_token_id)
@@ -1330,7 +1324,7 @@ class BertVits2TextEncoder(nn.Module):
 
 
 class BertVits2ReferenceEncoder(nn.Module):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__()
         self.config = config
         ref_enc_filters = [32, 32, 64, 64, 128, 128]
@@ -1464,7 +1458,7 @@ BERT_VITS2_INPUTS_DOCSTRING = r"""
     BERT_VITS2_START_DOCSTRING,
 )
 class BertVits2Model(BertVits2PreTrainedModel):
-    def __init__(self, config: BertVits2Config):
+    def __init__(self, config):
         super().__init__(config)
         self.config = config
         self.text_encoder = BertVits2TextEncoder(config)