Apply isort and black reformatting

Signed-off-by: Edresson <Edresson@users.noreply.github.com>

Author: Edresson

nemo/collections/speechlm2/models/duplex_ear_tts.py

@@ -321,7 +321,13 @@ def setup_rvq_audio_codec(model):
     with fp32_precision():
         if model.cfg.get("pretrained_ae_dir", None):
             model.audio_codec = (
-                RVQVAEModel.from_pretrained(model.cfg.pretrained_ae_dir, cfg=DictConfig(model.cfg.codec_config) if model.cfg.get("codec_config", None) else None, strict=False).eval().to(model.device)
+                RVQVAEModel.from_pretrained(
+                    model.cfg.pretrained_ae_dir,
+                    cfg=DictConfig(model.cfg.codec_config) if model.cfg.get("codec_config", None) else None,
+                    strict=False,
+                )
+                .eval()
+                .to(model.device)
             )
         else:
             # init codec from config

nemo/collections/speechlm2/modules/ear_tts_commons.py

@@ -24,11 +24,11 @@
 from collections.abc import Mapping
 from typing import Any
 
+from omegaconf import DictConfig
 from safetensors import safe_open
 from torch import nn
 
 from nemo.utils import logging
-from omegaconf import DictConfig
 
 # ==============================================================================
 # Contants
@@ -152,6 +152,7 @@ def get_config_from_dir(workdir_path: str) -> DictConfig:
 # Base Model Classes
 # ==============================================================================
 
+
 class PreTrainedModel(nn.Module):
     config_class = DictConfig
 

nemo/collections/speechlm2/modules/rvq_ear_tts_model.py

@@ -20,6 +20,7 @@
 
 import torch
 import transformers
+from omegaconf import DictConfig, OmegaConf
 from torch import Tensor, nn
 from torch.nn import functional as F
 from transformers import AutoConfig, AutoModel, AutoModelForTextEncoding, AutoTokenizer, Cache
@@ -29,7 +30,6 @@
 from nemo.collections.speechlm2.parts.precision import fp32_precision
 from nemo.collections.speechlm2.parts.pretrained import set_model_dict_for_partial_init
 from nemo.utils import logging
-from omegaconf import DictConfig, OmegaConf
 
 # ==============================================================================
 # MLP module and Norm
@@ -894,7 +894,9 @@ def __init__(
 
         # 2. Initialize the backbone model
         if backbone_type:
-            config = AutoConfig.for_model(backbone_type, **(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {}))
+            config = AutoConfig.for_model(
+                backbone_type, **(OmegaConf.to_container(backbone_config, resolve=True) if backbone_config else {})
+            )
             self.backbone = AutoModelForTextEncoding.from_config(config)
         else:
             assert backbone_model_class and backbone_config_class
@@ -1044,12 +1046,12 @@ class RVQEARTTSModel(PreTrainedModel):
     Args:
         config (DictConfig | dict[str, Any]): The configuration object for the model.
     """
+
     rvq_embs: Tensor
 
     def __init__(self, config: DictConfig | dict[str, Any]):
         super().__init__(config)
 
-
         # Backbone module
         if self.config.get("pretrained_text_name", None):
             # Load pretrained backbone from huggingface
@@ -1059,15 +1061,26 @@ def __init__(self, config: DictConfig | dict[str, Any]):
             self.backbone = llm.model  # fetch PretrainedBaseModel from model "ForCausalLM"
         else:
             if self.config.get("backbone_type", None) is None:
-                assert self.config.get("backbone_model_class", None) is not None and self.config.get("backbone_config_class", None) is not None
+                assert (
+                    self.config.get("backbone_model_class", None) is not None
+                    and self.config.get("backbone_config_class", None) is not None
+                )
                 backbone_config = getattr(transformers, self.config.backbone_config_class)(
-                    **(OmegaConf.to_container(self.config.backbone_config, resolve=True) if self.config.backbone_config else {}),
+                    **(
+                        OmegaConf.to_container(self.config.backbone_config, resolve=True)
+                        if self.config.backbone_config
+                        else {}
+                    ),
                 )
                 self.backbone = getattr(transformers, self.config.backbone_model_class)(backbone_config)
             else:
                 backbone_config = AutoConfig.for_model(
                     self.config.backbone_type,
-                    **(OmegaConf.to_container(self.config.backbone_config, resolve=True) if self.config.backbone_config else {}),
+                    **(
+                        OmegaConf.to_container(self.config.backbone_config, resolve=True)
+                        if self.config.backbone_config
+                        else {}
+                    ),
                 )
                 self.backbone = AutoModel.from_config(backbone_config)
 

nemo/collections/speechlm2/modules/rvq_ear_tts_vae.py

@@ -20,13 +20,14 @@
 
 # Third-party
 import torch
+from omegaconf import DictConfig
 from torch import Tensor, nn
 from torch.nn import functional as F
 from torchaudio import functional as ta_F
 
 # Project
 from nemo.collections.speechlm2.modules.ear_tts_commons import PreTrainedModel
-from omegaconf import DictConfig
+
 
 @contextmanager
 def disable_tf32():
@@ -37,6 +38,7 @@ def disable_tf32():
     finally:
         torch.backends.cudnn.allow_tf32 = prev
 
+
 # ==============================================================================
 # Utility Functions
 # ==============================================================================