Support local transformer in longform magpietts (#15296)

Signed-off-by: subhankar-ghosh <[email protected]>

Author: subhankar-ghosh

nemo/collections/tts/models/magpietts.py

@@ -3658,6 +3658,7 @@ def do_tts(
                         end_of_text=end_of_text,
                         beginning_of_text=beginning_of_text,
                         use_cfg=use_cfg,
+                        use_local_transformer_for_inference=True,
                     )
                     if output.predicted_codes_lens[0] > 0:
                         all_codes.append(output.predicted_codes[0, :, : output.predicted_codes_lens[0]])
@@ -4011,7 +4012,7 @@ def _run_longform_forward_with_cfg(
         dummy_additional_decoder_input: Optional[torch.Tensor],
         dummy_addition_dec_mask: Optional[torch.Tensor],
         batch_size: int,
-    ) -> Tuple[torch.Tensor, Any]:
+    ) -> Tuple[torch.Tensor, Any, torch.Tensor]:
         """
         Run forward pass with optional classifier-free guidance.
 
@@ -4029,7 +4030,7 @@ def _run_longform_forward_with_cfg(
             batch_size: Number of items in the batch.
 
         Returns:
-            Tuple of (logits, attention_probs).
+            Tuple of (logits, attention_probs, decoder_output).
         """
         if use_cfg:
             # Combine conditional and unconditional inputs
@@ -4049,7 +4050,7 @@ def _run_longform_forward_with_cfg(
                 )
                 cfg_audio_mask[batch_size:, : dummy_additional_decoder_input.size(1)] = dummy_addition_dec_mask
 
-            combined_logits, attn_probs, _ = self.forward(
+            combined_logits, attn_probs, dec_out = self.forward(
                 dec_input_embedded=cfg_audio_embedded,
                 dec_input_mask=cfg_audio_mask,
                 cond=cfg_cond,
@@ -4061,8 +4062,9 @@ def _run_longform_forward_with_cfg(
             cond_logits = combined_logits[:batch_size]
             uncond_logits = combined_logits[batch_size:]
             all_code_logits = (1 - cfg_scale) * uncond_logits + cfg_scale * cond_logits
+            # NOTE: Keep dec_out doubled for local transformer CFG handling
         else:
-            all_code_logits, attn_probs, _ = self.forward(
+            all_code_logits, attn_probs, dec_out = self.forward(
                 dec_input_embedded=audio_codes_embedded,
                 dec_input_mask=audio_codes_mask,
                 cond=context_tensors.cond,
@@ -4071,7 +4073,7 @@ def _run_longform_forward_with_cfg(
                 multi_encoder_mapping=context_tensors.multi_encoder_mapping,
             )
 
-        return all_code_logits, attn_probs
+        return all_code_logits, attn_probs, dec_out
 
     def _initialize_longform_attn_prior(
         self,
@@ -4242,6 +4244,12 @@ def generate_long_form_speech(
         end_of_text,
         beginning_of_text,
         use_cfg=True,
+        use_local_transformer_for_inference=False,
+        maskgit_n_steps=3,
+        maskgit_noise_scale=0.0,
+        maskgit_fixed_schedule=None,
+        maskgit_dynamic_cfg_scale=False,
+        maskgit_sampling_type=None,
     ):
         """
         Generates speech for long-form text by progressively shifting through text tokens.
@@ -4258,6 +4266,12 @@ def generate_long_form_speech(
             end_of_text (List[bool]): Whether entire text has been provided for each batch item.
             beginning_of_text (bool): Whether this is the first chunk.
             use_cfg (bool): Whether to use classifier-free guidance.
+            use_local_transformer_for_inference (bool): Whether to use local transformer for sampling.
+            maskgit_n_steps (int): Number of MaskGit refinement steps.
+            maskgit_noise_scale (float): Noise scale for MaskGit sampling.
+            maskgit_fixed_schedule (Optional[List[int]]): Fixed schedule for MaskGit.
+            maskgit_dynamic_cfg_scale (bool): Whether to use dynamic CFG scale in MaskGit.
+            maskgit_sampling_type (Optional[str]): Type of MaskGit sampling.
 
         Returns:
             InferBatchOutput: Contains predicted_codes, predicted_codes_lens, and empty audio fields.
@@ -4365,7 +4379,7 @@ def generate_long_form_speech(
                     attn_prior = [attn_prior, None]
 
                 # Run forward pass with optional CFG
-                all_code_logits, attn_probs = self._run_longform_forward_with_cfg(
+                all_code_logits, attn_probs, dec_out = self._run_longform_forward_with_cfg(
                     context_tensors=context_tensors,
                     audio_codes_embedded=_audio_codes_embedded,
                     audio_codes_mask=_audio_codes_mask,
@@ -4439,13 +4453,47 @@ def generate_long_form_speech(
                     unfinished_items = {k: v for k, v in state.unfinished_texts.items() if v}
 
                 all_code_logits_t = all_code_logits[:, -1, :]  # (B, num_codebooks * num_tokens_per_codebook)
-                audio_codes_next = self.sample_codes_from_logits(
-                    all_code_logits_t,
-                    temperature=self.inference_parameters.temperature,
-                    topk=self.inference_parameters.topk,
-                    unfinished_items=unfinished_items,
-                    finished_items=finished_items,
-                )  # (B, num_codebooks)
+
+                if use_local_transformer_for_inference:
+                    if self.local_transformer_type == LocalTransformerType.AR:
+                        # Autoregressive sampling with local transformer
+                        audio_codes_next = self.local_transformer_sample_autoregressive(
+                            dec_output=dec_out[:, -1, :],
+                            temperature=self.inference_parameters.temperature,
+                            topk=self.inference_parameters.topk,
+                            unfinished_items=unfinished_items,
+                            finished_items=finished_items,
+                            use_cfg=use_cfg,
+                            cfg_scale=cfg_scale,
+                            use_kv_cache=self.inference_parameters.use_LT_kv_cache,
+                        )
+                    elif self.local_transformer_type == LocalTransformerType.MASKGIT:
+                        audio_codes_next = self.local_transformer_sample_maskgit(
+                            dec_output=dec_out[:, -1, :],
+                            temperature=self.inference_parameters.temperature,
+                            topk=self.inference_parameters.topk,
+                            unfinished_items=unfinished_items,
+                            finished_items=finished_items,
+                            use_cfg=use_cfg,
+                            cfg_scale=cfg_scale,
+                            n_steps=maskgit_n_steps,
+                            noise_scale=maskgit_noise_scale,
+                            fixed_schedule=maskgit_fixed_schedule,
+                            dynamic_cfg_scale=maskgit_dynamic_cfg_scale,
+                            sampling_type=maskgit_sampling_type,
+                        )
+                    else:
+                        raise ValueError(
+                            f"Local transformer inference requested but local transformer type is {self.local_transformer_type}"
+                        )
+                else:
+                    audio_codes_next = self.sample_codes_from_logits(
+                        all_code_logits_t,
+                        temperature=self.inference_parameters.temperature,
+                        topk=self.inference_parameters.topk,
+                        unfinished_items=unfinished_items,
+                        finished_items=finished_items,
+                    )  # (B, num_codebooks)
                 all_codes_next_argmax = self.sample_codes_from_logits(
                     all_code_logits_t,
                     temperature=self.longform_config.argmax_temperature,

nemo/collections/tts/modules/magpietts_inference/inference.py

@@ -649,6 +649,11 @@ def _run_longform_inference(
                     end_of_text=is_end_of_text,
                     beginning_of_text=beginning_of_text,
                     use_cfg=self.config.use_cfg,
+                    use_local_transformer_for_inference=self.config.use_local_transformer,
+                    maskgit_n_steps=self.config.maskgit_n_steps,
+                    maskgit_noise_scale=self.config.maskgit_noise_scale,
+                    maskgit_fixed_schedule=self.config.maskgit_fixed_schedule,
+                    maskgit_sampling_type=self.config.maskgit_sampling_type,
                 )
 
                 # Unpack output - generate_long_form_speech returns InferBatchOutput