Skip to content

vad API

maai.models.vad

VadGPT

Bases: Module

Voice Activity Detection (VAD) model.

Detects current voice activity for each of the two input channels. The architecture is the VAP model without the projection head: the per-channel and cross-channel transformers are followed by a single va_classifier that is shared between both channels.

Unlike a per-channel energy based VAD, the cross-channel attention lets the model decide which speaker is actually talking even when the voice of the other speaker leaks into the channel.

Source code in src/maai/models/vad.py 
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
class VadGPT(nn.Module):
    """Voice Activity Detection (VAD) model.

    Detects *current* voice activity for each of the two input channels.
    The architecture is the VAP model without the projection head: the
    per-channel and cross-channel transformers are followed by a single
    ``va_classifier`` that is shared between both channels.

    Unlike a per-channel energy based VAD, the cross-channel attention lets
    the model decide which speaker is actually talking even when the voice
    of the other speaker leaks into the channel.
    """

    def __init__(self, conf: Optional[VapConfig] = None):
        """Initialize the VadGPT model.

        Args:
            conf (Optional[VapConfig]): Configuration object for the model.
                If None, default VapConfig is used.
        """
        super().__init__()
        if conf is None:
            conf = VapConfig()
        self.conf = conf
        self.sample_rate = conf.sample_rate
        self.frame_hz = conf.frame_hz

        # Single channel
        self.ar_channel = GPT(
            dim=conf.dim,
            dff_k=3,
            num_layers=conf.channel_layers,
            num_heads=conf.num_heads,
            dropout=conf.dropout,
            context_limit=conf.context_limit,
        )

        # Cross channel
        # The VAD model only uses the per-channel outputs (x1 / x2), so the
        # combinator is not built (it has no weights in the checkpoints).
        self.ar = GPTStereo(
            dim=conf.dim,
            dff_k=3,
            num_layers=conf.cross_layers,
            num_heads=conf.num_heads,
            dropout=conf.dropout,
            context_limit=conf.context_limit,
            combinator=False,
        )

        # Outputs
        # Voice activity objective -> x1, x2 -> logits -> BCE
        self.va_classifier = nn.Linear(conf.dim, 1)

    def load_encoder(self, cpc_model):
        """Load and build the audio encoders for both speakers.

        Args:
            cpc_model: Pre-trained CPC model to be used as feature extractor.
        """
        self.encoder1 = build_audio_encoder(self.conf, cpc_model=cpc_model)
        self.encoder1 = self.encoder1.eval()
        self.encoder2 = build_audio_encoder(self.conf, cpc_model=cpc_model)
        self.encoder2 = self.encoder2.eval()

        encoder_dim = getattr(self.encoder1, "output_dim", self.conf.dim)
        if encoder_dim != self.conf.dim:
            self.decrease_dimension = nn.Linear(encoder_dim, self.conf.dim)

        if self.conf.freeze_encoder == 1:
            print('freeze encoder')
            self.encoder1.freeze()
            self.encoder2.freeze()

    def encode_audio(self, audio1: torch.Tensor, audio2: torch.Tensor) -> Tuple[Tensor, Tensor]:
        """Encode the raw audio inputs into feature representations.

        Args:
            audio1 (torch.Tensor): Audio waveform for speaker 1.
            audio2 (torch.Tensor): Audio waveform for speaker 2.

        Returns:
            Tuple[Tensor, Tensor]: Encoded features for speaker 1 and speaker 2.
        """
        x1 = self.encoder1(audio1)  # speaker 1
        x2 = self.encoder2(audio2)  # speaker 2

        if hasattr(self, "decrease_dimension"):
            x1 = torch.relu(self.decrease_dimension(x1))
            x2 = torch.relu(self.decrease_dimension(x2))

        return x1, x2

    def forward(
        self,
        x1: Tensor,
        x2: Tensor,
        cache: Optional[dict] = None,
        return_all_frames: bool = False,
    ) -> Tuple[dict, dict]:
        """
        Forward pass for the VadGPT model.

        Args:
            x1 (Tensor): Input audio embedded tensor for speaker 1.
            x2 (Tensor): Input audio embedded tensor for speaker 2.
            cache (dict, optional): Cache of past keys/values.

        Returns:
            Tuple[dict, dict]: Model outputs and updated cache.
        """

        if cache is None:
            cache = {}

        o1 = self.ar_channel(x1, past_kv=cache.get("ar1"))
        o2 = self.ar_channel(x2, past_kv=cache.get("ar2"))
        out = self.ar(
            o1["x"],
            o2["x"],
            past_kv1=cache.get("cross1"),
            past_kv2=cache.get("cross2"),
            past_kv1_c=cache.get("cross1_c"),
            past_kv2_c=cache.get("cross2_c"),
        )

        new_cache = {
            "ar1": (o1["past_k"], o1["past_v"]),
            "ar2": (o2["past_k"], o2["past_v"]),
            "cross1": (out["past_k1"], out["past_v1"]),
            "cross2": (out["past_k2"], out["past_v2"]),
            "cross1_c": (out["past_k1_c"], out["past_v1_c"]),
            "cross2_c": (out["past_k2_c"], out["past_v2_c"]),
        }

        # Outputs
        vad1 = self.va_classifier(out["x1"])
        vad2 = self.va_classifier(out["x2"])

        # Get back to the CPU
        vad1_all = vad1.sigmoid().to("cpu").tolist()[0]
        vad2_all = vad2.sigmoid().to("cpu").tolist()[0]

        frames = [
            {"vad": [vad1_all[t][0], vad2_all[t][0]]}
            for t in range(x1.shape[1])
        ]
        ret = frames if return_all_frames else frames[-1]

        return ret, new_cache

