update

.env.webui

@@ -14,9 +14,9 @@ DEBUG_GENERATE=True
 PRELOAD_MODELS=True
 
 # Text-to-Speech (TTS) configuration
-TTS_MAX_LEN=1000
+TTS_MAX_LEN=2000
 SSML_MAX_LEN=3000
 MAX_BATCH_SIZE=12
 
-V_GIT_TAG="🤗hf(0.6.1-rc)"
+V_GIT_TAG="🤗hf(0.6.1)"
 V_GIT_COMMIT=main

README.md

@@ -16,7 +16,7 @@ sdk_version: 4.36.1
 
 | 类型 | 最大字符数 |
 |------|-----------|
-| TTS  | 1000 字符 |
+| TTS  | 2000 字符 |
 | SSML | 3000 字符（不计算 SSML 标签，只计算文本） |
 
 # HuggingFace Space Limit
@@ -25,7 +25,7 @@ Due to the runtime limit for GPU usage on HuggingFace, extremely long tasks will
 
 | Type | Maximum Characters |
 |------|---------------------|
-| TTS  | 1000 characters     |
+| TTS  | 2000 characters     |
 | SSML | 3000 characters (excluding SSML tags, only counting text) |
 
 # 🗣️ ChatTTS-Forge

launch.py

@@ -1,23 +1,24 @@
-import os
 import logging
+import os
 
-from modules.api.api_setup import setup_api_args, setup_model_args, setup_uvicon_args
 from modules.ffmpeg_env import setup_ffmpeg_path
 
-setup_ffmpeg_path()
-logging.basicConfig(
-    level=os.getenv("LOG_LEVEL", "INFO"),
-    format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
-)
+try:
+    setup_ffmpeg_path()
+    logging.basicConfig(
+        level=os.getenv("LOG_LEVEL", "INFO"),
+        format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
+    )
+except BaseException:
+    pass
 
 import argparse
+
 import uvicorn
 
-from modules import config
+from modules.api.api_setup import setup_api_args, setup_model_args, setup_uvicon_args
 from modules.utils import env
 
-from fastapi import FastAPI
-
 logger = logging.getLogger(__name__)
 
 if __name__ == "__main__":

modules/ChatTTS/ChatTTS/__init__.py

@@ -1 +1 @@
-from .core import Chat
+from .core import Chat

modules/ChatTTS/ChatTTS/core.py

@@ -1,21 +1,21 @@
-import os
 import logging
-from omegaconf import OmegaConf
+import os
 
 import torch
+from huggingface_hub import snapshot_download
+from omegaconf import OmegaConf
 from vocos import Vocos
+
+from .infer.api import infer_code, refine_text
 from .model.dvae import DVAE
 from .model.gpt import GPT_warpper
 from .utils.infer_utils import (
-    count_invalid_characters,
-    detect_language,
     apply_character_map,
     apply_half2full_map,
+    count_invalid_characters,
+    detect_language,
 )
 from .utils.io_utils import get_latest_modified_file
-from .infer.api import refine_text, infer_code
-
-from huggingface_hub import snapshot_download
 
 logging.basicConfig(level=logging.INFO)
 

modules/ChatTTS/ChatTTS/infer/api.py

@@ -1,6 +1,7 @@
 import torch
 import torch.nn.functional as F
 from transformers.generation import TopKLogitsWarper, TopPLogitsWarper
+
 from ..utils.infer_utils import CustomRepetitionPenaltyLogitsProcessorRepeat
 
 

modules/ChatTTS/ChatTTS/model/dvae.py

@@ -1,28 +1,36 @@
 import math
-from einops import rearrange
-from vector_quantize_pytorch import GroupedResidualFSQ
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from einops import rearrange
+from vector_quantize_pytorch import GroupedResidualFSQ
+
 
 class ConvNeXtBlock(nn.Module):
     def __init__(
         self,
         dim: int,
         intermediate_dim: int,
-        kernel, dilation,
+        kernel,
+        dilation,
         layer_scale_init_value: float = 1e-6,
     ):
         # ConvNeXt Block copied from Vocos.
         super().__init__()
