transcription_diarization.py

import os
import torch
import math
from moviepy.editor import VideoFileClip
from pyannote.audio import Pipeline
from transformers import AutoModelForSpeechSeq2Seq, AutoProcessor, pipeline
import librosa
import datetime
from collections import defaultdict
import numpy as np
import spaces

class LazyDiarizationPipeline:
    def __init__(self):
        self.pipeline = None

    @spaces.GPU(duration=100)
    def get_pipeline(self, diarization_access_token):
        if self.pipeline is None:
            self.pipeline = Pipeline.from_pretrained("pyannote/speaker-diarization-3.1", use_auth_token=diarization_access_token)
            self.pipeline = self.pipeline.to(torch.device("cuda"))
        return self.pipeline

class LazyTranscriptionPipeline:
    def __init__(self):
        self.model = None
        self.processor = None
        self.pipe = None

    @spaces.GPU(duration=100)
    def get_pipeline(self, language):
        if self.pipe is None:
            model_id = "openai/whisper-large-v3"
            self.model = AutoModelForSpeechSeq2Seq.from_pretrained(
                model_id, torch_dtype=torch.float16, low_cpu_mem_usage=True, use_safetensors=True
            )
            self.model.to(torch.device("cuda"))
            self.processor = AutoProcessor.from_pretrained(model_id)
            self.pipe = pipeline(
                "automatic-speech-recognition",
                model=self.model,
                tokenizer=self.processor.tokenizer,
                feature_extractor=self.processor.feature_extractor,
                max_new_tokens=128,
                chunk_length_s=30,
                batch_size=2,
                return_timestamps=True,
                torch_dtype=torch.float16,
                device=torch.device("cuda"),
                generate_kwargs={"language": language}
            )
        return self.pipe

lazy_diarization_pipeline = LazyDiarizationPipeline()
lazy_transcription_pipeline = LazyTranscriptionPipeline()

def extract_audio(video_path, audio_path):
    video = VideoFileClip(video_path)
    audio = video.audio
    audio.write_audiofile(audio_path, codec='pcm_s16le', fps=16000)

def format_timestamp(seconds):
    return str(datetime.timedelta(seconds=seconds)).split('.')[0]

@spaces.GPU(duration=100)
def transcribe_audio(audio_path, language):
    pipe = lazy_transcription_pipeline.get_pipeline(language)

    audio, sr = librosa.load(audio_path, sr=16000)
    duration = len(audio) / sr
    n_chunks = math.ceil(duration / 30)
    transcription_txt = ""
    transcription_chunks = []

    for i in range(n_chunks):
        start = i * 30 * sr
        end = min((i + 1) * 30 * sr, len(audio))
        audio_chunk = audio[start:end]
        audio_chunk = (audio_chunk * 32767).astype(np.float32)

        result = pipe(audio_chunk)
        transcription_txt += result["text"]
        for chunk in result["chunks"]:
            start_time, end_time = chunk["timestamp"]
            transcription_chunks.append({
                "start": start_time + i * 30,
                "end": end_time + i * 30,
                "text": chunk["text"]
            })

        print(f"Transcription Progress: {int(((i + 1) / n_chunks) * 100)}%")

    return transcription_txt, transcription_chunks

def create_combined_srt(transcription_chunks, diarization, output_path):
    speaker_segments = []
    speaker_durations = defaultdict(float)

    for segment, _, speaker in diarization.itertracks(yield_label=True):
        speaker_durations[speaker] += segment.end - segment.start
        speaker_segments.append((segment.start, segment.end, speaker))

    sorted_speakers = sorted(speaker_durations.items(), key=lambda x: x[1], reverse=True)
    
    speaker_map = {}
    for i, (speaker, _) in enumerate(sorted_speakers, start=1):
        speaker_map[speaker] = f"Speaker {i}"

    with open(output_path, 'w', encoding='utf-8') as srt_file:
        for i, chunk in enumerate(transcription_chunks, 1):
            start_time, end_time = chunk["start"], chunk["end"]
            text = chunk["text"]

            current_speaker = "Unknown"
            for seg_start, seg_end, speaker in speaker_segments:
                if seg_start <= start_time < seg_end:
                    current_speaker = speaker_map[speaker]
                    break

            start_str = format_timestamp(start_time).split('.')[0].lstrip('0')
            end_str = format_timestamp(end_time).split('.')[0].lstrip('0')

            srt_file.write(f"{i}\n")
            srt_file.write(f"{current_speaker}\n time: ({start_str} --> {end_str})\n text: {text}\n\n")

    with open(output_path, 'a', encoding='utf-8') as srt_file:
        for i, (speaker, duration) in enumerate(sorted_speakers[:2], start=1):
            duration_str = format_timestamp(duration).split('.')[0].lstrip('0')
            srt_file.write(f"Speaker {i} (originally {speaker}): total duration {duration_str}\n")

@spaces.GPU(duration=100)
def process_video(video_path, diarization_access_token, language):
    base_name = os.path.splitext(video_path)[0]
    audio_path = f"{base_name}.wav"
    extract_audio(video_path, audio_path)

    print("Performing diarization...")
    pipeline = lazy_diarization_pipeline.get_pipeline(diarization_access_token)
    diarization = pipeline(audio_path)
    print("Diarization complete.")

    print("Performing transcription...")
    transcription, chunks = transcribe_audio(audio_path, language)
    print("Transcription complete.")

    combined_srt_path = f"{base_name}_combined.srt"
    create_combined_srt(chunks, diarization, combined_srt_path)
    print(f"Combined SRT file created and saved to {combined_srt_path}")

    os.remove(audio_path)

    return combined_srt_path