import logging
import math
import os
import subprocess
from io import BytesIO
import librosa
import numpy as np
import torch
import torch.nn.functional as F
import torchaudio
from audio_separator.separator import Separator
from einops import rearrange
from funasr.download.download_from_hub import download_model
from funasr.models.emotion2vec.model import Emotion2vec
from transformers import Wav2Vec2FeatureExtractor
from memo.models.emotion_classifier import AudioEmotionClassifierModel
from memo.models.wav2vec import Wav2VecModel
logger = logging.getLogger(__name__)
def resample_audio(input_audio_file: str, output_audio_file: str, sample_rate: int = 16000):
p = subprocess.Popen(
[
"ffmpeg",
"-y",
"-v",
"error",
"-i",
input_audio_file,
"-ar",
str(sample_rate),
output_audio_file,
]
)
ret = p.wait()
assert ret == 0, f"Resample audio failed! Input: {input_audio_file}, Output: {output_audio_file}"
return output_audio_file
@torch.no_grad()
def preprocess_audio(
wav_path: str,
fps: int,
wav2vec_model: str,
vocal_separator_model: str = None,
cache_dir: str = "",
device: str = "cuda",
sample_rate: int = 16000,
num_generated_frames_per_clip: int = -1,
):
"""
Preprocess the audio file and extract audio embeddings.
Args:
wav_path (str): Path to the input audio file.
fps (int): Frames per second for the audio processing.
wav2vec_model (str): Path to the pretrained Wav2Vec model.
vocal_separator_model (str, optional): Path to the vocal separator model. Defaults to None.
cache_dir (str, optional): Directory for cached files. Defaults to "".
device (str, optional): Device to use ('cuda' or 'cpu'). Defaults to "cuda".
sample_rate (int, optional): Sampling rate for audio processing. Defaults to 16000.
num_generated_frames_per_clip (int, optional): Number of generated frames per clip for padding. Defaults to -1.
Returns:
tuple: A tuple containing:
- audio_emb (torch.Tensor): The processed audio embeddings.
- audio_length (int): The length of the audio in frames.
"""
# Initialize Wav2Vec model
audio_encoder = Wav2VecModel.from_pretrained(wav2vec_model).to(device=device)
audio_encoder.feature_extractor._freeze_parameters()
# Initialize Wav2Vec feature extractor
wav2vec_feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_model)
# Initialize vocal separator if provided
vocal_separator = None
if vocal_separator_model is not None:
os.makedirs(cache_dir, exist_ok=True)
vocal_separator = Separator(
output_dir=cache_dir,
output_single_stem="vocals",
model_file_dir=os.path.dirname(vocal_separator_model),
)
vocal_separator.load_model(os.path.basename(vocal_separator_model))
#vocal_separator.load_model("UVR-MDX-NET-Inst_HQ_3.onnx")
assert vocal_separator.model_instance is not None, "Failed to load audio separation model."
# Perform vocal separation if applicable
if vocal_separator is not None:
outputs = vocal_separator.separate(wav_path)
assert len(outputs) > 0, "Audio separation failed."
vocal_audio_file = outputs[0]
vocal_audio_name, _ = os.path.splitext(vocal_audio_file)
vocal_audio_file = os.path.join(vocal_separator.output_dir, vocal_audio_file)
vocal_audio_file = resample_audio(
vocal_audio_file,
os.path.join(vocal_separator.output_dir, f"{vocal_audio_name}-16k.wav"),
sample_rate,
)
else:
vocal_audio_file = wav_path
# Load audio and extract Wav2Vec features
speech_array, sampling_rate = librosa.load(vocal_audio_file, sr=sample_rate)
audio_feature = np.squeeze(wav2vec_feature_extractor(speech_array, sampling_rate=sampling_rate).input_values)
audio_length = math.ceil(len(audio_feature) / sample_rate * fps)
audio_feature = torch.from_numpy(audio_feature).float().to(device=device)
# Pad audio features to match the required length
if num_generated_frames_per_clip > 0 and audio_length % num_generated_frames_per_clip != 0:
audio_feature = torch.nn.functional.pad(
audio_feature,
(
0,
(num_generated_frames_per_clip - audio_length % num_generated_frames_per_clip) * (sample_rate // fps),
),
"constant",
0.0,
)
audio_length += num_generated_frames_per_clip - audio_length % num_generated_frames_per_clip
audio_feature = audio_feature.unsqueeze(0)
# Extract audio embeddings
with torch.no_grad():
embeddings = audio_encoder(audio_feature, seq_len=audio_length, output_hidden_states=True)
assert len(embeddings) > 0, "Failed to extract audio embeddings."
audio_emb = torch.stack(embeddings.hidden_states[1:], dim=1).squeeze(0)
audio_emb = rearrange(audio_emb, "b s d -> s b d")
# Concatenate embeddings with surrounding frames
audio_emb = audio_emb.cpu().detach()
concatenated_tensors = []
for i in range(audio_emb.shape[0]):
vectors_to_concat = [audio_emb[max(min(i + j, audio_emb.shape[0] - 1), 0)] for j in range(-2, 3)]
concatenated_tensors.append(torch.stack(vectors_to_concat, dim=0))
audio_emb = torch.stack(concatenated_tensors, dim=0)
if vocal_separator is not None:
del vocal_separator
del audio_encoder
return audio_emb, audio_length
@torch.no_grad()
def extract_audio_emotion_labels(
model: str,
wav_path: str,
emotion2vec_model: str,
audio_length: int,
sample_rate: int = 16000,
device: str = "cuda",
):
"""
Extract audio emotion labels from an audio file.
Args:
model (str): Path to the MEMO model.
wav_path (str): Path to the input audio file.
emotion2vec_model (str): Path to the Emotion2vec model.
audio_length (int): Target length for interpolated emotion labels.
sample_rate (int, optional): Sample rate of the input audio. Default is 16000.
device (str, optional): Device to use ('cuda' or 'cpu'). Default is "cuda".
Returns:
torch.Tensor: Processed emotion labels with shape matching the target audio length.
"""
# Load models
logger.info("Downloading emotion2vec models from modelscope")
kwargs = download_model(model=emotion2vec_model)
kwargs["tokenizer"] = None
kwargs["input_size"] = None
kwargs["frontend"] = None
emotion_model = Emotion2vec(**kwargs, vocab_size=-1).to(device)
init_param = kwargs.get("init_param", None)
load_emotion2vec_model(
model=emotion_model,
path=init_param,
ignore_init_mismatch=kwargs.get("ignore_init_mismatch", True),
oss_bucket=kwargs.get("oss_bucket", None),
scope_map=kwargs.get("scope_map", []),
)
emotion_model.eval()
classifier = AudioEmotionClassifierModel.from_pretrained(
model,
subfolder="misc/audio_emotion_classifier",
use_safetensors=True,
).to(device=device)
classifier.eval()
# Load audio
wav, sr = torchaudio.load(wav_path)
if sr != sample_rate:
wav = torchaudio.functional.resample(wav, sr, sample_rate)
wav = wav.view(-1) if wav.dim() == 1 else wav[0].view(-1)
emotion_labels = torch.full_like(wav, -1, dtype=torch.int32)
def extract_emotion(x):
"""
Extract emotion for a given audio segment.
"""
x = x.to(device=device)
x = F.layer_norm(x, x.shape).view(1, -1)
feats = emotion_model.extract_features(x)
x = feats["x"].mean(dim=1) # average across frames
x = classifier(x)
x = torch.softmax(x, dim=-1)
return torch.argmax(x, dim=-1)
# Process start, middle, and end segments
start_label = extract_emotion(wav[: sample_rate * 2]).item()
emotion_labels[:sample_rate] = start_label
for i in range(sample_rate, len(wav) - sample_rate, sample_rate):
mid_wav = wav[i - sample_rate : i - sample_rate + sample_rate * 3]
mid_label = extract_emotion(mid_wav).item()
emotion_labels[i : i + sample_rate] = mid_label
end_label = extract_emotion(wav[-sample_rate * 2 :]).item()
emotion_labels[-sample_rate:] = end_label
# Interpolate to match the target audio length
emotion_labels = emotion_labels.unsqueeze(0).unsqueeze(0).float()
emotion_labels = F.interpolate(emotion_labels, size=audio_length, mode="nearest").squeeze(0).squeeze(0).int()
num_emotion_classes = classifier.num_emotion_classes
del emotion_model
del classifier
return emotion_labels, num_emotion_classes
def load_emotion2vec_model(
path: str,
model: torch.nn.Module,
ignore_init_mismatch: bool = True,
map_location: str = "cpu",
oss_bucket=None,
scope_map=[],
):
obj = model
dst_state = obj.state_dict()
logger.debug(f"Emotion2vec checkpoint: {path}")
if oss_bucket is None:
src_state = torch.load(path, map_location=map_location)
else:
buffer = BytesIO(oss_bucket.get_object(path).read())
src_state = torch.load(buffer, map_location=map_location)
src_state = src_state["state_dict"] if "state_dict" in src_state else src_state
src_state = src_state["model_state_dict"] if "model_state_dict" in src_state else src_state
src_state = src_state["model"] if "model" in src_state else src_state
if isinstance(scope_map, str):
scope_map = scope_map.split(",")
scope_map += ["module.", "None"]
for k in dst_state.keys():
k_src = k
if scope_map is not None:
src_prefix = ""
dst_prefix = ""
for i in range(0, len(scope_map), 2):
src_prefix = scope_map[i] if scope_map[i].lower() != "none" else ""
dst_prefix = scope_map[i + 1] if scope_map[i + 1].lower() != "none" else ""
if dst_prefix == "" and (src_prefix + k) in src_state.keys():
k_src = src_prefix + k
if not k_src.startswith("module."):
logger.debug(f"init param, map: {k} from {k_src} in ckpt")
elif k.startswith(dst_prefix) and k.replace(dst_prefix, src_prefix, 1) in src_state.keys():
k_src = k.replace(dst_prefix, src_prefix, 1)
if not k_src.startswith("module."):
logger.debug(f"init param, map: {k} from {k_src} in ckpt")
if k_src in src_state.keys():
if ignore_init_mismatch and dst_state[k].shape != src_state[k_src].shape:
logger.debug(
f"ignore_init_mismatch:{ignore_init_mismatch}, dst: {k, dst_state[k].shape}, src: {k_src, src_state[k_src].shape}"
)
else:
dst_state[k] = src_state[k_src]
else:
logger.debug(f"Warning, miss key in ckpt: {k}, mapped: {k_src}")
obj.load_state_dict(dst_state, strict=True)