-        self.dwconv = nn.Conv1d(dim, dim, 
-                                kernel_size=kernel, padding=dilation*(kernel//2), 
-                                dilation=dilation, groups=dim
-                            )  # depthwise conv
-        
+        self.dwconv = nn.Conv1d(
+            dim,
+            dim,
+            kernel_size=kernel,
+            padding=dilation * (kernel // 2),
+            dilation=dilation,
+            groups=dim,
+        )  # depthwise conv
+
         self.norm = nn.LayerNorm(dim, eps=1e-6)
-        self.pwconv1 = nn.Linear(dim, intermediate_dim)  # pointwise/1x1 convs, implemented with linear layers
+        self.pwconv1 = nn.Linear(
+            dim, intermediate_dim
+        )  # pointwise/1x1 convs, implemented with linear layers
         self.act = nn.GELU()
         self.pwconv2 = nn.Linear(intermediate_dim, dim)
         self.gamma = (
@@ -31,7 +39,7 @@ class ConvNeXtBlock(nn.Module):
             else None
         )
 
-    def forward(self, x: torch.Tensor, cond = None) -> torch.Tensor:
+    def forward(self, x: torch.Tensor, cond=None) -> torch.Tensor:
         residual = x
         x = self.dwconv(x)
         x = x.transpose(1, 2)  # (B, C, T) -> (B, T, C)
@@ -45,14 +53,11 @@ class ConvNeXtBlock(nn.Module):
 
         x = residual + x
         return x
-    
 
 
 class GFSQ(nn.Module):
 
-    def __init__(self, 
-            dim, levels, G, R, eps=1e-5, transpose = True
-        ):
+    def __init__(self, dim, levels, G, R, eps=1e-5, transpose=True):
         super(GFSQ, self).__init__()
         self.quantizer = GroupedResidualFSQ(
             dim=dim,
@@ -65,50 +70,74 @@ class GFSQ(nn.Module):
         self.transpose = transpose
         self.G = G
         self.R = R
-        
+
     def _embed(self, x):
         if self.transpose:
-            x = x.transpose(1,2)
+            x = x.transpose(1, 2)
         x = rearrange(
-            x, "b t (g r) -> g b t r", g = self.G, r = self.R,
-        )  
+            x,
+            "b t (g r) -> g b t r",
+            g=self.G,
+            r=self.R,
+        )
         feat = self.quantizer.get_output_from_indices(x)
-        return feat.transpose(1,2) if self.transpose else feat
-        
-    def forward(self, x,):
+        return feat.transpose(1, 2) if self.transpose else feat
+
+    def forward(
+        self,
+        x,
+    ):
         if self.transpose:
-            x = x.transpose(1,2)
+            x = x.transpose(1, 2)
         feat, ind = self.quantizer(x)
         ind = rearrange(
-            ind, "g b t r ->b t (g r)",
-        )  
+            ind,
+            "g b t r ->b t (g r)",
+        )
         embed_onehot = F.one_hot(ind.long(), self.n_ind).to(x.dtype)
-        e_mean = torch.mean(embed_onehot, dim=[0,1])
+        e_mean = torch.mean(embed_onehot, dim=[0, 1])
         e_mean = e_mean / (e_mean.sum(dim=1) + self.eps).unsqueeze(1)
         perplexity = torch.exp(-torch.sum(e_mean * torch.log(e_mean + self.eps), dim=1))
-        
+
         return (
             torch.zeros(perplexity.shape, dtype=x.dtype, device=x.device),
-            feat.transpose(1,2) if self.transpose else feat,
+            feat.transpose(1, 2) if self.transpose else feat,
             perplexity,
             None,
-            ind.transpose(1,2) if self.transpose else ind,
+            ind.transpose(1, 2) if self.transpose else ind,
         )
-        
+
+
 class DVAEDecoder(nn.Module):
-    def __init__(self, idim, odim,
-                 n_layer = 12, bn_dim = 64, hidden = 256, 
-                 kernel = 7, dilation = 2, up = False
-                ):
+    def __init__(
+        self,
+        idim,
+        odim,
+        n_layer=12,
+        bn_dim=64,
+        hidden=256,
+        kernel=7,
+        dilation=2,
+        up=False,
+    ):
         super().__init__()
         self.up = up
         self.conv_in = nn.Sequential(
-            nn.Conv1d(idim, bn_dim, 3, 1, 1), nn.GELU(),
-            nn.Conv1d(bn_dim, hidden, 3, 1, 1)
+            nn.Conv1d(idim, bn_dim, 3, 1, 1),
+            nn.GELU(),
+            nn.Conv1d(bn_dim, hidden, 3, 1, 1),
+        )
+        self.decoder_block = nn.ModuleList(
+            [
+                ConvNeXtBlock(
+                    hidden,
+                    hidden * 4,
+                    kernel,
+                    dilation,
+                )
+                for _ in range(n_layer)
+            ]
         )
-        self.decoder_block = nn.ModuleList([
-            ConvNeXtBlock(hidden, hidden* 4, kernel, dilation,)
-            for _ in range(n_layer)])
         self.conv_out = nn.Conv1d(hidden, odim, kernel_size=1, bias=False)
 
     def forward(self, input, conditioning=None):
@@ -117,17 +146,15 @@ class DVAEDecoder(nn.Module):
         x = self.conv_in(x)
         for f in self.decoder_block:
             x = f(x, conditioning)
-        
+
         x = self.conv_out(x)
         return x.transpose(1, 2)
-    
+
 
 class DVAE(nn.Module):
-    def __init__(
-        self, decoder_config, vq_config, dim=512
-    ):
+    def __init__(self, decoder_config, vq_config, dim=512):
         super().__init__()
-        self.register_buffer('coef', torch.randn(1, 100, 1))
+        self.register_buffer("coef", torch.randn(1, 100, 1))
 
         self.decoder = DVAEDecoder(**decoder_config)
         self.out_conv = nn.Conv1d(dim, 100, 3, 1, 1, bias=False)
@@ -142,10 +169,14 @@ class DVAE(nn.Module):
             vq_feats = self.vq_layer._embed(inp)
         else:
             vq_feats = inp.detach().clone()
-            
-        vq_feats = vq_feats.view(
-            (vq_feats.size(0), 2, vq_feats.size(1)//2, vq_feats.size(2)),
-        ).permute(0, 2, 3, 1).flatten(2)
+
+        vq_feats = (
+            vq_feats.view(
+                (vq_feats.size(0), 2, vq_feats.size(1) // 2, vq_feats.size(2)),
+            )
+            .permute(0, 2, 3, 1)
+            .flatten(2)
+        )
 
         vq_feats = vq_feats.transpose(1, 2)
         dec_out = self.decoder(input=vq_feats)

modules/ChatTTS/ChatTTS/model/gpt.py

@@ -1,19 +1,20 @@
 import os
+
 os.environ["TOKENIZERS_PARALLELISM"] = "false"
 
 import logging
-from tqdm import tqdm
-from einops import rearrange
-from transformers.cache_utils import Cache
 
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import torch.nn.utils.parametrize as P
+from einops import rearrange
 from torch.nn.utils.parametrizations import weight_norm
-from transformers import LlamaModel, LlamaConfig
-    
-    
+from tqdm import tqdm
+from transformers import LlamaConfig, LlamaModel
+from transformers.cache_utils import Cache
+
+
 class LlamaMLP(nn.Module):
     def __init__(self, hidden_size, intermediate_size):
         super().__init__()
@@ -27,70 +28,106 @@ class LlamaMLP(nn.Module):
     def forward(self, x):
         down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))
         return down_proj
-    
-    
+
+
 class GPT_warpper(nn.Module):
     def __init__(
-        self, 
-        gpt_config, 
+        self,
+        gpt_config,
         num_audio_tokens,
         num_text_tokens,
         num_vq=4,
         **kwargs,
-        ):
+    ):
         super().__init__()
 
         self.logger = logging.getLogger(__name__)
         self.gpt = self.build_model(gpt_config)
-        self.model_dim = self.gpt.config.hidden_size 
+        self.model_dim = self.gpt.config.hidden_size
 
         self.num_vq = num_vq
-        self.emb_code = nn.ModuleList([nn.Embedding(num_audio_tokens, self.model_dim) for i in range(self.num_vq)])
+        self.emb_code = nn.ModuleList(
+            [nn.Embedding(num_audio_tokens, self.model_dim) for i in range(self.num_vq)]
+        )
         self.emb_text = nn.Embedding(num_text_tokens, self.model_dim)
-        self.head_text = weight_norm(nn.Linear(self.model_dim, num_text_tokens, bias=False), name='weight')
-        self.head_code = nn.ModuleList([weight_norm(nn.Linear(self.model_dim, num_audio_tokens, bias=False), name='weight') for i in range(self.num_vq)])
+        self.head_text = weight_norm(
+            nn.Linear(self.model_dim, num_text_tokens, bias=False), name="weight"
+        )
+        self.head_code = nn.ModuleList(
+            [
+                weight_norm(
+                    nn.Linear(self.model_dim, num_audio_tokens, bias=False),
+                    name="weight",
+                )
+                for i in range(self.num_vq)
+            ]
+        )
 
     def build_model(self, config):
