diff --git a/env_cpu.yml b/env_cpu.yml new file mode 100644 index 0000000..314f139 --- /dev/null +++ b/env_cpu.yml @@ -0,0 +1,19 @@ +name: notecast +channels: + - pytorch + - conda-forge + - defaults +dependencies: + - python=3.10 + - pytorch::pytorch=2.0.1 + - pytorch::torchvision=0.15.2 + - pytorch::torchaudio=2.0.2 + - transformers=4.30.0 + - accelerate=0.20.0 + - ffmpeg + - pyinstaller + - customtkinter + - isort + - mypy + - black + - tqdm \ No newline at end of file diff --git a/env_cuda.yml b/env_cuda.yml new file mode 100644 index 0000000..192b8ef --- /dev/null +++ b/env_cuda.yml @@ -0,0 +1,21 @@ +name: notecast +channels: + - pytorch + - nvidia + - conda-forge + - defaults +dependencies: + - python=3.10 + - pytorch::pytorch=2.0.1 + - pytorch::torchvision=0.15.2 + - pytorch::torchaudio=2.0.2 + - pytorch::cudatoolkit=11.8 + - transformers=4.30.0 + - accelerate=0.20.0 + - ffmpeg + - pyinstaller + - customtkinter + - isort + - mypy + - black + - tqdm \ No newline at end of file diff --git a/env_mps.yml b/env_mps.yml new file mode 100644 index 0000000..314f139 --- /dev/null +++ b/env_mps.yml @@ -0,0 +1,19 @@ +name: notecast +channels: + - pytorch + - conda-forge + - defaults +dependencies: + - python=3.10 + - pytorch::pytorch=2.0.1 + - pytorch::torchvision=0.15.2 + - pytorch::torchaudio=2.0.2 + - transformers=4.30.0 + - accelerate=0.20.0 + - ffmpeg + - pyinstaller + - customtkinter + - isort + - mypy + - black + - tqdm \ No newline at end of file diff --git a/environment.yml b/environment.yml deleted file mode 100644 index 5ed8d78..0000000 --- a/environment.yml +++ /dev/null @@ -1,16 +0,0 @@ -name: notecast -channels: - - defaults - - conda-forge -dependencies: - - torchvision - - pytorch - - torchaudio - - transformers - - accelerate - - ffmpeg - - pyinstaller - - customtkinter - - isort - - mypy - - black \ No newline at end of file diff --git a/main.py b/main.py index 51e50f3..9bcd0b0 100644 --- a/main.py +++ b/main.py @@ -2,4 +2,4 @@ from ui.ui import TranscriberApp if __name__ == "__main__": app = TranscriberApp() - app.mainloop() \ No newline at end of file + app.mainloop() diff --git a/transcription/audio_transcription.py b/transcription/audio_transcription.py index b2bee31..aff804e 100644 --- a/transcription/audio_transcription.py +++ b/transcription/audio_transcription.py @@ -1,43 +1,46 @@ -from transformers import WhisperProcessor, WhisperForConditionalGeneration +import logging +import math +import time + import torch import torchaudio -from ui.ui_log_handler import UILogHandler -from transcription.device_configuration import DeviceConfiguration -import logging -import time -import math from tqdm import tqdm +from transformers import WhisperForConditionalGeneration, WhisperProcessor + +from transcription.device_configuration import DeviceConfiguration +from ui.ui_log_handler import UILogHandler + # TODO: implement transcription with shift class AudioTranscription: - model_name = "openai/whisper-large-v2" - + model_name = "openai/whisper-large-v2" + filepath: str waveform: torch.Tensor sampling_rate: int - + chunks: list = [] batches: list = [] - chunk_size: int + chunk_size: int custom_chunk_length: int custom_batch_length: int - + device = "cuda" processor: WhisperProcessor model: WhisperForConditionalGeneration logger: logging.Logger torch_dtype: torch.dtype - + language = "ru" - + all_transcription: list = [] - + def __init__( - self, + self, filepath: str, device_configuration: DeviceConfiguration, logger: logging.Logger, - language = "ru" + language="ru", ) -> None: # TODO: add pretty docs here self.filepath = filepath @@ -50,44 +53,49 @@ class AudioTranscription: self.custom_chunk_length = device_configuration.chunk_length_s self.custom_batch_length = device_configuration.batch_size self.torch_dtype = device_configuration.torch_dtype - + self.chunks: list = [] self.batches: list = [] self.all_transcription: list = [] try: logger.info("Loading model WhisperProcessor...") self.processor = WhisperProcessor.from_pretrained(self.model_name) - + self.model = WhisperForConditionalGeneration.from_pretrained( - self.model_name, - torch_dtype=self.torch_dtype + self.model_name, torch_dtype=self.torch_dtype ).to(self.device) - + logger.info("Model loaded.") - - self.waveform, self.sampling_rate = torchaudio.load(filepath, format="mp3", backend="ffmpeg") + + self.waveform, self.sampling_rate = torchaudio.load( + filepath, format="mp3", backend="ffmpeg" + ) logger.info(f"Successfully loaded file {filepath}.") - + except Exception as e: logger.error(f"Unable to load file {self.filepath}: {e}") raise - + def _resample(self) -> None: - self.waveform = torchaudio.functional.resample(self.waveform, self.sampling_rate, 16000) - + self.waveform = torchaudio.functional.resample( + self.waveform, self.sampling_rate, 16000 + ) + def _to_mono(self): if self.waveform.shape[0] > 1: self.waveform = self.waveform.mean(dim=0, keepdim=True) self.waveform = self.waveform.squeeze(0) - + def _split_to_chunks(self, chunk_length_s: int = 30) -> None: self.logger.info(f"Splitting audio on chunks...") - + self.chunk_size = chunk_length_s * 16000 # 16kHz after resampling total_samples = self.waveform.shape[0] chunks_count = (total_samples + self.chunk_size - 1) // self.chunk_size - - self.logger.info(f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk.") + + self.logger.info( + f"File length - {total_samples / 16000:.1f} seconds, splitting on {chunks_count} chunks by {chunk_length_s} seconds per chunk." + ) self.chunks = [] for idx in tqdm(range(chunks_count)): @@ -95,46 +103,51 @@ class AudioTranscription: end = min((idx + 1) * self.chunk_size, total_samples) chunk = self.waveform[start:end].cpu().numpy().astype("float32") self.chunks.append(chunk) - + def _resplit_to_batches(self) -> None: self.logger.info(f"Splitting chunks into batches...") self.batches = [] for i in range(0, len(self.chunks), self.custom_batch_length): - batch = self.chunks[i:i + self.custom_batch_length] + batch = self.chunks[i : i + self.custom_batch_length] self.batches.append(batch) - + def _process_all_batches(self) -> None: start_time = time.time() try: self.all_transcription = [] - + for idx in tqdm(range(len(self.batches))): inputs = self.processor( - self.batches[idx], + self.batches[idx], sampling_rate=16000, - return_tensors="pt", - padding=True + return_tensors="pt", + padding=True, ) - - input_features = inputs.input_features.to(self.device).to(self.torch_dtype) - + + input_features = inputs.input_features.to(self.device).to( + self.torch_dtype + ) + with torch.no_grad(): predicted_ids = self.model.generate( input_features, language=self.language, task="transcribe", - temperature=0.0 + temperature=0.0, ) - - texts = self.processor.batch_decode(predicted_ids, skip_special_tokens=True) + + texts = self.processor.batch_decode( + predicted_ids, skip_special_tokens=True + ) self.all_transcription.extend(texts) end_time = time.time() - self.logger.info(f"Transcription completed in {end_time - start_time:.2f} seconds") + self.logger.info( + f"Transcription completed in {end_time - start_time:.2f} seconds" + ) except Exception as e: self.logger.error(f"Errors occured while processing chunks: {e}") - def transcribe_audio(self) -> str: # TODO: maybe something else, not str? self._resample() @@ -142,4 +155,4 @@ class AudioTranscription: self._split_to_chunks() self._resplit_to_batches() self._process_all_batches() - return " ".join(self.all_transcription) \ No newline at end of file + return " ".join(self.all_transcription) diff --git a/transcription/device_configuration.py b/transcription/device_configuration.py index 2f405c2..4e5edba 100644 --- a/transcription/device_configuration.py +++ b/transcription/device_configuration.py @@ -1,6 +1,8 @@ from dataclasses import dataclass + import torch + @dataclass class DeviceConfiguration: """ @@ -8,36 +10,37 @@ class DeviceConfiguration: Attributes: device (str): Type of device. Possible options: "cuda", "cpu", "mps". - + model_name (str): Whisper models. Possible options: - "openai/whisper-tiny" - "openai/whisper-small" - "openai/whisper-medium" - "openai/whisper-large" - "openai/whisper-large-v2" - + batch_size (int): Chunks in one batch. Selected for VRAM. - + chunk_length_s (int): Length of one audio chunk in seconds. Smaller -> less VRAM. - + data_type (str): custom data type of model. Variants: - torch.float16 - for GPUs - torch.float32 - for CPU / weak GPU - torch.bfloat16 - for GPUs which has BF16 support """ + device: str = "cuda" model_name: str = "openai/whisper-large-v2" batch_size: int = 16 chunk_length_s: int = 30 data_type: str = "torch.float16" - + _dtype_map = { "torch.float16": torch.float16, "torch.float32": torch.float32, - "torch.bfloat16": torch.bfloat16 + "torch.bfloat16": torch.bfloat16, } - + torch_dtype: torch.dtype = None - + def __post_init__(self): - self.torch_dtype = self._dtype_map[self.data_type] \ No newline at end of file + self.torch_dtype = self._dtype_map[self.data_type] diff --git a/transcription/torch_checker.py b/transcription/torch_checker.py index 16e2919..0a086ef 100644 --- a/transcription/torch_checker.py +++ b/transcription/torch_checker.py @@ -1,32 +1,38 @@ -import torch -from dataclasses import dataclass import logging +from dataclasses import dataclass + +import torch + def check_torch(logger: logging.Logger) -> None: logger.info("=== Checking PyTorch ===") logger.info(f"Torch version: {torch.__version__}") - - # === NVIDIA / AMD (CUDA API) === + + # NVIDIA / AMD (CUDA API) if torch.cuda.is_available(): backend = "CUDA" if torch.version.hip is not None: backend = "ROCm (AMD HIP)" - + logger.info(f"{backend} backend is available") - logger.info(f"Compiled with: CUDA {torch.version.cuda}, ROCm {torch.version.hip}") + logger.info( + f"Compiled with: CUDA {torch.version.cuda}, ROCm {torch.version.hip}" + ) logger.info(f"Number of devices: {torch.cuda.device_count()}") - + for i in range(torch.cuda.device_count()): logger.info(f"GPU {i}: {torch.cuda.get_device_name(i)}") - - # === Apple Silicon (MPS) === + + # Apple Silicon (MPS) elif hasattr(torch.backends, "mps") and torch.backends.mps.is_available(): logger.info("MPS backend is available (Apple Silicon)") - logger.info(f"MPS version: {getattr(torch.backends.mps, '__version__', 'unknown')}") + logger.info( + f"MPS version: {getattr(torch.backends.mps, '__version__', 'unknown')}" + ) logger.info("GPU: Apple Silicon (Metal)") - - # === CPU only mode === + + # CPU only mode else: logger.info("Only CPU is available") - + logger.info("=== Check completed ===") diff --git a/ui/configuration_actions.py b/ui/configuration_actions.py index 7872c8b..abbd3e2 100644 --- a/ui/configuration_actions.py +++ b/ui/configuration_actions.py @@ -1,9 +1,18 @@ from transcription.device_configuration import DeviceConfiguration + # TODO: implement saving & removing configuration def save_configuration(cfg: DeviceConfiguration): config = { "Model": cfg.model_name, "Batch Size": cfg.batch_size, - "Data Type": cfg.data_type - } \ No newline at end of file + "Data Type": cfg.data_type, + } + + +def load_config(): + pass + + +def delete_config(): + pass diff --git a/ui/ui.py b/ui/ui.py index 3c98b6c..ccb16a8 100644 --- a/ui/ui.py +++ b/ui/ui.py @@ -1,17 +1,16 @@ import os -import threading import queue +import threading import tkinter as tk -from tkinter import scrolledtext, filedialog, messagebox +from tkinter import filedialog, messagebox, scrolledtext import customtkinter as ctk import torch -from ui.ui_log_handler import setup_ui_logger -from transcription.torch_checker import check_torch -from transcription.device_configuration import DeviceConfiguration from transcription.audio_transcription import AudioTranscription - +from transcription.device_configuration import DeviceConfiguration +from transcription.torch_checker import check_torch +from ui.ui_log_handler import setup_ui_logger WINDOW_WIDTH = 900 WINDOW_HEIGHT = 650 @@ -175,7 +174,10 @@ class TranscriberApp(ctk.CTk): def _browse_input(self): path = filedialog.askopenfilename( title="Select input audio file", - filetypes=[("Audio files", "*.wav *.mp3 *.m4a *.flac"), ("All files", "*.*")] + filetypes=[ + ("Audio files", "*.wav *.mp3 *.m4a *.flac"), + ("All files", "*.*"), + ], ) if path: self.input_file_var.set(path) @@ -184,7 +186,7 @@ class TranscriberApp(ctk.CTk): path = filedialog.asksaveasfilename( title="Select output file", defaultextension=".txt", - filetypes=[("Text files", "*.txt"), ("All files", "*.*")] + filetypes=[("Text files", "*.txt"), ("All files", "*.*")], ) if path: self.output_file_var.set(path) @@ -242,4 +244,4 @@ class TranscriberApp(ctk.CTk): messagebox.showerror("Error", str(e)) finally: self.start_btn.configure(state="normal") - self.stop_btn.configure(state="disabled") \ No newline at end of file + self.stop_btn.configure(state="disabled") diff --git a/ui/ui_log_handler.py b/ui/ui_log_handler.py index 60e03d5..ec7fb10 100644 --- a/ui/ui_log_handler.py +++ b/ui/ui_log_handler.py @@ -1,5 +1,6 @@ -import tkinter as tk import logging +import tkinter as tk + class UILogHandler(logging.Handler): def __init__(self, text_widget: tk.Text): @@ -11,6 +12,7 @@ class UILogHandler(logging.Handler): self.text_widget.insert(tk.END, log_entry + "\n") self.text_widget.see(tk.END) + # TODO: maybe some tqdm here, not in console? def setup_ui_logger(text_widget: tk.Text, level=logging.INFO): logger = logging.getLogger("UI_LOGGER") @@ -21,4 +23,4 @@ def setup_ui_logger(text_widget: tk.Text, level=logging.INFO): handler.setFormatter(logging.Formatter("%(asctime)s - %(levelname)s - %(message)s")) logger.addHandler(handler) - return logger \ No newline at end of file + return logger