Spaces:
Build error
Build error
File size: 4,777 Bytes
f36e52e 7dccca6 f36e52e 7dccca6 f36e52e e30002c f36e52e e30002c f36e52e e30002c f36e52e |
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 |
import torch
import whisper
import numpy as np
import torchaudio as ta
import gradio as gr
import spaces
from model_utils import get_processor, get_model, get_whisper_model_small, get_device
from config import SAMPLING_RATE, CHUNK_LENGTH_S
import subprocess
import subprocess
import torchaudio as ta
def resample_with_ffmpeg(input_file, output_file, target_sr=16000):
command = [
'ffmpeg', '-i', input_file, '-ar', str(target_sr), output_file
]
subprocess.run(command, check=True)
@spaces.GPU
def detect_language(audio):
whisper_model = get_whisper_model_small()
# Save the input audio to a temporary file
ta.save("input_audio.wav", torch.tensor(audio[1]).unsqueeze(0), audio[0])
# Resample if necessary using ffmpeg
if audio[0] != SAMPLING_RATE:
resample_with_ffmpeg("input_audio.wav", "resampled_audio.wav", target_sr=SAMPLING_RATE)
audio_tensor, _ = ta.load("resampled_audio.wav")
else:
audio_tensor = torch.tensor(audio[1]).float()
# Ensure the audio is in the correct shape (mono)
if audio_tensor.dim() == 2:
audio_tensor = audio_tensor.mean(dim=0)
# Use Whisper's preprocessing
audio_tensor = whisper.pad_or_trim(audio_tensor)
print(f"Audio length after pad/trim: {audio_tensor.shape[-1] / SAMPLING_RATE} seconds")
mel = whisper.log_mel_spectrogram(audio_tensor).to(whisper_model.device)
# Detect language
_, probs = whisper_model.detect_language(mel)
detected_lang = max(probs, key=probs.get)
print(f"Audio shape: {audio_tensor.shape}")
print(f"Mel spectrogram shape: {mel.shape}")
print(f"Detected language: {detected_lang}")
print("Language probabilities:", probs)
return detected_lang
@spaces.GPU
def process_long_audio(audio, task="transcribe", language=None):
if audio[0] != SAMPLING_RATE:
# Save the input audio to a file for ffmpeg processing
ta.save("input_audio_1.wav", torch.tensor(audio[1]).unsqueeze(0), audio[0])
# Resample using ffmpeg
try:
resample_with_ffmpeg("input_audio_1.wav", "resampled_audio_2.wav", target_sr=SAMPLING_RATE)
except subprocess.CalledProcessError as e:
print(f"ffmpeg failed: {e.stderr}")
raise e
waveform, _ = ta.load("resampled_audio_2.wav")
else:
waveform = torch.tensor(audio[1]).float()
# Ensure the audio is in the correct shape (mono)
if waveform.dim() == 2:
waveform = waveform.mean(dim=0)
print(f"Waveform shape after processing: {waveform.shape}")
if waveform.numel() == 0:
raise ValueError("Waveform is empty. Please check the input audio file.")
input_length = waveform.shape[0] # Since waveform is 1D, access the length with shape[0]
chunk_length = int(CHUNK_LENGTH_S * SAMPLING_RATE)
# Corrected slicing for 1D tensor
chunks = [waveform[i:i + chunk_length] for i in range(0, input_length, chunk_length)]
# Initialize the processor
processor = get_processor()
model = get_model()
device = get_device()
results = []
for chunk in chunks:
input_features = processor(chunk, sampling_rate=SAMPLING_RATE, return_tensors="pt").input_features.to(device)
with torch.no_grad():
if task == "translate":
forced_decoder_ids = processor.get_decoder_prompt_ids(language=language, task="translate")
generated_ids = model.generate(input_features, forced_decoder_ids=forced_decoder_ids)
else:
generated_ids = model.generate(input_features)
transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)
results.extend(transcription)
# Clear GPU cache
torch.cuda.empty_cache()
return " ".join(results)
@spaces.GPU
def process_audio(audio):
if audio is None:
return "No file uploaded", "", ""
detected_lang = detect_language(audio)
transcription = process_long_audio(audio, task="transcribe")
translation = process_long_audio(audio, task="translate", language=detected_lang)
return detected_lang, transcription, translation
# Gradio interface
iface = gr.Interface(
fn=process_audio,
inputs=gr.Audio(),
outputs=[
gr.Textbox(label="Detected Language"),
gr.Textbox(label="Transcription", lines=5),
gr.Textbox(label="Translation", lines=5)
],
title="Audio Transcription and Translation",
description="Upload an audio file to detect its language, transcribe, and translate it.",
allow_flagging="never",
css=".output-textbox { font-family: 'Noto Sans Devanagari', sans-serif; font-size: 18px; }"
)
if __name__ == "__main__":
iface.launch()
|