import gradio as gr
from transformers import AutoProcessor, AutoModelForCTC
import torch
import numpy as np
import soundfile as sf
import librosa

# Load the model and processor
processor = AutoProcessor.from_pretrained("sazzadul/Shrutimala_Bangla_ASR")
model = AutoModelForCTC.from_pretrained("sazzadul/Shrutimala_Bangla_ASR")

def transcribe_audio_file(audio):
    speech_array, sampling_rate = sf.read(audio)

    # Resample the audio to 16000 Hz if necessary
    if sampling_rate != 16000:
        speech_array = librosa.resample(speech_array, orig_sr=sampling_rate, target_sr=16000)

    return process_audio(speech_array)

def process_audio(audio_array):
    # Preprocess the audio
    inputs = processor(audio_array, sampling_rate=16000, return_tensors="pt")

    # Perform inference
    with torch.no_grad():
        logits = model(**inputs).logits

    # Take the argmax over the logits to get the predicted tokens
    predicted_ids = torch.argmax(logits, dim=-1)

    # Decode the predicted tokens into text
    transcription = processor.batch_decode(predicted_ids)[0]

    return transcription

# Create a Gradio interface for file upload only
iface = gr.Interface(
    fn=transcribe_audio_file,
    inputs=gr.Audio(type="filepath"),
    outputs="text",
    title="File Upload Wav2Vec-BERT Transcription",
    description="Transcribe audio from an uploaded file using a Wav2Vec-BERT model."
)

# Launch the interface
iface.launch()

# import gradio as gr
# from transformers import AutoProcessor, AutoModelForCTC
# import torch
# import numpy as np
# import soundfile as sf
# import librosa
# import noisereduce as nr
# import logging

# # Setup logging
# logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')

# # Load the model and processor
# logging.info("Loading model and processor...")
# processor = AutoProcessor.from_pretrained("sazzadul/Shrutimala_Bangla_ASR")
# model = AutoModelForCTC.from_pretrained("sazzadul/Shrutimala_Bangla_ASR")

# def optimize_audio(audio):
#     logging.info("Audio processing started...")

#     # Load audio file
#     speech_array, sampling_rate = sf.read(audio)
#     logging.info(f"Loaded audio with sampling rate: {sampling_rate}")

#     # Resample the audio to 16000 Hz if necessary
#     if sampling_rate != 16000:
#         logging.info("Resampling audio to 16000 Hz...")
#         speech_array = librosa.resample(speech_array, orig_sr=sampling_rate, target_sr=16000)

#     # Apply noise reduction
#     logging.info("Applying noise reduction...")
#     reduced_noise = nr.reduce_noise(y=speech_array, sr=16000)

#     # Normalize the audio
#     logging.info("Normalizing audio...")
#     enhanced_audio = np.int16(reduced_noise / np.max(np.abs(reduced_noise)) * 32767)
#     enhanced_audio_path = "enhanced_audio.wav"
#     sf.write(enhanced_audio_path, enhanced_audio, 16000)

#     logging.info("Audio optimization completed.")
#     return enhanced_audio_path

# def transcribe_audio(enhanced_audio_path):
#     logging.info("Transcription process started...")

#     # Load the enhanced audio for transcription
#     speech_array, sampling_rate = sf.read(enhanced_audio_path)

#     # Preprocess the audio
#     inputs = processor(speech_array, sampling_rate=16000, return_tensors="pt")
#     logging.info("Audio preprocessed for model input.")

#     # Perform inference
#     with torch.no_grad():
#         logits = model(**inputs).logits
#     logging.info("Inference completed.")

#     # Take the argmax over the logits to get the predicted tokens
#     predicted_ids = torch.argmax(logits, dim=-1)

#     # Decode the predicted tokens into text
#     transcription = processor.batch_decode(predicted_ids)[0]
#     logging.info("Transcription completed.")

#     return transcription

# def process_pipeline(audio):
#     logging.info("Starting processing pipeline...")

#     # Step 1: Optimize Audio
#     enhanced_audio_path = optimize_audio(audio)
    
#     # Step 2: Transcribe Enhanced Audio
#     transcription = transcribe_audio(enhanced_audio_path)
    
#     logging.info("Processing pipeline completed.")
#     return transcription, enhanced_audio_path

# # Create a Gradio interface
# iface = gr.Interface(
#     fn=process_pipeline,
#     inputs=gr.Audio(type="filepath"),
#     outputs=[
#         "text",                   # Transcription output
#         gr.Audio(type="filepath") # Enhanced audio preview
#     ],
#     title="Bangla ASR with Optimized Audio",
#     description="Upload an audio file, preview the enhanced version, and get the transcription."
# )

# # Launch the interface
# iface.launch()