1577-2 / backend /utils.py
jts-ai-team's picture
Update backend/utils.py
6425857 verified
Raw History Blame
4.29 kB
import numpy as np
import librosa
import io
import os
import warnings
import tempfile
from pydub import AudioSegment
from dotenv import load_dotenv
from fastrtc import get_cloudflare_turn_credentials_async, get_cloudflare_turn_credentials
try:
import torch
except ModuleNotFoundError:
torch = None # type: ignore
warnings.filterwarnings("ignore")
# load_dotenv(override = True)
# --- Device Configuration ---
def get_device():
"""Gets the best available device for PyTorch."""
if torch is None:
return "cpu"
if torch.cuda.is_available():
return "cuda"
elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
return "mps"
else:
return "cpu"
device = get_device()
print(f"Using device: {device}")
# --- Cloud Credentials ---
async def get_async_credentials():
"""Asynchronously fetches Cloudflare TURN credentials."""
return await get_cloudflare_turn_credentials_async(hf_token=os.getenv('HF_TOKEN'))
def get_sync_credentials(ttl=360_000):
"""Synchronously fetches Cloudflare TURN credentials."""
return get_cloudflare_turn_credentials(ttl=ttl)
def setup_gcp_credentials():
"""Sets up Google Cloud credentials from an environment variable."""
gcp_service_account_json_str = os.getenv("GCP_SERVICE_ACCOUNT_JSON")
if gcp_service_account_json_str:
try:
# Create a temporary file to store the credentials
with tempfile.NamedTemporaryFile(mode='w', delete=False, suffix=".json") as temp_file:
temp_file.write(gcp_service_account_json_str)
gcp_credential_path = temp_file.name # Get the path to the temporary file
# Set the environment variable that Google Cloud libraries expect
os.environ["GOOGLE_APPLICATION_CREDENTIALS"] = gcp_credential_path
print(f"Google Cloud credentials set from secret to: {gcp_credential_path}")
except Exception as e:
print(f"Error setting up Google Cloud credentials: {e}")
else:
print("Warning: GCP_SERVICE_ACCOUNT_JSON secret not found. Google Cloud services may fail.")
# if gcp_service_account_json_str:
# print("GCP service account JSON loaded from environment variable.")
# else:
# print("Warning: GCP_SERVICE_ACCOUNT_JSON is not set; Google Cloud clients may fail.")
# return gcp_service_account_json_str
# --- Audio Processing ---
def audiosegment_to_numpy(audio, target_sample_rate=16000):
samples = np.array(audio.get_array_of_samples(), dtype=np.float32)
if audio.channels > 1:
samples = samples.reshape((-1, audio.channels)).mean(axis=1)
samples /= np.iinfo(audio.array_type).max
if audio.frame_rate != target_sample_rate:
samples = librosa.resample(samples, orig_sr=audio.frame_rate, target_sr=target_sample_rate)
return samples
def preprocess_audio(audio, target_channels=1, target_frame_rate=16000):
"""
Preprocess the audio using pydub AudioSegment by setting the number of channels and frame rate.
Args:
audio (tuple): A tuple (sample_rate, audio_array) where audio_array is a NumPy array.
target_channels (int): Desired number of channels (default is 1 for mono).
target_frame_rate (int): Desired frame rate (default is 16000 Hz).
Returns:
np.ndarray: The processed audio as a NumPy array.
"""
sample_rate, audio_array = audio
target_frame_rate = sample_rate
audio_array_int16 = audio_array.astype(np.int16)
audio_bytes = audio_array_int16.tobytes()
audio_io = io.BytesIO(audio_bytes)
segment = AudioSegment.from_raw(audio_io, sample_width=2, frame_rate=sample_rate, channels=1)
segment = segment.set_channels(target_channels)
segment = segment.set_frame_rate(target_frame_rate)
return audiosegment_to_numpy(segment)
# --- Conversation Utilities ---
def is_valid_turn(turn):
"""Return True if turn is a valid dict with non-empty 'role' and 'content' strings."""
return (
isinstance(turn, dict)
and "role" in turn
and "content" in turn
and isinstance(turn["role"], str)
and isinstance(turn["content"], str)
and turn["role"].strip() != ""
and turn["content"].strip() != ""
)