Skip to content

Ensure fusion gating is identity when audio disabled #14

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Robbyswimmer wants to merge 1 commit into
base: master
Choose a base branch
from
Open

Ensure fusion gating is identity when audio disabled #14

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
99 changes: 86 additions & 13 deletions safe/models/audio_encoders.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,6 @@
import torch
import torch.nn as nn
import numpy as np
import torch.nn as nn
import librosa
from typing import Optional, Union, List, Tuple
import laion_clap
Expand All @@ -21,11 +21,14 @@ def __init__(
max_length: float = 10.0 # seconds
):
super().__init__()

self.sample_rate = sample_rate
self.max_length = max_length
self.max_samples = int(sample_rate * max_length)

self.debug_logging = False
self._waveform_log_limit = 5
self._waveform_logs_emitted = 0

# Load CLAP model (suppress verbose output)
import logging
import sys
Expand All @@ -51,7 +54,40 @@ def __init__(

# Get audio embedding dimension
self.audio_embed_dim = 512 # CLAP audio embedding dimension


def set_debug_logging(self, enabled: bool, max_waveform_logs: int = 5) -> None:
"""Enable or disable verbose waveform statistics logging."""

self.debug_logging = bool(enabled)
self._waveform_log_limit = int(max(0, max_waveform_logs))
self._waveform_logs_emitted = 0

# ------------------------------------------------------------------
def _log_waveform_stats(self, audio_data: np.ndarray, source: str) -> None:
if not self.debug_logging:
return
if self._waveform_logs_emitted >= self._waveform_log_limit:
return

flattened = audio_data.astype(np.float32).ravel()
if flattened.size == 0:
print(f"[AudioDebug] {source}: empty waveform")
self._waveform_logs_emitted += 1
return

max_val = float(np.max(flattened))
min_val = float(np.min(flattened))
mean_val = float(np.mean(flattened))
mean_abs = float(np.mean(np.abs(flattened)))
zero_fraction = float(np.mean(np.isclose(flattened, 0.0)))

print(
f"[AudioDebug] {source}: max={max_val:.6f} min={min_val:.6f} "
f"mean={mean_val:.6f} mean|x|={mean_abs:.6f} zero_frac={zero_fraction:.3f}",
flush=True,
)
self._waveform_logs_emitted += 1

def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> torch.Tensor:
"""
Preprocess audio to the format expected by CLAP.
Expand All @@ -68,20 +104,22 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> torch
elif isinstance(audio, np.ndarray):
audio_data = audio
elif isinstance(audio, torch.Tensor):
audio_data = audio.cpu().numpy()
audio_data = audio.detach().cpu().numpy()
else:
raise ValueError(f"Unsupported audio type: {type(audio)}")

# Ensure correct sample rate
if len(audio_data.shape) > 1:
audio_data = audio_data.mean(axis=0) # Convert to mono

# Truncate or pad to max_length
if len(audio_data) > self.max_samples:
audio_data = audio_data[:self.max_samples]
else:
audio_data = np.pad(audio_data, (0, self.max_samples - len(audio_data)))


self._log_waveform_stats(audio_data, "CLAP.preprocess_audio")

return torch.from_numpy(audio_data).float().unsqueeze(0) # (1, max_samples)

def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> torch.Tensor:
Expand All @@ -97,7 +135,7 @@ def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> to
if isinstance(audio, list):
# Process batch of audio files/arrays
processed_audio = []
for a in audio:
for idx, a in enumerate(audio):
processed = self.preprocess_audio(a)
processed_audio.append(processed)
audio_batch = torch.stack(processed_audio) # (batch_size, 1, max_samples)
Expand Down Expand Up @@ -125,7 +163,7 @@ def forward(self, audio: Union[torch.Tensor, List[str], List[np.ndarray]]) -> to
audio_numpy = audio_batch.detach().cpu().numpy()
else:
audio_numpy = audio_batch

audio_embeddings = self.model.get_audio_embedding_from_data(
x=audio_numpy,
use_tensor=False
Expand Down Expand Up @@ -157,11 +195,14 @@ def __init__(
max_length: float = 30.0 # seconds
):
super().__init__()

self.sample_rate = sample_rate
self.max_length = max_length
self.extract_transcript = extract_transcript
self.max_samples = int(sample_rate * max_length)
self.debug_logging = False
self._waveform_log_limit = 5
self._waveform_logs_emitted = 0

# Load Whisper model
self.model = whisper.load_model(model_name)
Expand All @@ -174,6 +215,36 @@ def __init__(

# Get embedding dimension from Whisper encoder
self.audio_embed_dim = self.model.dims.n_audio_state # 512 for small, 768 for base, etc.

def set_debug_logging(self, enabled: bool, max_waveform_logs: int = 5) -> None:
self.debug_logging = bool(enabled)
self._waveform_log_limit = int(max(0, max_waveform_logs))
self._waveform_logs_emitted = 0

def _log_waveform_stats(self, audio_data: np.ndarray, source: str) -> None:
if not self.debug_logging:
return
if self._waveform_logs_emitted >= self._waveform_log_limit:
return

flattened = audio_data.astype(np.float32).ravel()
if flattened.size == 0:
print(f"[AudioDebug] {source}: empty waveform", flush=True)
self._waveform_logs_emitted += 1
return

max_val = float(np.max(flattened))
min_val = float(np.min(flattened))
mean_val = float(np.mean(flattened))
mean_abs = float(np.mean(np.abs(flattened)))
zero_fraction = float(np.mean(np.isclose(flattened, 0.0)))

print(
f"[AudioDebug] {source}: max={max_val:.6f} min={min_val:.6f} "
f"mean={mean_val:.6f} mean|x|={mean_abs:.6f} zero_frac={zero_fraction:.3f}",
flush=True,
)
self._waveform_logs_emitted += 1

def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.ndarray:
"""
Expand All @@ -188,7 +259,7 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.nd
if isinstance(audio, str):
audio_data, sr = librosa.load(audio, sr=self.sample_rate)
elif isinstance(audio, torch.Tensor):
audio_data = audio.cpu().numpy()
audio_data = audio.detach().cpu().numpy()
else:
audio_data = audio

Expand All @@ -201,7 +272,9 @@ def preprocess_audio(self, audio: Union[torch.Tensor, np.ndarray, str]) -> np.nd
audio_data = audio_data[:self.max_samples]
else:
audio_data = np.pad(audio_data, (0, self.max_samples - len(audio_data)))


self._log_waveform_stats(audio_data, "Whisper.preprocess_audio")

return audio_data

def forward(
Expand Down
Loading