File size: 9,029 Bytes
eca55dc | 1 2 3 4 5 6 7 8 9 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 160 161 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 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 | import glob
import os
import sys
import torch
import torch.nn as nn
from omegaconf import OmegaConf
from safetensors.torch import load_file
# Add audio-embeddings to path dynamically
# We assume audio-embeddings is a sibling directory to xares-llm or provided via env var
# Prioritize absolute path if known, otherwise relative
POSSIBLE_PATHS = [
# "/media/ltuncay/Shared-4TB/dev/audio-embeddings",
os.path.abspath(os.path.join(os.path.dirname(__file__), "audio-embeddings")),
# os.path.abspath(os.path.join(os.getcwd(), "../audio-embeddings")),
]
AUDIO_EMBEDDINGS_PATH = None
for p in POSSIBLE_PATHS:
if os.path.exists(p):
AUDIO_EMBEDDINGS_PATH = p
break
if AUDIO_EMBEDDINGS_PATH:
if AUDIO_EMBEDDINGS_PATH not in sys.path:
sys.path.append(AUDIO_EMBEDDINGS_PATH)
print(f"Added {AUDIO_EMBEDDINGS_PATH} to sys.path")
else:
print(
"Warning: audio-embeddings path not found. Imports may fail if not installed in environment."
)
try:
from src.models.best_rq2_module import BestRQ2Module
except ImportError as e:
raise ImportError(
f"Could not import src.models.best_rq2_module. Ensure audio-embeddings is correctly located or installed. Error: {e}"
)
class BestRQ2Encoder(nn.Module):
def __init__(self, checkpoint_path=None, model_config_path=None, **kwargs):
super().__init__()
base_path = os.path.dirname(__file__)
model_config_path = os.path.join(base_path, "config.yaml")
checkpoint_path = os.path.join(base_path, "BEST-RQ-2.safetensors")
if not os.path.exists(model_config_path):
raise FileNotFoundError(f"Config not found at {model_config_path}")
if not checkpoint_path or not os.path.exists(checkpoint_path):
raise FileNotFoundError(f"Checkpoint not found at {checkpoint_path}")
print(f"Loading BestRQ2 config from {model_config_path}")
cfg = OmegaConf.load(model_config_path)
print(f"Loading BestRQ2 checkpoint from {checkpoint_path}")
# Reconstruct model args from config
model_cfg = cfg.model
net_cfg = model_cfg.net
# Instantiate model
# Note: BestRQ2Module inherits from LightningModule
self.module = BestRQ2Module(
optimizer=None, # Not needed for inference
net=net_cfg,
warmup_pct=model_cfg.get("warmup_pct", 0.1),
final_lr_ratio=model_cfg.get("final_lr_ratio", 0.001),
spectrogram_adjustment_mode=model_cfg.get(
"spectrogram_adjustment_mode", "pad"
),
codebook_dim=model_cfg.get("codebook_dim", 16),
vocab_size=model_cfg.get("vocab_size", 8192),
criterion=None,
)
# Load weights
try:
state_dict = load_file(checkpoint_path)
except Exception as e:
print(f"Error loading safetensors: {e}. Trying torch.load...")
state_dict = torch.load(checkpoint_path, map_location="cpu")
if "state_dict" in state_dict:
state_dict = state_dict["state_dict"]
# Handle 'module.' prefix if present in checkpoint vs model
# Usually LightningModules save with state_dict keys matching model attributes.
# But sometimes they might be wrapped.
# We will try loading strict=False and inspect.
missing, unexpected = self.module.load_state_dict(state_dict, strict=False)
if missing:
# Check if prefixes match
# If all missing keys start with something common, or if state_dict has prefixes
print(f"Warning: {len(missing)} keys missing during loading.")
# print(missing[:5])
if unexpected:
print(f"Warning: {len(unexpected)} keys unexpected during loading.")
self.module.eval()
self.output_dim = net_cfg.encoder.embed_dim
# Extract dynamic parameters for length handling
try:
# 1. Sample Rate & Hop Length (from Spectrogram)
# BestRQ2Module -> Spectrogram -> MelSpectrogram -> hop_length
self.sample_rate = self.module.spectrogram.mel_spec.sample_rate
self.hop_length = self.module.spectrogram.mel_spec.hop_length
# 2. Patch Size (Time dimension)
# BestRQ2Module -> PatchEmbed -> patch_size (H, W) -> W is time
self.patch_size_time = self.module.patch_embed.patch_size[1]
# 3. Max Input Frames (Time dimension)
# BestRQ2Module -> PatchEmbed -> img_size (H, W) -> W is time frames
self.max_frames = self.module.patch_embed.img_size[1]
# Calculations
# Minimum samples required to get at least 1 patch width in spectrogram
# We need T_spec >= patch_size_time
# T_spec = T_samples // hop_length (roughly)
# So T_samples >= patch_size_time * hop_length
self.min_samples = self.patch_size_time * self.hop_length
# Chunk size: The maximum audio length the model's positional embeddings can handle
# T_samples_max = max_frames * hop_length
self.chunk_samples = self.max_frames * self.hop_length
print(
f"BestRQ2Encoder constraints: Min Samples={self.min_samples}, Chunk Samples={self.chunk_samples}"
)
except Exception as e:
print(f"Warning: Could not extract dynamic length constraints: {e}")
print("Falling back to safe defaults (1s min, 10s chunk)")
self.min_samples = 16000
self.chunk_samples = 16000 * 10
def _forward_chunk(self, audio_chunk: torch.Tensor) -> torch.Tensor:
"""Helper to process a single time-chunk of audio."""
# Determine target device from the spectrogram window (safest for STFT)
try:
target_device = self.module.spectrogram.mel_spec.spectrogram.window.device
except AttributeError:
if hasattr(self.module.spectrogram.mel_spec, "window"):
target_device = self.module.spectrogram.mel_spec.window.device
else:
target_device = self.module.device
if audio_chunk.device != target_device:
audio_chunk = audio_chunk.to(target_device)
# BestRQ2Module expects [B, C, T]
if audio_chunk.ndim == 2:
audio_chunk = audio_chunk.unsqueeze(1) # [B, 1, T]
# _process_audio returns (patches, grid_size)
patches, grid_size = self.module._process_audio(audio_chunk)
# Create Dummy Mask (all False = keep all)
B, N, D = patches.shape
mask = torch.zeros((B, N), dtype=torch.bool, device=patches.device)
# Compute encoder
encoder_out = self.module.compute_encoder(patches, mask, grid_size)
return encoder_out
def forward(
self, audio: torch.Tensor, audio_attention_mask=None
) -> tuple[torch.Tensor, torch.Tensor | None]:
# audio: [B, T]
if audio.ndim == 1:
audio = audio.unsqueeze(0)
B, T = audio.shape
# 1. Handle Short Audio (Whole Batch)
if T < self.min_samples:
pad_amt = self.min_samples - T
audio = torch.nn.functional.pad(audio, (0, pad_amt))
T = self.min_samples # Update T
# 2. Sequential Chunking
if T <= self.chunk_samples:
# Single chunk processing
return self._forward_chunk(audio), None
else:
# Split into chunks of max length
chunks = torch.split(audio, self.chunk_samples, dim=1)
outputs = []
for chunk in chunks:
# Handle potentially short last chunk
chunk_len = chunk.shape[1]
if chunk_len < self.min_samples:
pad_amt = self.min_samples - chunk_len
chunk = torch.nn.functional.pad(chunk, (0, pad_amt))
# Process
out_chunk = self._forward_chunk(chunk)
# If we padded the last chunk solely to meet min_samples,
# should we slice? BestRQ2 output is patches.
# 1 patch covers `min_samples`.
# If original was < 1 patch, we produced 1 patch.
# We can't slice sub-patch. We just return the 1 patch.
outputs.append(out_chunk)
# Concatenate along sequence dimension (dim=1)
final_output = torch.cat(outputs, dim=1)
return final_output, None
if __name__ == "__main__":
try:
mdl = BestRQ2Encoder()
print("Model initialized successfully")
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
mdl.module.to(device)
x = torch.randn(1, 160000).to(device)
y, _ = mdl(x)
print(f"Output shape: {y.shape}")
except Exception as e:
print(f"Error testing model: {e}")
import traceback
traceback.print_exc()
|