-        
+
         configuration = LlamaConfig(**config)
         model = LlamaModel(configuration)
         del model.embed_tokens
-        
+
         return model
-    
+
     def get_emb(self, input_ids, text_mask, **kwargs):
 
         emb_text = self.emb_text(input_ids[text_mask][:, 0])
-        
-        emb_code = [self.emb_code[i](input_ids[~text_mask][:, i]) for i in range(self.num_vq)]
+
+        emb_code = [
+            self.emb_code[i](input_ids[~text_mask][:, i]) for i in range(self.num_vq)
+        ]
         emb_code = torch.stack(emb_code, 2).sum(2)
-        
-        emb = torch.zeros((input_ids.shape[:-1])+(emb_text.shape[-1],), device=emb_text.device, dtype=emb_text.dtype)
+
+        emb = torch.zeros(
+            (input_ids.shape[:-1]) + (emb_text.shape[-1],),
+            device=emb_text.device,
+            dtype=emb_text.dtype,
+        )
         emb[text_mask] = emb_text
         emb[~text_mask] = emb_code.to(emb.dtype)
-        
+
         return emb
-    
+
     def prepare_inputs_for_generation(
-        self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, cache_position=None, **kwargs
+        self,
+        input_ids,
+        past_key_values=None,
+        attention_mask=None,
+        inputs_embeds=None,
+        cache_position=None,
+        **kwargs,
     ):
         # With static cache, the `past_key_values` is None
         # TODO joao: standardize interface for the different Cache classes and remove of this if
         has_static_cache = False
         if past_key_values is None:
-            past_key_values = getattr(self.gpt.layers[0].self_attn, "past_key_value", None)
+            past_key_values = getattr(
+                self.gpt.layers[0].self_attn, "past_key_value", None
+            )
             has_static_cache = past_key_values is not None
 
         past_length = 0
         if past_key_values is not None:
             if isinstance(past_key_values, Cache):
-                past_length = cache_position[0] if cache_position is not None else past_key_values.get_seq_length()
+                past_length = (
+                    cache_position[0]
+                    if cache_position is not None
+                    else past_key_values.get_seq_length()
+                )
                 max_cache_length = (
-                    torch.tensor(past_key_values.get_max_length(), device=input_ids.device)
+                    torch.tensor(
+                        past_key_values.get_max_length(), device=input_ids.device
+                    )
                     if past_key_values.get_max_length() is not None
                     else None
                 )
-                cache_length = past_length if max_cache_length is None else torch.min(max_cache_length, past_length)
+                cache_length = (
+                    past_length
+                    if max_cache_length is None
+                    else torch.min(max_cache_length, past_length)
+                )
             # TODO joao: remove this `else` after `generate` prioritizes `Cache` objects
             else:
                 cache_length = past_length = past_key_values[0][0].shape[2]
@@ -100,7 +137,10 @@ class GPT_warpper(nn.Module):
             # 1 - If the length of the attention_mask exceeds the length of input_ids, then we are in a setting where
             # some of the inputs are exclusively passed as part of the cache (e.g. when passing input_embeds as
             # input)
-            if attention_mask is not None and attention_mask.shape[1] > input_ids.shape[1]:
+            if (
+                attention_mask is not None
+                and attention_mask.shape[1] > input_ids.shape[1]
+            ):
                 input_ids = input_ids[:, -(attention_mask.shape[1] - past_length) :]
             # 2 - If the past_length is smaller than input_ids', then input_ids holds all input tokens. We can discard
             # input_ids based on the past_length.
@@ -133,9 +173,13 @@ class GPT_warpper(nn.Module):
             # TODO: use `next_tokens` directly instead.
             model_inputs = {"input_ids": input_ids.contiguous()}
 
-        input_length = position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
+        input_length = (
+            position_ids.shape[-1] if position_ids is not None else input_ids.shape[-1]
+        )
         if cache_position is None:
-            cache_position = torch.arange(past_length, past_length + input_length, device=input_ids.device)
+            cache_position = torch.arange(
+                past_length, past_length + input_length, device=input_ids.device
+            )
         else:
             cache_position = cache_position[-input_length:]
 
@@ -152,118 +196,154 @@ class GPT_warpper(nn.Module):
             }
         )
         return model_inputs
-    
+
     def generate(
-        self, 
-        emb, 
-        inputs_ids, 
-        temperature, 
-        eos_token, 
-        attention_mask = None,
-        max_new_token = 2048, 
-        min_new_token = 0,
-        LogitsWarpers = [],
-        LogitsProcessors = [],
+        self,
+        emb,
+        inputs_ids,
+        temperature,
+        eos_token,
+        attention_mask=None,
+        max_new_token=2048,
+        min_new_token=0,
+        LogitsWarpers=[],
+        LogitsProcessors=[],
         infer_text=False,
         return_attn=False,
         return_hidden=False,
-        disable_tqdm=False
+        disable_tqdm=False,
     ):
         if disable_tqdm:
             tqdm = lambda x: x
         else:
             from tqdm import tqdm
-            
-        with torch.no_grad():   
-        
+
+        with torch.no_grad():
+
             attentions = []
             hiddens = []
-            
-            start_idx, end_idx = inputs_ids.shape[1], torch.zeros(inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long)
+
+            start_idx, end_idx = inputs_ids.shape[1], torch.zeros(
+                inputs_ids.shape[0], device=inputs_ids.device, dtype=torch.long
+            )
             finish = torch.zeros(inputs_ids.shape[0], device=inputs_ids.device).bool()
-            
+
             temperature = temperature[None].expand(inputs_ids.shape[0], -1)
             temperature = rearrange(temperature, "b n -> (b n) 1")
 
-            attention_mask_cache = torch.ones((inputs_ids.shape[0], inputs_ids.shape[1]+max_new_token,), dtype=torch.bool, device=inputs_ids.device)
+            attention_mask_cache = torch.ones(
+                (
+                    inputs_ids.shape[0],
+                    inputs_ids.shape[1] + max_new_token,
+                ),
+                dtype=torch.bool,
+                device=inputs_ids.device,
+            )
             if attention_mask is not None:
