app.py

from time import perf_counter
import glob
import os
from dotenv import load_dotenv
import shutil
from pathlib import Path
import gradio as gr
import spaces
import torch
import torchaudio
import transformers
import numpy as np
import noisereduce as nr
from pyannote.audio import Pipeline
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
from speechbrain.inference.separation import SepformerSeparation
from pydub import AudioSegment

# Disables TF32 for to prevent lower accuracy in pyannote
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False

# Passing .env file as environment variable on machine
load_dotenv()

# Hugging Face authentication token grab from .env
token = os.getenv("HF_TOKEN")

# GPU processing availability check
device = "cuda" if torch.cuda.is_available() else "cpu"

# Name of mixture folder
mixtures = Path(os.getcwd())/"mixtures"

# Name of separated file folder
sources = Path(os.getcwd()) / "sources"


# Path checking for removing source directory every run
def reset_dir():
    remove_dir = [mixtures,sources]
    # Getting current source directory
    for dir in remove_dir:
        if os.path.exists(dir):
            try:
                # Remove directorty recursively
                shutil.rmtree(sources)
            except:
                raise gr.Error("Reset Failed")
    
    gr.Info("System Reset",title="Notice")

    # To clear all Gradio Blocks
    return None, None, None, None, None


# Formatting the transcription log to Gradio chat display
def chatFormat(transcription_log):
    chat = []
    for i in transcription_log:
        name = i.get("speaker")
        text = i.get("text")
        start = i.get("start")
        end = i.get("end")

        # Title format for gradio chatbot block
        title_format = f"**({name}) [{start}s - {end}s]**"
        role = "user" if name == "speaker_2" else "assistant"

        chat.append({"role": role, "content":text, "metadata":{"title":title_format}})

    return chat


@spaces.GPU
def process_audio(mixture, progress=gr.Progress()):

    # Inference Time Processing
    start_time = perf_counter()

    progress(0, desc="Starting")

    # Gradio processing alert
    gr.Info("Processing Has Started", title="Notice")

    # Creates mixture directory
    Path.mkdir(mixtures, exist_ok=True)

    # Creates separated directory
    Path.mkdir(sources, exist_ok=True)

# ---------------------------------------------------------
# 1. INITIALIZE MODELS
# ---------------------------------------------------------

    progress(0.12, "Loading RE-SepFormer Model...")
    print("Intializing models...")

    # Loading Resepformer Model
    try:
        separator = SepformerSeparation.from_hparams(
            source=f"Jahmori-R/resepformer-librimix2spk",
            hparams_file="hyperparams.yaml",
            run_opts={"device": device},
        )
        print("RE-SepFormer model successfully loaded")
    except Exception as e:
        raise RuntimeError("Failed to load RE-SepFormer model") from e

    progress(0.25, "Loading Pyannote Model...")
    # Loading Pyannote Model

    try:
        diarize = Pipeline.from_pretrained(
            "pyannote/speaker-diarization-community-1", token=token
        )
        # Assigning GPU processing to model
        diarize.to(torch.device("cuda"))
        print("Pyannote model successfully loaded")
    except Exception as e:
        raise RuntimeError("Failed to load Pyannote model") from e

    progress(0.38, "Loading Whisper Model...")

    # Loading Whisper Model
    transformers.logging.set_verbosity_error()  # Prevents WhisperAI logging in cmd line
    try:
        transcription_model_id = "openai/whisper-large-v3-turbo"
        torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
        model = AutoModelForSpeechSeq2Seq.from_pretrained(
            transcription_model_id,
            dtype=torch_dtype,
            low_cpu_mem_usage=True,
            use_safetensors=True,
        )

        processor = AutoProcessor.from_pretrained(transcription_model_id)

        transcribe = pipeline(
            "automatic-speech-recognition",
            model=model,
            tokenizer=processor.tokenizer,
            feature_extractor=processor.feature_extractor,
            dtype=torch_dtype,
            device="cuda",
        )

        print("Whisper model successfully loaded")
    except Exception as e:
        raise RuntimeError("Failed to load Whisper model") from e

# ---------------------------------------------------------
# LOAD & SEPARATE AUDIO
# ---------------------------------------------------------

    # Calling ReSepformer model onto audio source
    assert separator is not None, "Failed to load the Sepformer model."

    progress(0.40,"Denoising Mixture...")

    #Denoising mixture
    sr, mixture_np = mixture

    if len(mixture_np.shape) > 1:
        # Swaps channels location with sample rate in numpy array
        if mixture_np.shape[0] > mixture_np.shape[1]:
            mixture_np = mixture_np.T

    # If it has more than 2 channels take the first one
    if len(mixture_np.shape) > 1 and mixture_np.shape[0] > 2:
        mixture_np = mixture_np[0] 
    
    # Converting numpy to float32 for torchaudio 
    if mixture_np.dtype != np.float32:

        # Converting integer value to float between -1.0 and 1.0
        if np.issubdtype(mixture_np.dtype, np.integer):
            mixture_np = mixture_np.astype(np.float32) / np.iinfo(mixture_np.dtype).max
        else:
            mixture_np = mixture_np.astype(np.float32)

    # Apply noise gate onto waveform
    denoised_mixture_np = nr.reduce_noise(y=mixture_np,sr=sr, prop_decrease=0.5,n_fft=512)

    # Convert numpy data back to torch for saving
    denoised_mixture = torch.from_numpy(denoised_mixture_np)

    # Saving new denoised mixture
    torchaudio.save(f"{mixtures}/mixture.wav",denoised_mixture,sr)

    progress(0.50,"Separating Mixtures...")

    # Running mixed audio onto separation model
    print("Separating Audio Mixture")

    est_sources = separator.separate_file(f"{mixtures}/mixture.wav")
    print("Mixture separated")
    
    # Saving separated sources to directory
    torchaudio.save(
        f"{sources}/speaker_1.wav", est_sources[:, :, 0].detach().cpu(), 16000
    )
    torchaudio.save(
        f"{sources}/speaker_2.wav", est_sources[:, :, 1].detach().cpu(), 16000
    )