__init__(conf=None)

Initialize the VadGPT model.

Parameters:

Name Type Description Default
conf Optional[VapConfig]

Configuration object for the model. If None, default VapConfig is used.

 None
Source code in src/maai/models/vad.py 
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
def __init__(self, conf: Optional[VapConfig] = None):
    """Initialize the VadGPT model.

    Args:
        conf (Optional[VapConfig]): Configuration object for the model.
            If None, default VapConfig is used.
    """
    super().__init__()
    if conf is None:
        conf = VapConfig()
    self.conf = conf
    self.sample_rate = conf.sample_rate
    self.frame_hz = conf.frame_hz

    # Single channel
    self.ar_channel = GPT(
        dim=conf.dim,
        dff_k=3,
        num_layers=conf.channel_layers,
        num_heads=conf.num_heads,
        dropout=conf.dropout,
        context_limit=conf.context_limit,
    )

    # Cross channel
    # The VAD model only uses the per-channel outputs (x1 / x2), so the
    # combinator is not built (it has no weights in the checkpoints).
    self.ar = GPTStereo(
        dim=conf.dim,
        dff_k=3,
        num_layers=conf.cross_layers,
        num_heads=conf.num_heads,
        dropout=conf.dropout,
        context_limit=conf.context_limit,
        combinator=False,
    )

    # Outputs
    # Voice activity objective -> x1, x2 -> logits -> BCE
    self.va_classifier = nn.Linear(conf.dim, 1)

encode_audio(audio1, audio2)

Encode the raw audio inputs into feature representations.

Parameters:

Name Type Description Default
audio1 Tensor

Audio waveform for speaker 1.

 required
audio2 Tensor

Audio waveform for speaker 2.

 required

Returns:

Type Description
Tuple[Tensor, Tensor]

Tuple[Tensor, Tensor]: Encoded features for speaker 1 and speaker 2.
Source code in src/maai/models/vad.py 
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
def encode_audio(self, audio1: torch.Tensor, audio2: torch.Tensor) -> Tuple[Tensor, Tensor]:
    """Encode the raw audio inputs into feature representations.

    Args:
        audio1 (torch.Tensor): Audio waveform for speaker 1.
        audio2 (torch.Tensor): Audio waveform for speaker 2.

    Returns:
        Tuple[Tensor, Tensor]: Encoded features for speaker 1 and speaker 2.
    """
    x1 = self.encoder1(audio1)  # speaker 1
    x2 = self.encoder2(audio2)  # speaker 2

    if hasattr(self, "decrease_dimension"):
        x1 = torch.relu(self.decrease_dimension(x1))
        x2 = torch.relu(self.decrease_dimension(x2))

    return x1, x2

forward(x1, x2, cache=None, return_all_frames=False)

Forward pass for the VadGPT model.

Parameters:

Name Type Description Default
x1 Tensor

Input audio embedded tensor for speaker 1.

 required
x2 Tensor

Input audio embedded tensor for speaker 2.

 required
cache dict

Cache of past keys/values.

 None

Returns:

Type Description
Tuple[dict, dict]

Tuple[dict, dict]: Model outputs and updated cache.
Source code in src/maai/models/vad.py 
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
def forward(
    self,
    x1: Tensor,
    x2: Tensor,
    cache: Optional[dict] = None,
    return_all_frames: bool = False,
) -> Tuple[dict, dict]:
    """
    Forward pass for the VadGPT model.

    Args:
        x1 (Tensor): Input audio embedded tensor for speaker 1.
        x2 (Tensor): Input audio embedded tensor for speaker 2.
        cache (dict, optional): Cache of past keys/values.

    Returns:
        Tuple[dict, dict]: Model outputs and updated cache.
    """

    if cache is None:
        cache = {}

    o1 = self.ar_channel(x1, past_kv=cache.get("ar1"))
    o2 = self.ar_channel(x2, past_kv=cache.get("ar2"))
    out = self.ar(
        o1["x"],
        o2["x"],
        past_kv1=cache.get("cross1"),
        past_kv2=cache.get("cross2"),
        past_kv1_c=cache.get("cross1_c"),
        past_kv2_c=cache.get("cross2_c"),
    )

    new_cache = {
        "ar1": (o1["past_k"], o1["past_v"]),
        "ar2": (o2["past_k"], o2["past_v"]),
        "cross1": (out["past_k1"], out["past_v1"]),
        "cross2": (out["past_k2"], out["past_v2"]),
        "cross1_c": (out["past_k1_c"], out["past_v1_c"]),
        "cross2_c": (out["past_k2_c"], out["past_v2_c"]),
    }

    # Outputs
    vad1 = self.va_classifier(out["x1"])
    vad2 = self.va_classifier(out["x2"])

    # Get back to the CPU
    vad1_all = vad1.sigmoid().to("cpu").tolist()[0]
    vad2_all = vad2.sigmoid().to("cpu").tolist()[0]

    frames = [
        {"vad": [vad1_all[t][0], vad2_all[t][0]]}
        for t in range(x1.shape[1])
    ]
    ret = frames if return_all_frames else frames[-1]

    return ret, new_cache

