change to from_pretrained

app.py

@@ -1,9 +1,9 @@
 import gradio as gr
 import torch
 from model import ECAPA_gender
-
-model = ECAPA_gender({"C": 1024})
-model.load_state_dict(torch.load("gender_classifier.model", map_location="cpu"))
+# Load the model
+model = ECAPA_gender.from_pretrained("JaesungHuh/ecapa-gender")
+# model.load_state_dict(torch.load("gender_classifier.model", map_location="cpu"))
 
 model.eval()
 

model.py

@@ -1,14 +1,18 @@
+import math
+from typing import Optional
+
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
 
 import torchaudio
 from torchaudio.functional import resample
-import math
+
+from huggingface_hub import PyTorchModelHubMixin
 
 
 class SEModule(nn.Module):
-    def __init__(self, channels, bottleneck=128):
+    def __init__(self, channels : int , bottleneck : int = 128) -> None:
         super(SEModule, self).__init__()
         self.se = nn.Sequential(
             nn.AdaptiveAvgPool1d(1),
@@ -19,13 +23,13 @@ class SEModule(nn.Module):
             nn.Sigmoid(),
             )
 
-    def forward(self, input):
+    def forward(self, input : torch.Tensor) -> torch.Tensor:
         x = self.se(input)
         return input * x
 
-class Bottle2neck(nn.Module):
 
-    def __init__(self, inplanes, planes, kernel_size=None, dilation=None, scale = 8):
+class Bottle2neck(nn.Module):
+    def __init__(self, inplanes : int, planes : int, kernel_size : Optional[int] = None, dilation : Optional[int] = None, scale : int = 8) -> None:
         super(Bottle2neck, self).__init__()
         width       = int(math.floor(planes / scale))
         self.conv1  = nn.Conv1d(inplanes, width*scale, kernel_size=1)
@@ -45,7 +49,7 @@ class Bottle2neck(nn.Module):
         self.width  = width
         self.se     = SEModule(planes)
 
-    def forward(self, x):
+    def forward(self, x : torch.Tensor) -> torch.Tensor:
         residual = x
         out = self.conv1(x)
         out = self.relu(out)
@@ -73,34 +77,12 @@ class Bottle2neck(nn.Module):
         out = self.se(out)
         out += residual
         return out 
+    
 
-class PreEmphasis(torch.nn.Module):
-
-    def __init__(self, coef: float = 0.97):
-        super().__init__()
-        self.coef = coef
-        self.register_buffer(
-            'flipped_filter', torch.FloatTensor([-self.coef, 1.]).unsqueeze(0).unsqueeze(0)
-        )
-
-    def forward(self, input: torch.tensor) -> torch.tensor:
-        input = input.unsqueeze(1)
-        input = F.pad(input, (1, 0), 'reflect')
-        return F.conv1d(input, self.flipped_filter).squeeze(1)
-
-
-class ECAPA_gender(nn.Module):
-    def __init__(self, config):
+class ECAPA_gender(nn.Module, PyTorchModelHubMixin):
+    def __init__(self, C : int = 1024):
         super(ECAPA_gender, self).__init__()
-        self.config = config
-        C = config["C"]
-
-        self.torchfbank = torch.nn.Sequential(
-            PreEmphasis(),            
-            torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, \
-                                                 f_min = 20, f_max = 7600, window_fn=torch.hamming_window, n_mels=80),
-            )
-
+        self.C = C
         self.conv1  = nn.Conv1d(80, C, kernel_size=5, stride=1, padding=2)
         self.relu   = nn.ReLU()
         self.bn1    = nn.BatchNorm1d(C)
@@ -121,13 +103,26 @@ class ECAPA_gender(nn.Module):
         self.fc6 = nn.Linear(3072, 192)
         self.bn6 = nn.BatchNorm1d(192)
         self.fc7 = nn.Linear(192, 2)
-        self.pred2gender = {0 : 'Male', 1 : 'Female'}
-
-    def forward(self, x):
-        with torch.no_grad():
-            x = self.torchfbank(x)+1e-6
-            x = x.log()   
-            x = x - torch.mean(x, dim=-1, keepdim=True)
+        self.pred2gender = {0 : 'male', 1 : 'female'}
+
+    def logtorchfbank(self, x : torch.Tensor) -> torch.Tensor:
+        # Preemphasis
+        flipped_filter = torch.FloatTensor([-0.97, 1.]).unsqueeze(0).unsqueeze(0)
+        x = x.unsqueeze(1)
+        x = F.pad(x, (1, 0), 'reflect')
+        x = F.conv1d(x, flipped_filter).squeeze(1)
+
+        # Melspectrogram
+        x = torchaudio.transforms.MelSpectrogram(sample_rate=16000, n_fft=512, win_length=400, hop_length=160, \
+                                                 f_min = 20, f_max = 7600, window_fn=torch.hamming_window, n_mels=80)(x) + 1e-6
+        
+        # Log and normalize
+        x = x.log()   
+        x = x - torch.mean(x, dim=-1, keepdim=True)
+        return x
+    
+    def forward(self, x : torch.Tensor) -> torch.Tensor:
+        x = self.logtorchfbank(x)
 
         x = self.conv1(x)
         x = self.relu(x)
@@ -158,17 +153,16 @@ class ECAPA_gender(nn.Module):
         
         return x
     
-    def load_audio(self, path):
+    def load_audio(self, path : str) -> torch.Tensor:
         audio, sr = torchaudio.load(path)
         if sr != 16000:
             audio = resample(audio, sr, 16000)
         return audio
     
-    def predict(self, audio):
+    def predict(self, audio : torch.Tensor) -> torch.Tensor:
         audio = self.load_audio(audio)
         self.eval()
         with torch.no_grad():
             output = self.forward(audio)
             _, pred = output.max(1)
-        return self.pred2gender[pred.item()]
-    
+        return self.pred2gender[pred.item()]

requirements.txt

@@ -1,2 +1,3 @@
 torch
-torchaudio
+torchaudio
+pysoundfile