-                attention_mask_cache[:, :attention_mask.shape[1]] = attention_mask
-            
+                attention_mask_cache[:, : attention_mask.shape[1]] = attention_mask
+
             for i in tqdm(range(max_new_token)):
                 if finish.all():
                     continue
-        
-                model_input = self.prepare_inputs_for_generation(inputs_ids, 
-                    outputs.past_key_values if i!=0 else None, 
-                    attention_mask_cache[:, :inputs_ids.shape[1]], use_cache=True)
-            
+
+                model_input = self.prepare_inputs_for_generation(
+                    inputs_ids,
+                    outputs.past_key_values if i != 0 else None,
+                    attention_mask_cache[:, : inputs_ids.shape[1]],
+                    use_cache=True,
+                )
+
                 if i == 0:
-                    model_input['inputs_embeds'] = emb
+                    model_input["inputs_embeds"] = emb
                 else:
                     if infer_text:
-                        model_input['inputs_embeds'] = self.emb_text(model_input['input_ids'][:,:,0])
+                        model_input["inputs_embeds"] = self.emb_text(
+                            model_input["input_ids"][:, :, 0]
+                        )
                     else:
-                        code_emb = [self.emb_code[i](model_input['input_ids'][:,:,i]) for i in range(self.num_vq)]
-                        model_input['inputs_embeds'] = torch.stack(code_emb, 3).sum(3)
-                
-                model_input['input_ids'] = None
+                        code_emb = [
+                            self.emb_code[i](model_input["input_ids"][:, :, i])
+                            for i in range(self.num_vq)
+                        ]
+                        model_input["inputs_embeds"] = torch.stack(code_emb, 3).sum(3)
+
+                model_input["input_ids"] = None
                 outputs = self.gpt.forward(**model_input, output_attentions=return_attn)
                 attentions.append(outputs.attentions)
-                hidden_states = outputs[0] # 🐻
+                hidden_states = outputs[0]  # 🐻
                 if return_hidden:
                     hiddens.append(hidden_states[:, -1])
 
                 with P.cached():
                     if infer_text:
-                        logits = self.head_text(hidden_states) 
+                        logits = self.head_text(hidden_states)
                     else:
-                        logits = torch.stack([self.head_code[i](hidden_states) for i in range(self.num_vq)], 3)
-        
+                        logits = torch.stack(
+                            [
+                                self.head_code[i](hidden_states)
+                                for i in range(self.num_vq)
+                            ],
+                            3,
+                        )
+
                 logits = logits[:, -1].float()
 
                 if not infer_text:
                     logits = rearrange(logits, "b c n -> (b n) c")
-                    logits_token = rearrange(inputs_ids[:, start_idx:], "b c n -> (b n) c")
+                    logits_token = rearrange(
+                        inputs_ids[:, start_idx:], "b c n -> (b n) c"
+                    )
                 else:
                     logits_token = inputs_ids[:, start_idx:, 0]
-                    
+
                 logits = logits / temperature
-                
+
                 for logitsProcessors in LogitsProcessors:
                     logits = logitsProcessors(logits_token, logits)
-                    
+
                 for logitsWarpers in LogitsWarpers:
                     logits = logitsWarpers(logits_token, logits)
-                    
+
                 if i < min_new_token:
                     logits[:, eos_token] = -torch.inf
-                
+
                 scores = F.softmax(logits, dim=-1)
-            
+
                 idx_next = torch.multinomial(scores, num_samples=1)
-                
+
                 if not infer_text:
                     idx_next = rearrange(idx_next, "(b n) 1 -> b n", n=self.num_vq)
                     finish = finish | (idx_next == eos_token).any(1)
                     inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(1)], 1)
                 else:
                     finish = finish | (idx_next == eos_token).any(1)
-                    inputs_ids = torch.cat([inputs_ids, idx_next.unsqueeze(-1).expand(-1, -1, self.num_vq)], 1)
+                    inputs_ids = torch.cat(
+                        [
+                            inputs_ids,
+                            idx_next.unsqueeze(-1).expand(-1, -1, self.num_vq),
+                        ],
+                        1,
+                    )
 
                 end_idx = end_idx + (~finish).int()
-            
-            inputs_ids = [inputs_ids[idx, start_idx: start_idx+i] for idx, i in enumerate(end_idx.int())]
+
+            inputs_ids = [
+                inputs_ids[idx, start_idx : start_idx + i]
+                for idx, i in enumerate(end_idx.int())
+            ]
             inputs_ids = [i[:, 0] for i in inputs_ids] if infer_text else inputs_ids
-            
+
             if return_hidden:
                 hiddens = torch.stack(hiddens, 1)
                 hiddens = [hiddens[idx, :i] for idx, i in enumerate(end_idx.int())]
-                    
+
             if not finish.all():
-                self.logger.warn(f'Incomplete result. hit max_new_token: {max_new_token}')    
-                   
+                self.logger.warn(
+                    f"Incomplete result. hit max_new_token: {max_new_token}"
+                )
+
             return {
-                'ids': inputs_ids, 
-                'attentions': attentions,
-                'hiddens':hiddens,
-            }
+                "ids": inputs_ids,
+                "attentions": attentions,
+                "hiddens": hiddens,
+            }

modules/ChatTTS/ChatTTS/utils/infer_utils.py

@@ -1,4 +1,5 @@
 import re
+
 import torch
 import torch.nn.functional as F
 

modules/ChatTTS/ChatTTS/utils/io_utils.py

@@ -1,14 +1,14 @@
-
-import os
 import logging
+import os
+
 
 def get_latest_modified_file(directory):
     logger = logging.getLogger(__name__)
-    
-    files = [os.path.join(directory, f) for f in os.listdir(directory)] 
+
+    files = [os.path.join(directory, f) for f in os.listdir(directory)]
     if not files:
-        logger.log(logging.WARNING, f'No files found in the directory: {directory}')
+        logger.log(logging.WARNING, f"No files found in the directory: {directory}")
         return None
     latest_file = max(files, key=os.path.getmtime)
 
-    return latest_file
+    return latest_file

modules/Denoiser/AudioDenoiser.py

@@ -1,15 +1,17 @@
 import logging
 import math
 from typing import Union
+
 import torch
 import torchaudio
-from torch import nn
-from audio_denoiser.helpers.torch_helper import batched_apply
-from modules.Denoiser.AudioNosiseModel import load_audio_denosier_model
 from audio_denoiser.helpers.audio_helper import (
     create_spectrogram,
     reconstruct_from_spectrogram,
 )
+from audio_denoiser.helpers.torch_helper import batched_apply
+from torch import nn
+
+from modules.Denoiser.AudioNosiseModel import load_audio_denosier_model
 
 _expected_t_std = 0.23
 _recommended_backend = "soundfile"

modules/Denoiser/AudioNosiseModel.py

@@ -1,12 +1,11 @@
+import json
+
 import torch
 import torch.nn as nn
-
 from audio_denoiser.modules.Permute import Permute
 from audio_denoiser.modules.SimpleRoberta import SimpleRoberta
 from audio_denoiser.modules.SpectrogramScaler import SpectrogramScaler
 
-import json
-
 
 class AudioNoiseModel(nn.Module):
     def __init__(self, config: dict):

modules/Enhancer/ResembleEnhance.py

@@ -1,20 +1,17 @@
 import gc
+import logging
+from pathlib import Path
+from threading import Lock
 from typing import Literal
 
 import numpy as np
+import torch
+
 from modules.devices import devices
 from modules.repos_static.resemble_enhance.enhancer.enhancer import Enhancer
 from modules.repos_static.resemble_enhance.enhancer.hparams import HParams
 from modules.repos_static.resemble_enhance.inference import inference
-
-import torch
-
 from modules.utils.constants import MODELS_DIR
-from pathlib import Path
-
-from threading import Lock
-
-import logging
 
 logger = logging.getLogger(__name__)
 
@@ -155,8 +152,8 @@ def apply_audio_enhance(
 
 
 if __name__ == "__main__":
-    import torchaudio
     import gradio as gr
+    import torchaudio
 
     device = torch.device("cuda")
 

modules/SentenceSplitter.py

@@ -1,4 +1,5 @@
 import re
+
 import zhon
 
 

modules/SynthesizeSegments.py

@@ -1,31 +1,37 @@
 import copy
+import json
+import logging
 import re
+from typing import List, Union
+
+import numpy as np
 from box import Box
 from pydub import AudioSegment
-from typing import List, Union
-from scipy.io.wavfile import write
-import io
-from modules.SentenceSplitter import SentenceSplitter
-from modules.api.utils import calc_spk_style
-from modules.ssml_parser.SSMLParser import SSMLSegment, SSMLBreak, SSMLContext
-from modules.utils import rng
-from modules.utils.audio import time_stretch, pitch_shift
+
 from modules import generate_audio
+from modules.api.utils import calc_spk_style
 from modules.normalization import text_normalize
-import logging
-import json
-
-from modules.speaker import Speaker, speaker_mgr
+from modules.SentenceSplitter import SentenceSplitter
+from modules.speaker import Speaker
+from modules.ssml_parser.SSMLParser import SSMLBreak, SSMLContext, SSMLSegment
+from modules.utils import rng
+from modules.utils.audio import pitch_shift, time_stretch
 
 logger = logging.getLogger(__name__)
 
 
-def audio_data_to_segment(audio_data, sr):
-    byte_io = io.BytesIO()
-    write(byte_io, rate=sr, data=audio_data)
-    byte_io.seek(0)
-
-    return AudioSegment.from_file(byte_io, format="wav")
+def audio_data_to_segment(audio_data: np.ndarray, sr: int):
+    """
+    optimize: https://github.com/lenML/ChatTTS-Forge/issues/57
+    """
+    audio_data = (audio_data * 32767).astype(np.int16)
+    audio_segment = AudioSegment(
+        audio_data.tobytes(),
+        frame_rate=sr,
+        sample_width=audio_data.dtype.itemsize,
+        channels=1,
+    )
+    return audio_segment
 
 
 def combine_audio_segments(audio_segments: list[AudioSegment]) -> AudioSegment:

modules/api/Api.py

@@ -1,12 +1,10 @@
-from fastapi import FastAPI
-from fastapi.middleware.cors import CORSMiddleware
-
+import fnmatch
 import logging
 
+from fastapi import FastAPI
+from fastapi.middleware.cors import CORSMiddleware
 from fastapi.staticfiles import StaticFiles
 
-import fnmatch
-
 
 def is_excluded(path, exclude_patterns):
     """

modules/api/api_setup.py

@@ -1,26 +1,24 @@
-import logging
-from modules.Enhancer.ResembleEnhance import load_enhancer
-from modules.devices import devices
 import argparse
+import logging
 
-from modules import config
-from modules.models import load_chat_tts
-from modules.utils import env
-from modules import generate_audio
+from modules import config, generate_audio
 from modules.api.Api import APIManager
-
 from modules.api.impl import (
-    style_api,
-    tts_api,
-    ssml_api,
     google_api,
+    models_api,
     openai_api,
+    ping_api,
     refiner_api,
     speaker_api,
-    ping_api,
-    models_api,
+    ssml_api,
+    style_api,
+    tts_api,
     xtts_v2_api,
 )
+from modules.devices import devices
+from modules.Enhancer.ResembleEnhance import load_enhancer
+from modules.models import load_chat_tts
+from modules.utils import env
 
 logger = logging.getLogger(__name__)
 

modules/api/impl/google_api.py

@@ -1,22 +1,18 @@
 from typing import Union
-from fastapi import HTTPException
 
+from fastapi import HTTPException
 from pydantic import BaseModel
 
-
+from modules.api import utils as api_utils
 from modules.api.Api import APIManager
 from modules.api.impl.handler.SSMLHandler import SSMLHandler
 from modules.api.impl.handler.TTSHandler import TTSHandler
 from modules.api.impl.model.audio_model import AdjustConfig, AudioFormat
 from modules.api.impl.model.chattts_model import ChatTTSConfig, InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
-
 from modules.speaker import Speaker, speaker_mgr
 
 
-from modules.api import utils as api_utils
-
-
 class SynthesisInput(BaseModel):
     text: Union[str, None] = None
     ssml: Union[str, None] = None

modules/api/impl/handler/AudioHandler.py

@@ -1,10 +1,11 @@
 import base64
 import io
+
 import numpy as np
 import soundfile as sf
 
-from modules.api.impl.model.audio_model import AudioFormat
 from modules.api import utils as api_utils
+from modules.api.impl.model.audio_model import AudioFormat
 
 
 class AudioHandler:

modules/api/impl/handler/SSMLHandler.py

@@ -1,14 +1,14 @@
-from fastapi import HTTPException
 import numpy as np
+from fastapi import HTTPException
 
-from modules.Enhancer.ResembleEnhance import apply_audio_enhance_full
-from modules.SynthesizeSegments import SynthesizeSegments, combine_audio_segments
 from modules.api.impl.handler.AudioHandler import AudioHandler
 from modules.api.impl.model.audio_model import AdjustConfig
 from modules.api.impl.model.chattts_model import InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
+from modules.Enhancer.ResembleEnhance import apply_audio_enhance_full
 from modules.normalization import text_normalize
 from modules.ssml_parser.SSMLParser import create_ssml_parser
+from modules.SynthesizeSegments import SynthesizeSegments, combine_audio_segments
 from modules.utils import audio
 
 

modules/api/impl/handler/TTSHandler.py

@@ -1,13 +1,13 @@
 import numpy as np
-from modules.Enhancer.ResembleEnhance import apply_audio_enhance_full
+
 from modules.api.impl.handler.AudioHandler import AudioHandler
 from modules.api.impl.model.audio_model import AdjustConfig
 from modules.api.impl.model.chattts_model import ChatTTSConfig, InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
+from modules.Enhancer.ResembleEnhance import apply_audio_enhance_full
 from modules.normalization import text_normalize
 from modules.speaker import Speaker
 from modules.synthesize_audio import synthesize_audio
-
 from modules.utils.audio import apply_prosody_to_audio_data
 
 

modules/api/impl/model/enhancer_model.py

@@ -1,4 +1,5 @@
 from typing import Literal
+
 from pydantic import BaseModel
 
 

modules/api/impl/models_api.py

@@ -1,6 +1,6 @@
-from modules.Enhancer.ResembleEnhance import reload_enhancer, unload_enhancer
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
+from modules.Enhancer.ResembleEnhance import reload_enhancer, unload_enhancer
 from modules.models import reload_chat_tts, unload_chat_tts
 
 

modules/api/impl/openai_api.py

@@ -1,23 +1,18 @@
-from fastapi import File, Form, HTTPException, Body, UploadFile
+from typing import List, Optional
 
+from fastapi import Body, File, Form, HTTPException, UploadFile
+from fastapi.responses import StreamingResponse
 from numpy import clip
 from pydantic import BaseModel, Field
-from fastapi.responses import StreamingResponse
-
 
+from modules.api import utils as api_utils
+from modules.api.Api import APIManager
 from modules.api.impl.handler.TTSHandler import TTSHandler
 from modules.api.impl.model.audio_model import AdjustConfig, AudioFormat
 from modules.api.impl.model.chattts_model import ChatTTSConfig, InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
-
-
-from typing import List, Optional
-
-from modules.api import utils as api_utils
-from modules.api.Api import APIManager
-
-from modules.speaker import Speaker, speaker_mgr
 from modules.data import styles_mgr
+from modules.speaker import Speaker, speaker_mgr
 
 
 class AudioSpeechRequest(BaseModel):

modules/api/impl/ping_api.py

@@ -1,8 +1,7 @@
+from modules import config
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
 
-from modules import config
-
 
 def setup(app: APIManager):
     @app.get("/v1/ping", response_model=api_utils.BaseResponse)

modules/api/impl/refiner_api.py

@@ -1,10 +1,7 @@
 from fastapi import HTTPException
-
 from pydantic import BaseModel
 
-
 from modules import refiner
-
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
 from modules.normalization import text_normalize

modules/api/impl/speaker_api.py

@@ -1,9 +1,10 @@
+import torch
 from fastapi import HTTPException
 from pydantic import BaseModel
-import torch
-from modules.speaker import speaker_mgr
+
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
+from modules.speaker import speaker_mgr
 
 
 class CreateSpeaker(BaseModel):

modules/api/impl/ssml_api.py

@@ -1,19 +1,14 @@
-from fastapi import HTTPException, Body
-from fastapi.responses import StreamingResponse
-
+from fastapi import Body, HTTPException
+from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel
-from fastapi.responses import FileResponse
-
 
+from modules.api.Api import APIManager
 from modules.api.impl.handler.SSMLHandler import SSMLHandler
 from modules.api.impl.model.audio_model import AdjustConfig, AudioFormat
 from modules.api.impl.model.chattts_model import InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
 
 
-from modules.api.Api import APIManager
-
-
 class SSMLRequest(BaseModel):
     ssml: str
     format: AudioFormat = "mp3"

modules/api/impl/style_api.py

@@ -1,6 +1,6 @@
-from modules.data import styles_mgr
 from modules.api import utils as api_utils
 from modules.api.Api import APIManager
+from modules.data import styles_mgr
 
 
 async def list_styles():

modules/api/impl/tts_api.py

@@ -1,17 +1,13 @@
 from fastapi import Depends, HTTPException, Query
-from fastapi.responses import StreamingResponse
-
+from fastapi.responses import FileResponse, StreamingResponse
 from pydantic import BaseModel
-from fastapi.responses import FileResponse
-
 
+from modules.api import utils as api_utils
+from modules.api.Api import APIManager
 from modules.api.impl.handler.TTSHandler import TTSHandler
 from modules.api.impl.model.audio_model import AdjustConfig, AudioFormat
 from modules.api.impl.model.chattts_model import ChatTTSConfig, InferConfig
 from modules.api.impl.model.enhancer_model import EnhancerConfig
-
-from modules.api import utils as api_utils
-from modules.api.Api import APIManager
 from modules.speaker import Speaker
 
 

modules/api/impl/xtts_v2_api.py

@@ -1,19 +1,17 @@
 import io
+import logging
+
+import soundfile as sf
 from fastapi import HTTPException
 from fastapi.responses import StreamingResponse
 from pydantic import BaseModel
-from modules.api import utils as api_utils
-from modules.api.Api import APIManager
-
-import soundfile as sf
 
 from modules import config
+from modules.api import utils as api_utils
+from modules.api.Api import APIManager
 from modules.normalization import text_normalize
 from modules.speaker import speaker_mgr
 from modules.synthesize_audio import synthesize_audio
-
-import logging
-
 from modules.utils.audio import apply_prosody_to_audio_data
 
 logger = logging.getLogger(__name__)

modules/api/utils.py

@@ -1,14 +1,10 @@
-from pydantic import BaseModel
 from typing import Any, Union
 
-
-from modules.speaker import speaker_mgr
-
-
-from modules.data import styles_mgr
-
+from pydantic import BaseModel
 from pydub import AudioSegment
 
+from modules.data import styles_mgr
+from modules.speaker import speaker_mgr
 from modules.ssml import merge_prompt
 
 

modules/api/worker.py

@@ -1,6 +1,7 @@
 import argparse
 import logging
 import os
+
 import dotenv
 from fastapi import FastAPI
 
@@ -12,6 +13,7 @@ logging.basicConfig(
     format="%(asctime)s - %(name)s - %(levelname)s - %(message)s",
 )
 
+from modules import config
 from modules.api.api_setup import (
     process_api_args,
     process_model_args,
@@ -20,7 +22,6 @@ from modules.api.api_setup import (
     setup_uvicon_args,
 )
 from modules.api.app_config import app_description, app_title, app_version
-from modules import config
 from modules.utils.torch_opt import configure_torch_optimizations
 
 dotenv.load_dotenv(

modules/config.py

@@ -1,9 +1,9 @@
 import sys
 
 import torch
-from modules.utils.JsonObject import JsonObject
 
-from modules.utils import git, ffmpeg
+from modules.utils import ffmpeg, git
+from modules.utils.JsonObject import JsonObject
 
 # TODO impl RuntimeEnvVars() class
 runtime_env_vars = JsonObject({})

modules/data.py

@@ -1,6 +1,5 @@
 from modules.utils.CsvMgr import BaseManager
 
-
 # speakers_mgr = BaseManager("./data/speakers.csv")
 styles_mgr = BaseManager("./data/styles.csv")
 

modules/denoise.py

@@ -1,15 +1,13 @@
 import os
 from typing import Union
 
+import soundfile as sf
 import torch
 import torchaudio
-from modules.Denoiser.AudioDenoiser import AudioDenoiser
-
-from modules.utils.constants import MODELS_DIR
 
+from modules.Denoiser.AudioDenoiser import AudioDenoiser
 from modules.devices import devices
-
-import soundfile as sf
+from modules.utils.constants import MODELS_DIR
 
 ad: Union[AudioDenoiser, None] = None
 

modules/devices/devices.py

@@ -1,9 +1,10 @@
-from functools import lru_cache
+import logging
 import sys
+from functools import lru_cache
+
 import torch
-from modules import config
 
-import logging
+from modules import config
 
 logger = logging.getLogger(__name__)
 

modules/devices/mac_devices.py

@@ -1,8 +1,9 @@
-import torch
 import logging
-from packaging import version
+
+import torch
 import torch.backends
 import torch.backends.mps
+from packaging import version
 
 logger = logging.getLogger(__name__)
 

modules/ffmpeg_env.py

@@ -1,6 +1,7 @@
+import logging
 import os
+
 from modules.utils.constants import ROOT_DIR
-import logging
 
 logger = logging.getLogger(__name__)
 

modules/finetune/train_speaker.py

@@ -3,9 +3,10 @@ import torch.nn.functional as F
 import transformers
 
 from modules.finetune.model.encoder import DVAEEncoder, get_encoder_config
-from modules.finetune.utils.output import get_ansi_len, output_iter, ansi
-from .utils.logger import MetricLogger
+from modules.finetune.utils.output import ansi, get_ansi_len, output_iter
+
 from .utils.dataset import AudioCollator, XzListTar
+from .utils.logger import MetricLogger
 from .utils.model import quantize
 
 IGNORE_TOKEN_ID = transformers.trainer_pt_utils.LabelSmoother.ignore_index
@@ -201,11 +202,13 @@ def train_speaker_embeddings(
 if __name__ == "__main__":
     import argparse
     import os
-    import numpy as np
     import pathlib
-    from modules.models import load_chat_tts
-    from modules.devices import devices
+
+    import numpy as np
+
     from modules import config
+    from modules.devices import devices
+    from modules.models import load_chat_tts
     from modules.speaker import Speaker
 
     config.runtime_env_vars.no_half = True

modules/finetune/utils/dataset.py

@@ -1,21 +1,21 @@
-import os
+import abc
 import functools
-import json
-import tarfile
 import io
+import json
 import logging
-import abc
+import os
+import tarfile
 import typing
 
 import torch.utils.data
 import torchaudio
-from torchvision.datasets.utils import download_url
 import transformers
 import vocos
+from torchvision.datasets.utils import download_url
 
 from modules.ChatTTS.ChatTTS.utils.infer_utils import (
-    count_invalid_characters,
     apply_character_map,
+    count_invalid_characters,
 )
 
 

modules/finetune/utils/logger.py

@@ -3,15 +3,14 @@
 import statistics
 import time
 from collections import defaultdict, deque
-from tqdm import tqdm as tqdm_class
-
 from typing import Generator, Iterable, TypeVar
-from typing_extensions import Self
 
 import torch
 import torch.distributed as dist
+from tqdm import tqdm as tqdm_class
+from typing_extensions import Self
 
-from .output import ansi, prints, get_ansi_len
+from .output import ansi, get_ansi_len, prints
 
 __all__ = ["SmoothedValue", "MetricLogger"]
 

modules/generate_audio.py

@@ -1,18 +1,15 @@
+import gc
+import logging
+from typing import Union
+
 import numpy as np
 import torch
 
-from modules.speaker import Speaker
-from modules.utils.SeedContext import SeedContext
-
-from modules import models, config
-
-import logging
-import gc
-
+from modules import config, models
 from modules.devices import devices
-from typing import Union
-
+from modules.speaker import Speaker
 from modules.utils.cache import conditional_cache
+from modules.utils.SeedContext import SeedContext
 
 logger = logging.getLogger(__name__)
 

modules/models.py

@@ -1,13 +1,13 @@
+import gc
+import logging
 import threading
+
 import torch
-from modules.ChatTTS import ChatTTS
+
 from modules import config
+from modules.ChatTTS import ChatTTS
 from modules.devices import devices
 
-import logging
-import gc
-
-
 logger = logging.getLogger(__name__)
 
 chat_tts = None

modules/normalization.py

@@ -1,9 +1,11 @@
+import re
 from functools import lru_cache
-from modules.utils.zh_normalization.text_normlization import *
+
 import emojiswitch
-from modules.utils.markdown import markdown_to_text
+
 from modules import models
-import re
+from modules.utils.markdown import markdown_to_text
+from modules.utils.zh_normalization.text_normlization import *
 
 # 是否关闭 unk token 检查
 # NOTE: 单测的时候用于跳过模型加载

modules/prompts/news_oral_prompt.txt

@@ -1,7 +1,7 @@
-# 任务要求
-任务： 新闻稿口播化
+#任务要求
+任务：新闻稿口播化
 
-你需要将一个新闻稿改写为口语化的口播文本
+你需要将一个新闻稿改写为口语化的口播文本，以提供给新闻主播在晚间新闻节目中播报
 同时，适当的添加一些 附语言 标签为文本增加多样性
 
 目前可以使用的附语言标签如下:
@@ -10,5 +10,24 @@
 - `[v_break]`: 表示有声停顿，如“嗯”、“啊”等
 - `[lbreak]`: 表示一个长停顿一般表示段落结束
 
-# 输入
+# examples
+## case 1
+- input: `天气预报显示，今天会有小雨，请大家出门时记得带伞。降温的天气也提醒我们要适时添衣保暖`
+- output: `天气预报显示，今天会有小雨，请大家出门时记得带伞[uv_break]。那降温的天气[uv_break]也提醒我们要适时添衣保暖[lbreak]`
+
+## case 2
+- input: `请注意，电梯将在下午两点进行例行维护，预计需要一个小时的时间，请大家在此期间使用楼梯`
+- output: `请注意啊，这个电梯将在下午两点进行[uv_break]例行维护[uv_break]，预计需要一个小时的时间[uv_break]，请大家在此期间使用楼梯[lbreak]`
+
+## case 3
+- input: `它的任务是简化记者编辑的工作流程。记者写稿时可以用标签来标明关键词、标题或主题。随着时间推移，数据积累到一定程度后，机器编辑就能自动识别这些标签`
+- output: `它的任务呢是简化记者编辑的工作流程[uv_break]。记者写稿时呢可以用标签来标明关键词啊、标题啊或主题[uv_break]。那随着时间推移呢，数据积累到一定程度后[uv_break]，机器编辑就能自动识别这些标签[uv_break]`
+
+## case 4
+- input: `有一天，小明问他爸爸：“爸爸，我是不是傻孩子啊？”
+
+爸爸说：“傻孩子，你怎么会是傻孩子呢？”`
+- output: `然后有一天呢，小明问他[uv_break]爸爸[uv_break]，爸爸，我是不是傻孩[uv_break]子啊？爸爸说，傻孩[laugh]子啊，你怎么会是傻孩子呢[laugh]？`
+
+# 用户输入
 {{USER_INPUT}}

modules/refiner.py

@@ -1,10 +1,9 @@
 import numpy as np
 import torch
 
+from modules import config, models
 from modules.utils.SeedContext import SeedContext
 
-from modules import models, config
-
 
 @torch.inference_mode()
 def refine_text(

modules/repos_static/resemble_enhance/common.py

@@ -42,7 +42,9 @@ class Normalizer(nn.Module):
             self.running_var_unsafe = x.var()
         else:
             self.running_mean_unsafe = self._ema(self.running_mean_unsafe, x.mean())
-            self.running_var_unsafe = self._ema(self.running_var_unsafe, (x - self.running_mean).pow(2).mean())
+            self.running_var_unsafe = self._ema(
+                self.running_var_unsafe, (x - self.running_mean).pow(2).mean()
+            )
 
     def forward(self, x: Tensor, update=True):
         if self.training and update:

modules/repos_static/resemble_enhance/data/dataset.py

@@ -44,7 +44,9 @@ def praat_augment(wav, sr):
     sound = parselmouth.Sound(wav, sr)
     formant_shift_ratio = random.uniform(1.1, 1.5)
     pitch_range_factor = random.uniform(0.5, 2.0)
-    sound = parselmouth.praat.call(sound, "Change gender", 75, 600, formant_shift_ratio, 0, pitch_range_factor, 1.0)
+    sound = parselmouth.praat.call(
+        sound, "Change gender", 75, 600, formant_shift_ratio, 0, pitch_range_factor, 1.0
+    )
     wav = np.array(sound.values)[0].astype(np.float32)
     return wav
 
@@ -73,7 +75,9 @@ class Dataset(DatasetBase):
         if len(self.bg_paths) == 0:
             raise ValueError(f"No background audio files found in {hp.bg_dir}")
 
-        logger.info(f"Found {len(self.fg_paths)} foreground files and {len(self.bg_paths)} background files")
+        logger.info(
+            f"Found {len(self.fg_paths)} foreground files and {len(self.bg_paths)} background files"
+        )
 
         self.training = training
         self.max_retries = max_retries
@@ -121,7 +125,9 @@ class Dataset(DatasetBase):
         fg_path = self.fg_paths[index]
 
         if self.training and random.random() < self.silent_fg_prob:
-            fg_wav = np.zeros(int(self.hp.training_seconds * self.hp.wav_rate), dtype=np.float32)
+            fg_wav = np.zeros(
+                int(self.hp.training_seconds * self.hp.wav_rate), dtype=np.float32
+            )
         else:
             fg_wav = self._load_wav(fg_path)
             if random.random() < self.hp.praat_augment_prob and self.training:
@@ -132,14 +138,20 @@ class Dataset(DatasetBase):
             fg_dwav = None
             bg_dwav = None
         else:
-            fg_dwav = _normalize(self.distorter(fg_wav, self.hp.wav_rate)).astype(np.float32)
+            fg_dwav = _normalize(self.distorter(fg_wav, self.hp.wav_rate)).astype(
+                np.float32
+            )
             if self.training:
                 bg_path = random.choice(self.bg_paths)
             else:
                 # Deterministic for validation
                 bg_path = self.bg_paths[index % len(self.bg_paths)]
-            bg_wav = self._load_wav(bg_path, length=len(fg_wav), random_crop=self.training)
-            bg_dwav = _normalize(self.distorter(bg_wav, self.hp.wav_rate)).astype(np.float32)
+            bg_wav = self._load_wav(
+                bg_path, length=len(fg_wav), random_crop=self.training
+            )
+            bg_dwav = _normalize(self.distorter(bg_wav, self.hp.wav_rate)).astype(
+                np.float32
+            )
 
         return dict(
             fg_wav=fg_wav,
@@ -154,7 +166,9 @@ class Dataset(DatasetBase):
                 return self._getitem_unsafe(index)
             except Exception as e:
                 if i == self.max_retries - 1:
-                    raise RuntimeError(f"Failed to load {self.fg_paths[index]} after {self.max_retries} retries") from e
+                    raise RuntimeError(
+                        f"Failed to load {self.fg_paths[index]} after {self.max_retries} retries"
+                    ) from e
                 logger.debug(f"Error loading {self.fg_paths[index]}: {e}, skipping")
                 index = np.random.randint(0, len(self))
 

modules/repos_static/resemble_enhance/data/distorter/base.py

@@ -2,8 +2,8 @@ import itertools
 import os
 import random
 import time
-from typing import Union
 import warnings
+from typing import Union
 
 import numpy as np