Restore stable text style conditioning

Author: mera-chan[bot]

src/python/piper_train/vits/lightning.py

@@ -144,12 +144,11 @@ def __init__(
         self.automatic_optimization = (
             False  # Multiple optimizers require manual optimization
         )
-        if style_condition_mode not in (None, "global"):
+        if style_condition_mode not in {"global", "text"}:
             raise ValueError(
-                "Only global style conditioning is supported; "
+                "style_condition_mode must be either 'global' or 'text', "
                 f"got {style_condition_mode!r}"
             )
-        style_condition_mode = "global"
 
         # Fix gin_channels BEFORE save_hyperparameters() so the correct value is saved
         # This fixes the bug where gin_channels=0 was saved for multi-speaker models
@@ -185,6 +184,7 @@ def __init__(
             prosody_dim=self.hparams.prosody_dim,
             style_vector_dim=self.hparams.style_vector_dim,
             style_condition_dropout=self.hparams.style_condition_dropout,
+            style_condition_mode=self.hparams.style_condition_mode,
         )
         self.model_d = MultiPeriodDiscriminator(
             use_spectral_norm=self.hparams.use_spectral_norm
@@ -1130,11 +1130,12 @@ def add_model_specific_args(parent_parser):
         )
         parser.add_argument(
             "--style-condition-mode",
-            choices=("global",),
+            choices=("text", "global"),
             default="global",
             help=(
-                "Style-vector injection mode. Only 'global' is supported; "
-                "style vectors are added to VITS global conditioning."
+                "Where to inject utterance-level style vectors. "
+                "'global' adds style to VITS global conditioning; 'text' adds "
+                "projected style to the scaled text encoder input."
             ),
         )
         parser.add_argument(

src/python/piper_train/vits/models.py

@@ -194,6 +194,8 @@ def __init__(
         kernel_size: int,
         p_dropout: float,
         gin_channels: int = 0,
+        style_vector_dim: int = 0,
+        style_condition_dropout: float = 0.0,
     ):
         super().__init__()
         self.n_vocab = n_vocab
@@ -205,9 +207,19 @@ def __init__(
         self.kernel_size = kernel_size
         self.p_dropout = p_dropout
         self.gin_channels = gin_channels
+        self.style_vector_dim = style_vector_dim
+        self.style_condition_dropout = style_condition_dropout
 
         self.emb = nn.Embedding(n_vocab, hidden_channels)
         nn.init.normal_(self.emb.weight, 0.0, hidden_channels**-0.5)
+        self.style_proj = (
+            nn.Linear(style_vector_dim, hidden_channels)
+            if style_vector_dim > 0
+            else None
+        )
+        if self.style_proj is not None:
+            nn.init.zeros_(self.style_proj.weight)
+            nn.init.zeros_(self.style_proj.bias)
 
         self.encoder = attentions.Encoder(
             hidden_channels, filter_channels, n_heads, n_layers, kernel_size, p_dropout
@@ -217,8 +229,41 @@ def __init__(
         if gin_channels != 0:
             self.cond_layer = nn.Conv1d(gin_channels, hidden_channels, 1)
 
-    def forward(self, x, x_lengths, g=None):
+    def _style_embedding(self, style_vector, batch_size: int, device, dtype):
+        if self.style_proj is None:
+            return None
+
+        if style_vector is None:
+            style_vector = torch.zeros(
+                batch_size,
+                self.style_vector_dim,
+                device=device,
+                dtype=dtype,
+            )
+
+        proj_weight = next(self.style_proj.parameters())
+        style_emb = self.style_proj(
+            style_vector.to(device=proj_weight.device, dtype=proj_weight.dtype)
+        )
+        if self.training and self.style_condition_dropout > 0.0:
+            keep = (
+                torch.rand(style_emb.size(0), device=style_emb.device)
+                >= self.style_condition_dropout
+            ).to(style_emb.dtype)
+            style_emb = style_emb * keep[:, None]
+
+        return style_emb.to(device=device, dtype=dtype).unsqueeze(1)
+
+    def forward(self, x, x_lengths, g=None, style_vector=None):
         x = self.emb(x) * math.sqrt(self.hidden_channels)  # [b, t, h]
+        style_emb = self._style_embedding(
+            style_vector, batch_size=x.size(0), device=x.device, dtype=x.dtype
+        )
+        if style_emb is not None:
+            # Add text-level style after token embedding scaling. PE-A emotion
+            # vectors can be intentionally amplified at inference, and scaling
+            # the projected style by sqrt(hidden_channels) destabilizes duration.
+            x = x + style_emb
         x = torch.transpose(x, 1, -1)  # [b, h, t]
         x_mask = torch.unsqueeze(
             commons.sequence_mask(x_lengths, x.size(2)), 1
@@ -776,8 +821,14 @@ def __init__(
         prosody_language_ids: "set[int] | None" = None,
         style_vector_dim: int = 0,
         style_condition_dropout: float = 0.0,
+        style_condition_mode: str = "global",
     ):
         super().__init__()
+        if style_condition_mode not in {"global", "text"}:
+            raise ValueError(
+                "style_condition_mode must be either 'global' or 'text', "
+                f"got {style_condition_mode!r}"
+            )
         self.n_vocab = n_vocab
         self.spec_channels = spec_channels
         self.inter_channels = inter_channels
@@ -800,6 +851,7 @@ def __init__(
         self.prosody_dim = prosody_dim
         self.style_vector_dim = style_vector_dim
         self.style_condition_dropout = style_condition_dropout
+        self.style_condition_mode = style_condition_mode
         # Language IDs with real prosody features (others are zeroed).
         # Default: {0} (JA only). Configurable via prosody_language_ids param.
         self.prosody_language_ids: set[int] = (
@@ -818,6 +870,10 @@ def __init__(
             kernel_size,
             p_dropout,
             gin_channels=gin_channels,
+            style_vector_dim=(
+                style_vector_dim if style_condition_mode == "text" else 0
+            ),
+            style_condition_dropout=style_condition_dropout,
         )
         self.dec = Generator(
             inter_channels,
@@ -872,7 +928,7 @@ def __init__(
         self.spk_proj = None
 
         self.style_proj = None
-        if style_vector_dim > 0:
+        if style_vector_dim > 0 and style_condition_mode == "global":
             if gin_channels <= 0:
                 raise ValueError(
                     "style_vector_dim > 0 requires gin_channels > 0 so style vectors "
@@ -1008,10 +1064,16 @@ def forward(
         # training (emb_g(sid) is always used).  The parameter is reserved for
         # future extensions such as joint speaker-encoder fine-tuning.
         g = self._get_global_conditioning(sid, lid)
-        g = self._add_style_condition(
-            g, style_vector, batch_size=x.size(0), device=x.device, dtype=x.dtype
+        if self.style_condition_mode == "global":
+            g = self._add_style_condition(
+                g, style_vector, batch_size=x.size(0), device=x.device, dtype=x.dtype
+            )
+            text_style_vector = None
+        else:
+            text_style_vector = style_vector
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, g=g, style_vector=text_style_vector
         )
-        x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, g=g)
 
         z, m_q, logs_q, y_mask = self.enc_q(y, y_lengths, g=g)
         z_p = self.flow(z, y_mask, g=g)
@@ -1157,10 +1219,16 @@ def infer(
             if self.n_speakers > 1:
                 assert sid is not None, "Missing speaker id"
             g = self._get_global_conditioning(sid, lid)
-        g = self._add_style_condition(
-            g, style_vector, batch_size=x.size(0), device=x.device, dtype=x.dtype
+        if self.style_condition_mode == "global":
+            g = self._add_style_condition(
+                g, style_vector, batch_size=x.size(0), device=x.device, dtype=x.dtype
+            )
+            text_style_vector = None
+        else:
+            text_style_vector = style_vector
+        x, m_p, logs_p, x_mask = self.enc_p(
+            x, x_lengths, g=g, style_vector=text_style_vector
         )
-        x, m_p, logs_p, x_mask = self.enc_p(x, x_lengths, g=g)
 
         # Prepare input for duration predictor with prosody features
         x_dp = self._prepare_prosody_input(x, x_mask, prosody_features, lid=lid)