load_encoder(cpc_model)

Load and build the audio encoders for both speakers.

Parameters:

Name Type Description Default
cpc_model

Pre-trained CPC model to be used as feature extractor.

 required
Source code in src/maai/models/vad.py 
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
def load_encoder(self, cpc_model):
    """Load and build the audio encoders for both speakers.

    Args:
        cpc_model: Pre-trained CPC model to be used as feature extractor.
    """
    self.encoder1 = build_audio_encoder(self.conf, cpc_model=cpc_model)
    self.encoder1 = self.encoder1.eval()
    self.encoder2 = build_audio_encoder(self.conf, cpc_model=cpc_model)
    self.encoder2 = self.encoder2.eval()

    encoder_dim = getattr(self.encoder1, "output_dim", self.conf.dim)
    if encoder_dim != self.conf.dim:
        self.decrease_dimension = nn.Linear(encoder_dim, self.conf.dim)

    if self.conf.freeze_encoder == 1:
        print('freeze encoder')
        self.encoder1.freeze()
        self.encoder2.freeze()

VadGPT_mono

Bases: VadGPT

Single-channel (mono) variant of the VAD model.

The underlying model is identical to :class:VadGPT (the same pretrained vad checkpoints are loaded), but the interface is mono: only channel 1 carries speech while channel 2 is fed silence (MaaiInput.Zero). The output vad is therefore a single float for channel 1 instead of a value per speaker.

Source code in src/maai/models/vad.py 
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
class VadGPT_mono(VadGPT):
    """Single-channel (mono) variant of the VAD model.

    The underlying model is identical to :class:`VadGPT` (the same
    pretrained ``vad`` checkpoints are loaded), but the interface is
    mono: only channel 1 carries speech while channel 2 is fed silence
    (``MaaiInput.Zero``). The output ``vad`` is therefore a single float
    for channel 1 instead of a value per speaker.
    """

    def forward(
        self,
        x1: Tensor,
        x2: Tensor,
        cache: Optional[dict] = None,
        return_all_frames: bool = False,
    ) -> Tuple[dict, dict]:
        """
        Forward pass for the mono VAD model.

        Args:
            x1 (Tensor): Input audio embedded tensor for the single speaker.
            x2 (Tensor): Input audio embedded tensor for the silent channel.
            cache (dict, optional): Cache of past keys/values.

        Returns:
            Tuple[dict, dict]: Model outputs (scalar ``vad``) and updated cache.
        """
        ret, new_cache = super().forward(
            x1, x2, cache, return_all_frames=return_all_frames
        )

        if return_all_frames:
            for frame in ret:
                frame["vad"] = frame["vad"][0]
            return ret, new_cache

        # ch1 のみを返す
        ret["vad"] = ret["vad"][0]

        return ret, new_cache

forward(x1, x2, cache=None, return_all_frames=False)

Forward pass for the mono VAD model.

Parameters:

Name Type Description Default
x1 Tensor

Input audio embedded tensor for the single speaker.

 required
x2 Tensor

Input audio embedded tensor for the silent channel.

 required
cache dict

Cache of past keys/values.

 None

Returns:

Type Description
Tuple[dict, dict]

Tuple[dict, dict]: Model outputs (scalar vad) and updated cache.
Source code in src/maai/models/vad.py 
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def forward(
    self,
    x1: Tensor,
    x2: Tensor,
    cache: Optional[dict] = None,
    return_all_frames: bool = False,
) -> Tuple[dict, dict]:
    """
    Forward pass for the mono VAD model.

    Args:
        x1 (Tensor): Input audio embedded tensor for the single speaker.
        x2 (Tensor): Input audio embedded tensor for the silent channel.
        cache (dict, optional): Cache of past keys/values.

    Returns:
        Tuple[dict, dict]: Model outputs (scalar ``vad``) and updated cache.
    """
    ret, new_cache = super().forward(
        x1, x2, cache, return_all_frames=return_all_frames
    )

    if return_all_frames:
        for frame in ret:
            frame["vad"] = frame["vad"][0]
        return ret, new_cache

    # ch1 のみを返す
    ret["vad"] = ret["vad"][0]

    return ret, new_cache