mirror of
https://github.com/chidiwilliams/buzz.git
synced 2026-09-05 23:55:55 +00:00
refactor(file_transcriber_queue_worker): extract run() into smaller methods to fix R0915 (#1523)
Co-authored-by: Raivis Dejus <raivisd@scandiweb.com>
This commit is contained in:
Josias Batista
GitHub
parent
951a0ef596
commit
570e03510d
2 changed files with 177 additions and 171 deletions
|
|
@ -134,75 +134,79 @@ class FileTranscriberQueueWorker(QObject):
|
|||
|
||||
logging.debug("Waiting for next transcription task")
|
||||
|
||||
# Clean up of previous run.
|
||||
self._cleanup_previous_transcriber()
|
||||
|
||||
if not self._get_next_task():
|
||||
self.is_running = False
|
||||
self.completed.emit()
|
||||
return
|
||||
|
||||
self.is_running = True
|
||||
|
||||
if self.current_task.transcription_options.extract_speech:
|
||||
status = self._setup_speech_extraction()
|
||||
if status == "error":
|
||||
self.is_running = False
|
||||
return
|
||||
|
||||
self._run_plugins()
|
||||
|
||||
logging.debug("Starting next transcription task")
|
||||
self.task_progress.emit(self.current_task, 0)
|
||||
|
||||
self._create_transcriber()
|
||||
self._setup_transcriber_thread()
|
||||
|
||||
def _cleanup_previous_transcriber(self):
|
||||
if self.current_transcriber is not None:
|
||||
self.current_transcriber.stop()
|
||||
self.current_transcriber = None
|
||||
|
||||
# Get next non-canceled task from queue
|
||||
def _get_next_task(self) -> bool:
|
||||
while True:
|
||||
self.current_task: Optional[FileTranscriptionTask] = self.tasks_queue.get()
|
||||
|
||||
# Stop listening when a "None" task is received
|
||||
self.current_task = self.tasks_queue.get()
|
||||
if self.current_task is None:
|
||||
self.is_running = False
|
||||
self.completed.emit()
|
||||
return
|
||||
|
||||
return False
|
||||
if self.current_task.uid in self.canceled_tasks:
|
||||
continue
|
||||
return True
|
||||
|
||||
break
|
||||
def _setup_speech_extraction(self) -> str:
|
||||
logging.debug("Will extract speech")
|
||||
|
||||
# Set is_running AFTER we have a valid task to process
|
||||
self.is_running = True
|
||||
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false").lower() == "true"
|
||||
if force_cpu:
|
||||
device = "cpu"
|
||||
else:
|
||||
import torch
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
|
||||
if self.current_task.transcription_options.extract_speech:
|
||||
logging.debug("Will extract speech")
|
||||
task_file_path = Path(self.current_task.file_path)
|
||||
speech_path = task_file_path.with_name(f"{task_file_path.stem}_speech.mp3")
|
||||
|
||||
# Force CPU if specified, otherwise use CUDA if available
|
||||
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false").lower() == "true"
|
||||
if force_cpu:
|
||||
device = "cpu"
|
||||
else:
|
||||
import torch
|
||||
device = "cuda" if torch.cuda.is_available() else "cpu"
|
||||
status = self._extract_speech(str(task_file_path), str(speech_path), device)
|
||||
|
||||
task_file_path = Path(self.current_task.file_path)
|
||||
speech_path = task_file_path.with_name(f"{task_file_path.stem}_speech.mp3")
|
||||
if status == "error":
|
||||
self.task_error.emit(
|
||||
self.current_task,
|
||||
_("Speech extraction failed! Check your internet connection \u2014 a model may need to be downloaded."),
|
||||
)
|
||||
elif status == "ok":
|
||||
self.speech_path = speech_path
|
||||
if not self.current_task.original_file_path:
|
||||
self.current_task.original_file_path = str(task_file_path)
|
||||
self.current_task.file_path = str(speech_path)
|
||||
|
||||
status = self._extract_speech(str(task_file_path), str(speech_path), device)
|
||||
return status
|
||||
|
||||
if status == "error":
|
||||
self.task_error.emit(
|
||||
self.current_task,
|
||||
_("Speech extraction failed! Check your internet connection — a model may need to be downloaded."),
|
||||
)
|
||||
self.is_running = False
|
||||
return
|
||||
|
||||
if status == "ok":
|
||||
self.speech_path = speech_path
|
||||
# Remember the original audio path: file_path is about to point
|
||||
# at the temporary "_speech.mp3", which is deleted once the
|
||||
# transcription completes. Plugins (e.g. the transcript resizer)
|
||||
# need the original file in their post-completion hooks.
|
||||
if not self.current_task.original_file_path:
|
||||
self.current_task.original_file_path = str(task_file_path)
|
||||
self.current_task.file_path = str(speech_path)
|
||||
# status == "no_audio": transcribe the original file as-is.
|
||||
|
||||
# Let plugins process / replace the source audio before transcription.
|
||||
# Runs on this worker thread; plugins may overwrite current_task.file_path.
|
||||
def _run_plugins(self):
|
||||
if self.plugin_manager is not None:
|
||||
try:
|
||||
self.plugin_manager.run_before_transcription(self.current_task)
|
||||
except Exception as e:
|
||||
logging.error(f"Plugin before_transcription failed: {e}", exc_info=True)
|
||||
|
||||
logging.debug("Starting next transcription task")
|
||||
self.task_progress.emit(self.current_task, 0)
|
||||
|
||||
def _create_transcriber(self):
|
||||
model_type = self.current_task.transcription_options.model.model_type
|
||||
if model_type == ModelType.OPEN_AI_WHISPER_API:
|
||||
self.current_transcriber = OpenAIWhisperAPIFileTranscriber(
|
||||
|
|
@ -218,6 +222,7 @@ class FileTranscriberQueueWorker(QObject):
|
|||
else:
|
||||
raise Exception(f"Unknown model type: {model_type}")
|
||||
|
||||
def _setup_transcriber_thread(self):
|
||||
self.current_transcriber_thread = QThread(self)
|
||||
|
||||
self.current_transcriber.moveToThread(self.current_transcriber_thread)
|
||||
|
|
@ -240,7 +245,6 @@ class FileTranscriberQueueWorker(QObject):
|
|||
|
||||
self.current_transcriber.completed.connect(self.on_task_completed)
|
||||
|
||||
# Wait for next item on the queue
|
||||
self.current_transcriber.error.connect(lambda: self._on_task_finished())
|
||||
self.current_transcriber.completed.connect(lambda: self._on_task_finished())
|
||||
|
||||
|
|
|
|||
|
|
@ -649,133 +649,135 @@ class ModelDownloader(QRunnable):
|
|||
def _register_process(self, proc: multiprocessing.Process):
|
||||
self._download_process = proc
|
||||
|
||||
def _download_whisper_cpp(self) -> None:
|
||||
if self.custom_model_url:
|
||||
url = self.custom_model_url
|
||||
file_path = get_whisper_cpp_file_path(
|
||||
size=self.model.whisper_model_size)
|
||||
self.download_model_to_path(url=url, file_path=file_path)
|
||||
return
|
||||
|
||||
repo_id = WHISPER_CPP_REPO_ID
|
||||
|
||||
if self.model.whisper_model_size == WhisperModelSize.LUMII:
|
||||
repo_id = WHISPER_CPP_LUMII_REPO_ID
|
||||
|
||||
model_name = self.model.whisper_model_size.to_whisper_cpp_model_size()
|
||||
|
||||
whisper_cpp_model_files = [
|
||||
f"ggml-{model_name}.bin",
|
||||
"README.md"
|
||||
]
|
||||
if self.is_coreml_supported:
|
||||
whisper_cpp_model_files = [
|
||||
f"ggml-{model_name}.bin",
|
||||
f"ggml-{model_name}-encoder.mlmodelc.zip",
|
||||
"README.md"
|
||||
]
|
||||
|
||||
model_path = download_from_huggingface(
|
||||
repo_id=repo_id,
|
||||
allow_patterns=whisper_cpp_model_files,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if self.is_coreml_supported:
|
||||
import tempfile
|
||||
|
||||
target_dir = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc")
|
||||
zip_path = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc.zip")
|
||||
|
||||
if os.path.exists(target_dir):
|
||||
shutil.rmtree(target_dir)
|
||||
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
||||
zip_ref.extractall(temp_dir)
|
||||
|
||||
macosx_path = os.path.join(temp_dir, "__MACOSX")
|
||||
if os.path.exists(macosx_path):
|
||||
shutil.rmtree(macosx_path)
|
||||
|
||||
temp_contents = os.listdir(temp_dir)
|
||||
if len(temp_contents) == 1 and os.path.isdir(os.path.join(temp_dir, temp_contents[0])):
|
||||
nested_dir = os.path.join(temp_dir, temp_contents[0])
|
||||
shutil.move(nested_dir, target_dir)
|
||||
else:
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
for item in temp_contents:
|
||||
src = os.path.join(temp_dir, item)
|
||||
dst = os.path.join(target_dir, item)
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dst)
|
||||
else:
|
||||
shutil.copy2(src, dst)
|
||||
|
||||
self.signals.finished.emit(os.path.join(
|
||||
model_path, f"ggml-{model_name}.bin"))
|
||||
|
||||
def _download_whisper(self) -> None:
|
||||
url = whisper._MODELS[self.model.whisper_model_size.value]
|
||||
file_path = get_whisper_file_path(
|
||||
size=self.model.whisper_model_size)
|
||||
expected_sha256 = url.split("/")[-2]
|
||||
self.download_model_to_path(
|
||||
url=url, file_path=file_path, expected_sha256=expected_sha256
|
||||
)
|
||||
|
||||
def _download_faster_whisper(self) -> None:
|
||||
model_path = download_faster_whisper_model(
|
||||
model=self.model,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if model_path == "":
|
||||
self.signals.error.emit(_("Error"))
|
||||
|
||||
self.signals.finished.emit(model_path)
|
||||
|
||||
def _download_hugging_face(self) -> None:
|
||||
model_path = download_from_huggingface(
|
||||
self.model.hugging_face_model_id,
|
||||
allow_patterns=HUGGING_FACE_MODEL_ALLOW_PATTERNS,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if model_path == "":
|
||||
self.signals.error.emit(_("Error"))
|
||||
|
||||
self.signals.finished.emit(model_path)
|
||||
|
||||
def _download_openai_whisper_api(self) -> None:
|
||||
self.signals.finished.emit("")
|
||||
|
||||
def run(self) -> None:
|
||||
logging.debug("Downloading model: %s, %s", self.model,
|
||||
self.model.hugging_face_model_id)
|
||||
|
||||
if self.model.model_type == ModelType.WHISPER_CPP:
|
||||
if self.custom_model_url:
|
||||
url = self.custom_model_url
|
||||
file_path = get_whisper_cpp_file_path(
|
||||
size=self.model.whisper_model_size)
|
||||
return self.download_model_to_path(url=url, file_path=file_path)
|
||||
|
||||
repo_id = WHISPER_CPP_REPO_ID
|
||||
|
||||
if self.model.whisper_model_size == WhisperModelSize.LUMII:
|
||||
repo_id = WHISPER_CPP_LUMII_REPO_ID
|
||||
|
||||
model_name = self.model.whisper_model_size.to_whisper_cpp_model_size()
|
||||
|
||||
whisper_cpp_model_files = [
|
||||
f"ggml-{model_name}.bin",
|
||||
"README.md"
|
||||
]
|
||||
if self.is_coreml_supported:
|
||||
whisper_cpp_model_files = [
|
||||
f"ggml-{model_name}.bin",
|
||||
f"ggml-{model_name}-encoder.mlmodelc.zip",
|
||||
"README.md"
|
||||
]
|
||||
|
||||
model_path = download_from_huggingface(
|
||||
repo_id=repo_id,
|
||||
allow_patterns=whisper_cpp_model_files,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if self.is_coreml_supported:
|
||||
import tempfile
|
||||
|
||||
target_dir = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc")
|
||||
zip_path = os.path.join(model_path, f"ggml-{model_name}-encoder.mlmodelc.zip")
|
||||
|
||||
# Remove target directory if it exists
|
||||
if os.path.exists(target_dir):
|
||||
shutil.rmtree(target_dir)
|
||||
|
||||
# Extract to a temporary directory first
|
||||
with tempfile.TemporaryDirectory() as temp_dir:
|
||||
with zipfile.ZipFile(zip_path, 'r') as zip_ref:
|
||||
zip_ref.extractall(temp_dir)
|
||||
|
||||
# Remove __MACOSX metadata folders if present
|
||||
macosx_path = os.path.join(temp_dir, "__MACOSX")
|
||||
if os.path.exists(macosx_path):
|
||||
shutil.rmtree(macosx_path)
|
||||
|
||||
# Check if there's a single top-level directory
|
||||
temp_contents = os.listdir(temp_dir)
|
||||
if len(temp_contents) == 1 and os.path.isdir(os.path.join(temp_dir, temp_contents[0])):
|
||||
# Single directory - move its contents to target
|
||||
nested_dir = os.path.join(temp_dir, temp_contents[0])
|
||||
shutil.move(nested_dir, target_dir)
|
||||
else:
|
||||
# Multiple items or files - copy everything to target
|
||||
os.makedirs(target_dir, exist_ok=True)
|
||||
for item in temp_contents:
|
||||
src = os.path.join(temp_dir, item)
|
||||
dst = os.path.join(target_dir, item)
|
||||
if os.path.isdir(src):
|
||||
shutil.copytree(src, dst)
|
||||
else:
|
||||
shutil.copy2(src, dst)
|
||||
|
||||
self.signals.finished.emit(os.path.join(
|
||||
model_path, f"ggml-{model_name}.bin"))
|
||||
return
|
||||
|
||||
if self.model.model_type == ModelType.WHISPER:
|
||||
url = whisper._MODELS[self.model.whisper_model_size.value]
|
||||
file_path = get_whisper_file_path(
|
||||
size=self.model.whisper_model_size)
|
||||
expected_sha256 = url.split("/")[-2]
|
||||
return self.download_model_to_path(
|
||||
url=url, file_path=file_path, expected_sha256=expected_sha256
|
||||
)
|
||||
|
||||
if self.model.model_type == ModelType.FASTER_WHISPER:
|
||||
model_path = download_faster_whisper_model(
|
||||
model=self.model,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if model_path == "":
|
||||
self.signals.error.emit(_("Error"))
|
||||
|
||||
self.signals.finished.emit(model_path)
|
||||
return
|
||||
|
||||
if self.model.model_type == ModelType.HUGGING_FACE:
|
||||
model_path = download_from_huggingface(
|
||||
self.model.hugging_face_model_id,
|
||||
allow_patterns=HUGGING_FACE_MODEL_ALLOW_PATTERNS,
|
||||
progress=self.signals.progress,
|
||||
on_process=self._register_process,
|
||||
)
|
||||
|
||||
if self.stopped:
|
||||
return
|
||||
|
||||
if model_path == "":
|
||||
self.signals.error.emit(_("Error"))
|
||||
|
||||
self.signals.finished.emit(model_path)
|
||||
return
|
||||
|
||||
if self.model.model_type == ModelType.OPEN_AI_WHISPER_API:
|
||||
self.signals.finished.emit("")
|
||||
return
|
||||
|
||||
raise Exception("Invalid model type: " + self.model.model_type.value)
|
||||
self._download_whisper_cpp()
|
||||
elif self.model.model_type == ModelType.WHISPER:
|
||||
self._download_whisper()
|
||||
elif self.model.model_type == ModelType.FASTER_WHISPER:
|
||||
self._download_faster_whisper()
|
||||
elif self.model.model_type == ModelType.HUGGING_FACE:
|
||||
self._download_hugging_face()
|
||||
elif self.model.model_type == ModelType.OPEN_AI_WHISPER_API:
|
||||
self._download_openai_whisper_api()
|
||||
else:
|
||||
raise Exception("Invalid model type: " + self.model.model_type.value)
|
||||
|
||||
def download_model_to_path(
|
||||
self, url: str, file_path: str, expected_sha256: Optional[str] = None
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue