refactor: Simple code cleanups (from upstream PR yl4579#219)

- Remove unused imports and variables
- Replace wildcard imports with specific imports
- Fix Flake8 errors
- Remove IPython tracing in scripts

Original PR: yl4579#219
Author: Arvind Suresh <arvind@free-speech.ai>

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 4.5 <noreply@anthropic.com>

Author: ayutaz

Modules/diffusion/diffusion.py

@@ -1,14 +1,7 @@
-from math import pi
-from random import randint
-from typing import Any, Optional, Sequence, Tuple, Union
-
-import torch
-from einops import rearrange
 from torch import Tensor, nn
-from tqdm import tqdm
 
-from .utils import *
-from .sampler import *
+from .utils import groupby 
+from .sampler import UniformDistribution, LinearSchedule, VSampler
 
 """
 Diffusion Classes (generic for 1d data)

Modules/diffusion/modules.py

@@ -1,16 +1,16 @@
-from math import floor, log, pi
-from typing import Any, List, Optional, Sequence, Tuple, Union
+from math import log, pi
+from typing import Optional
 
-from .utils import *
+from .utils import default, exists, rand_bool
 
 import torch
 import torch.nn as nn
+import torch.nn.functional as F
 from einops import rearrange, reduce, repeat
 from einops.layers.torch import Rearrange
 from einops_exts import rearrange_many
 from torch import Tensor, einsum
 
-
 """
 Utils
 """

Modules/diffusion/sampler.py

@@ -7,7 +7,7 @@
 from einops import rearrange, reduce
 from torch import Tensor
 
-from .utils import *
+from .utils import default, exists
 
 """
 Diffusion Training
@@ -213,7 +213,6 @@ def loss_weight(self, sigmas: Tensor) -> Tensor:
 
     def forward(self, x: Tensor, noise: Tensor = None, **kwargs) -> Tensor:
         batch_size, device = x.shape[0], x.device
-        from einops import rearrange, reduce
 
         # Sample amount of noise to add for each batch element
         sigmas = self.sigma_distribution(num_samples=batch_size, device=device)

Modules/diffusion/utils.py

@@ -1,12 +1,9 @@
 from functools import reduce
 from inspect import isfunction
-from math import ceil, floor, log2, pi
+from math import ceil, floor, log2
 from typing import Callable, Dict, List, Optional, Sequence, Tuple, TypeVar, Union
 
 import torch
-import torch.nn.functional as F
-from einops import rearrange
-from torch import Generator, Tensor
 from typing_extensions import TypeGuard
 
 T = TypeVar("T")

Modules/discriminators.py

@@ -1,7 +1,7 @@
 import torch
 import torch.nn.functional as F
 import torch.nn as nn
-from torch.nn import Conv1d, AvgPool1d, Conv2d
+from torch.nn import Conv1d, Conv2d
 from torch.nn.utils import weight_norm, spectral_norm
 
 from .utils import get_padding
@@ -21,8 +21,6 @@ def stft(x, fft_size, hop_size, win_length, window):
     """
     x_stft = torch.stft(x, fft_size, hop_size, win_length, window,
             return_complex=True)
-    real = x_stft[..., 0]
-    imag = x_stft[..., 1]
 
     return torch.abs(x_stft).transpose(2, 1)
 
@@ -31,7 +29,7 @@ class SpecDiscriminator(nn.Module):
 
     def __init__(self, fft_size=1024, shift_size=120, win_length=600, window="hann_window", use_spectral_norm=False):
         super(SpecDiscriminator, self).__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.fft_size = fft_size
         self.shift_size = shift_size
         self.win_length = win_length
@@ -97,7 +95,7 @@ class DiscriminatorP(torch.nn.Module):
     def __init__(self, period, kernel_size=5, stride=3, use_spectral_norm=False):
         super(DiscriminatorP, self).__init__()
         self.period = period
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.convs = nn.ModuleList([
             norm_f(Conv2d(1, 32, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
             norm_f(Conv2d(32, 128, (kernel_size, 1), (stride, 1), padding=(get_padding(5, 1), 0))),
@@ -118,8 +116,8 @@ def forward(self, x):
             t = t + n_pad
         x = x.view(b, c, t // self.period, self.period)
 
-        for l in self.convs:
-            x = l(x)
+        for layer in self.convs:
+            x = layer(x)
             x = F.leaky_relu(x, LRELU_SLOPE)
             fmap.append(x)
         x = self.conv_post(x)
@@ -163,7 +161,7 @@ def __init__(self, slm_hidden=768,
                  initial_channel=64, 
                  use_spectral_norm=False):
         super(WavLMDiscriminator, self).__init__()
-        norm_f = weight_norm if use_spectral_norm == False else spectral_norm
+        norm_f = weight_norm if use_spectral_norm is False else spectral_norm
         self.pre = norm_f(Conv1d(slm_hidden * slm_layers, initial_channel, 1, 1, padding=0))
 
         self.convs = nn.ModuleList([
@@ -178,11 +176,11 @@ def forward(self, x):
         x = self.pre(x)
 
         fmap = []
-        for l in self.convs:
-            x = l(x)
+        for layer in self.convs:
+            x = layer(x)
             x = F.leaky_relu(x, LRELU_SLOPE)
             fmap.append(x)
         x = self.conv_post(x)
         x = torch.flatten(x, 1, -1)
 
-        return x
+        return x

Modules/hifigan.py

@@ -1,8 +1,8 @@
 import torch
 import torch.nn.functional as F
 import torch.nn as nn
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
 from .utils import init_weights, get_padding
 
 import math
@@ -74,10 +74,10 @@ def forward(self, x, s):
         return x
 
     def remove_weight_norm(self):
-        for l in self.convs1:
-            remove_weight_norm(l)
-        for l in self.convs2:
-            remove_weight_norm(l)
+        for layer in self.convs1:
+            remove_weight_norm(layer)
+        for layer in self.convs2:
+            remove_weight_norm(layer)
 
 class SineGen(torch.nn.Module):
     """ Definition of sine generator
@@ -193,8 +193,7 @@ def forward(self, f0):
         output sine_tensor: tensor(batchsize=1, length, dim)
         output uv: tensor(batchsize=1, length, 1)
         """
-        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
-                             device=f0.device)
+        torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
         # fundamental component
         fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
 
@@ -348,10 +347,10 @@ def forward(self, x, s, f0):
 
     def remove_weight_norm(self):
         print('Removing weight norm...')
-        for l in self.ups:
-            remove_weight_norm(l)
-        for l in self.resblocks:
-            l.remove_weight_norm()
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
         remove_weight_norm(self.conv_pre)
         remove_weight_norm(self.conv_post)
 
@@ -474,4 +473,4 @@ def forward(self, asr, F0_curve, N, s):
         x = self.generator(x, s, F0_curve)
         return x
 
-    
+    

Modules/istftnet.py

@@ -1,8 +1,8 @@
 import torch
 import torch.nn.functional as F
 import torch.nn as nn
-from torch.nn import Conv1d, ConvTranspose1d, AvgPool1d, Conv2d
-from torch.nn.utils import weight_norm, remove_weight_norm, spectral_norm
+from torch.nn import Conv1d, ConvTranspose1d
+from torch.nn.utils import weight_norm, remove_weight_norm
 from .utils import init_weights, get_padding
 
 import math
@@ -75,10 +75,10 @@ def forward(self, x, s):
         return x
 
     def remove_weight_norm(self):
-        for l in self.convs1:
-            remove_weight_norm(l)
-        for l in self.convs2:
-            remove_weight_norm(l)
+        for layer in self.convs1:
+            remove_weight_norm(layer)
+        for layer in self.convs2:
+            remove_weight_norm(layer)
 
 class TorchSTFT(torch.nn.Module):
     def __init__(self, filter_length=800, hop_length=200, win_length=800, window='hann'):
@@ -222,8 +222,7 @@ def forward(self, f0):
         output sine_tensor: tensor(batchsize=1, length, dim)
         output uv: tensor(batchsize=1, length, 1)
         """
-        f0_buf = torch.zeros(f0.shape[0], f0.shape[1], self.dim,
-                             device=f0.device)
+        torch.zeros(f0.shape[0], f0.shape[1], self.dim, device=f0.device)
         # fundamental component
         fn = torch.multiply(f0, torch.FloatTensor([[range(1, self.harmonic_num + 2)]]).to(f0.device))
 
@@ -399,10 +398,10 @@ def fw_phase(self, x, s):
 
     def remove_weight_norm(self):
         print('Removing weight norm...')
-        for l in self.ups:
-            remove_weight_norm(l)
-        for l in self.resblocks:
-            l.remove_weight_norm()
+        for layer in self.ups:
+            remove_weight_norm(layer)
+        for layer in self.resblocks:
+            layer.remove_weight_norm()
         remove_weight_norm(self.conv_pre)
         remove_weight_norm(self.conv_post)
 
@@ -527,4 +526,4 @@ def forward(self, asr, F0_curve, N, s):
         x = self.generator(x, s, F0_curve)
         return x
 
-    
+    

Modules/slmadv.py

@@ -41,7 +41,7 @@ def forward(self, iters, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_t
                          num_steps=num_steps).squeeze(1)
 
         s_dur = s_preds[:, 128:]
-        s = s_preds[:, :128]
+        s_preds[:, :128]
 
         d, _ = self.model.predictor(d_en, s_dur, 
                                                 ref_lengths, 
@@ -61,20 +61,20 @@ def forward(self, iters, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_t
             _s2s_pred = torch.sigmoid(_s2s_pred_org)
             _dur_pred = _s2s_pred.sum(axis=-1)
 
-            l = int(torch.round(_s2s_pred.sum()).item())
-            t = torch.arange(0, l).expand(l)
+            length = int(torch.round(_s2s_pred.sum()).item())
+            t = torch.arange(0, length).expand(length)
 
-            t = torch.arange(0, l).unsqueeze(0).expand((len(_s2s_pred), l)).to(ref_text.device)
+            t = torch.arange(0, length).unsqueeze(0).expand((len(_s2s_pred), length)).to(ref_text.device)
             loc = torch.cumsum(_dur_pred, dim=0) - _dur_pred / 2
 
-            h = torch.exp(-0.5 * torch.square(t - (l - loc.unsqueeze(-1))) / (self.sig)**2)
+            h = torch.exp(-0.5 * torch.square(t - (length - loc.unsqueeze(-1))) / (self.sig)**2)
 
             out = torch.nn.functional.conv1d(_s2s_pred_org.unsqueeze(0), 
                                          h.unsqueeze(1), 
-                                         padding=h.shape[-1] - 1, groups=int(_text_length))[..., :l]
+                                         padding=h.shape[-1] - 1, groups=int(_text_length))[..., :length]
             attn_preds.append(F.softmax(out.squeeze(), dim=0))
 
-            output_lengths.append(l)
+            output_lengths.append(length)
 
         max_len = max(output_lengths)
 
@@ -96,14 +96,9 @@ def forward(self, iters, y_rec_gt, y_rec_gt_pred, waves, mel_input_length, ref_t
         mel_len = min(mel_len, self.max_len // 2)
 
         # get clips
-        
         en = []
         p_en = []
         sp = []
-        
-        F0_fakes = []
-        N_fakes = []
-        
         wav = []
 
         for bib in range(len(output_lengths)):

Modules/utils.py

@@ -3,12 +3,5 @@ def init_weights(m, mean=0.0, std=0.01):
     if classname.find("Conv") != -1:
         m.weight.data.normal_(mean, std)
 
-
-def apply_weight_norm(m):
-    classname = m.__class__.__name__
-    if classname.find("Conv") != -1:
-        weight_norm(m)
-
-
 def get_padding(kernel_size, dilation=1):
-    return int((kernel_size*dilation - dilation)/2)
+    return int((kernel_size*dilation - dilation)/2)

Utils/ASR/layers.py

@@ -1,10 +1,6 @@
-import math
 import torch
 from torch import nn
-from typing import Optional, Any
-from torch import Tensor
 import torch.nn.functional as F
-import torchaudio
 import torchaudio.functional as audio_F
 
 import random