feat: expose RoPE theta via config API

configs/acoustic.yaml

@@ -67,6 +67,7 @@ max_beta: 0.02
 enc_ffn_kernel_size: 3
 use_rope: true
 rope_interleaved: false
+rope_theta: 10000
 use_stretch_embed: true
 use_variance_scaling: true
 rel_pos: true

configs/templates/config_acoustic.yaml

@@ -75,6 +75,7 @@ diffusion_type: reflow
 enc_ffn_kernel_size: 3
 use_rope: true
 rope_interleaved: false
+rope_theta: 10000
 use_stretch_embed: true
 use_variance_scaling: true
 use_shallow_diffusion: true

configs/templates/config_variance.yaml

@@ -66,6 +66,7 @@ tension_logit_max: 10.0
 enc_ffn_kernel_size: 3
 use_rope: true
 rope_interleaved: false
+rope_theta: 10000
 use_stretch_embed: false
 use_variance_scaling: true
 hidden_size: 384

configs/variance.yaml

@@ -37,6 +37,7 @@ predict_tension: false
 enc_ffn_kernel_size: 3
 use_rope: true
 rope_interleaved: false
+rope_theta: 10000
 use_stretch_embed: false
 use_variance_scaling: true
 rel_pos: true

modules/fastspeech/acoustic_encoder.py

@@ -58,6 +58,7 @@ def __init__(self, vocab_size):
             dropout=hparams['dropout'], num_heads=hparams['num_heads'],
             use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False),
             use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
+            rope_theta=hparams.get('rope_theta', 10000),
             mix_ln_layer=self.mix_ln_layer
         )
 

modules/fastspeech/tts_modules.py

@@ -373,7 +373,7 @@ def __init__(
             self, hidden_size, num_layers,
             ffn_kernel_size=9, ffn_act='gelu',
             dropout=None, num_heads=2, use_pos_embed=True, rel_pos=True,
-            use_rope=False, rope_interleaved=True, mix_ln_layer=None
+            use_rope=False, rope_interleaved=True, rope_theta=10000, mix_ln_layer=None
     ):
         super().__init__()
         self.num_layers = num_layers
@@ -386,7 +386,9 @@ def __init__(
                     "RoPE requires the hidden size to be multiple of "
                     f"num_heads * 2 = {num_heads * 2}, but got {embed_dim}."
                 )
-            rotary_embed = RotaryEmbedding(dim=embed_dim // num_heads, interleaved=rope_interleaved)
+            rotary_embed = RotaryEmbedding(
+                dim=embed_dim // num_heads, theta=rope_theta, interleaved=rope_interleaved
+            )
         else:
             rotary_embed = None
         self.layers = nn.ModuleList([

modules/fastspeech/variance_encoder.py

@@ -34,7 +34,8 @@ def __init__(self, vocab_size):
             ffn_kernel_size=hparams['enc_ffn_kernel_size'], ffn_act=hparams['ffn_act'],
             dropout=hparams['dropout'], num_heads=hparams['num_heads'],
             use_pos_embed=hparams['use_pos_embed'], rel_pos=hparams.get('rel_pos', False), 
-            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True)
+            use_rope=hparams.get('use_rope', False), rope_interleaved=hparams.get('rope_interleaved', True),
+            rope_theta=hparams.get('rope_theta', 10000)
         )
 
         dur_hparams = hparams['dur_prediction_args']
@@ -128,7 +129,8 @@ def get_hparam(key):
             ffn_kernel_size=get_hparam('enc_ffn_kernel_size'), ffn_act=get_hparam('ffn_act'),
             dropout=get_hparam('dropout'), num_heads=get_hparam('num_heads'),
             use_pos_embed=get_hparam('use_pos_embed'), rel_pos=get_hparam('rel_pos'),
-            use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True)
+            use_rope=get_hparam('use_rope'), rope_interleaved=hparams.get('rope_interleaved', True),
+            rope_theta=hparams.get('rope_theta', 10000)
         )
         self.out_proj = Linear(hidden_size, hparams['hidden_size'])