# ---------------------------------------------------------
# VAD + TRANSCRIPTION ON SOURCES
# ---------------------------------------------------------

    # Dictionary to store speaker information
    unfiltered_log = [[], []]

    progress(0.62, "Sorting Sources...")

    # Storing both speakers
    file_paths = sorted(
        # Glob finds all similar files
        glob.glob(f"{sources}/speaker_*.wav")
    )

    spk_1 = file_paths[0]
    spk_2 = file_paths[1]

    progress(0.75, "VAD + Transcribing...")

    # Individual Speaker Processing
    for i, path in enumerate(file_paths):
        speaker_key = f"speaker_{i + 1}"

        # Diarize speaker audio for time segments
        print(f"Diarizing {speaker_key}")
        diarization = diarize(path, num_speakers=1)

        # Minimum timestamp duration in seconds
        minimum_duration = 0.8

        # Storing speaker voice activity timestamps to dictionary
        for turn, _ in diarization.exclusive_speaker_diarization:
            duration = turn.end - turn.start
            if duration <  minimum_duration:
                continue
        
            # Adding timestamp to list
            unfiltered_log[i].append(
                {
                    "start": float(round(turn.start, 2)),
                    "end": float(round(turn.end, 2)),
                    "speaker": speaker_key,
                }
            )

        # Creating audio segment from detected voice activity
        audio_seg = AudioSegment.from_file(path)

        print("Transcribing Chunks...")
        # Transcribing each voice activity detected by Pyannote
        for segment in unfiltered_log[i]:
            # Converting segment timestamps to milliseconds for pydub segmentation
            start_ms = segment["start"] * 1000
            end_ms = segment["end"] * 1000
            chunk = audio_seg[start_ms:end_ms]

            # Converting pydub AudioSegment to numpy array for whisper model
            samples = np.array(chunk.get_array_of_samples())
            samples = samples.astype(np.float32) / 32768.0

            # Whisper Configutation
            kwargs = {
                "language": "english",  # Language transcription
            }
            # Generating transcription given chunk segment
            transcription = transcribe(samples, generate_kwargs=kwargs)

            # Appending each segment transcription to dictionary
            segment.update(transcription)

    progress(0.88, "Sorting Speaker Logs....")
    # Combined both speaker logs
    combined_log = unfiltered_log[0] + unfiltered_log[1]

    # Sorting by start time of timestamps
    filtered_log = sorted(combined_log, key=lambda x: x["start"])
    chat = chatFormat(filtered_log)

    progress(1.00, "Task Finished")
    print("Task Finished")

    end_time = perf_counter()

    # Inference time calculation
    inference_time = f"{end_time - start_time:.2f} seconds"
    return (spk_1, spk_2, chat, inference_time)

# ---------------------------------------------------------
# GRADIO INTERFACE CUSTOMIZATION
# ---------------------------------------------------------

# Custom css for gradio interface
css = """

#centered-row {
    justify-content: center;
    display: flex;
}
"""

# Gradio layout configuration
with gr.Blocks() as demo:
    gr.Markdown("PureVoiceAI")
    with gr.Row(variant="panel"):
        audio_input = gr.Audio(
            type="numpy", label="Input Audio Mixture | 2 Speakers"
        )
    with gr.Row(elem_id="centered-row"):
        process_button = gr.Button(
            "Process",
            variant="primary",
            scale=0,
        )
        reset_button = gr.Button("Restart", variant="stop", interactive=False, scale=0)
    with gr.Row(variant="panel", visible=False) as source_row:
        with gr.Column():
            speaker_1_output = gr.Audio(label="Speaker 1", interactive=False)
            speaker_2_output = gr.Audio(label="Speaker 2", interactive=False)
        with gr.Column():
            inference_time_log = gr.Textbox(label="Inference Time", interactive=False)

    with gr.Row(variant="panel") as log_row:
        logs = gr.Chatbot(label="Conversation Log",layout="bubble")

    # Process Button Functionality
    process_button.click(
        fn=process_audio,
        inputs=audio_input,
        outputs=[speaker_1_output, speaker_2_output, logs, inference_time_log],
    ).then(
        fn=lambda: (
            gr.Row(visible=True),
            gr.Button(interactive=False),
            gr.Button(interactive=True),
        ),
        inputs=None,
        outputs=[source_row, process_button, reset_button],
    )

    # Reset Button Functionality
    reset_button.click(
        fn=reset_dir,
        outputs=[
            audio_input,
            speaker_1_output,
            speaker_2_output,
            inference_time_log,
            logs,
        ],
    ).then(
        fn=lambda: (
            gr.Row(visible=False),
            gr.Button(interactive=True),
            gr.Button(interactive=False),
        ),
        inputs=None,
        outputs=[source_row, process_button, reset_button],
    )

if __name__ == "__main__":
    demo.launch(theme=gr.Theme.from_hub("pseudolab/huggingface-korea-theme"), css=css)