chore(format): run black on main (#510)

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
2noise · Jul 2, 2024 · a13daa4 · a13daa4
1 parent 46200b3
commit a13daa4
Show file tree

Hide file tree

Showing 2 changed files with 12 additions and 4 deletions.
diff --git a/ChatTTS/core.py b/ChatTTS/core.py
@@ -295,7 +295,9 @@ def _load(
 
         if gpt_config_path:
             cfg = OmegaConf.load(gpt_config_path)
-            gpt = GPT(**cfg, use_flash_attn=use_flash_attn, device=device, logger=self.logger).eval()
+            gpt = GPT(
+                **cfg, use_flash_attn=use_flash_attn, device=device, logger=self.logger
+            ).eval()
             assert gpt_ckpt_path, "gpt_ckpt_path should not be None"
             gpt.load_state_dict(torch.load(gpt_ckpt_path, weights_only=True, mmap=True))
             gpt.prepare(compile=compile and "cuda" in str(device))

diff --git a/ChatTTS/model/gpt.py b/ChatTTS/model/gpt.py
@@ -99,7 +99,9 @@ def get(self) -> bool:
             return self._interrupt
 
     def _build_llama(
-        self, config: omegaconf.DictConfig, device: torch.device,
+        self,
+        config: omegaconf.DictConfig,
+        device: torch.device,
     ) -> LlamaModel:
 
         model = None
@@ -122,7 +124,9 @@ def _build_llama(
                     **config,
                     attn_implementation="flash_attention_2",
                 )
-                self.logger.warn("enabling flash_attention_2 may make gpt be even slower")
+                self.logger.warn(
+                    "enabling flash_attention_2 may make gpt be even slower"
+                )
             else:
                 llama_config = LlamaConfig(**config)
             model = LlamaModel(llama_config)
@@ -439,7 +443,9 @@ def generate(
                 )
                 del_all(model_input)
                 attentions.append(outputs.attentions)
-                hidden_states = outputs.last_hidden_state.to(self.device, dtype=torch.float)  # 🐻
+                hidden_states = outputs.last_hidden_state.to(
+                    self.device, dtype=torch.float
+                )  # 🐻
                 past_key_values = outputs.past_key_values
                 del_all(outputs)
                 if return_hidden: