Skip to content

Commit efd6e2d

Browse files
doqhuclauderaivisdejus
authored
Import heavy ML libraries on demand to speed up startup (#1669)
Co-authored-by: Claude Opus 5.5 <noreply@anthropic.com> Co-authored-by: Raivis Dejus <orvils@gmail.com>
1 parent ee01182 commit efd6e2d

21 files changed

Lines changed: 291 additions & 117 deletions

‎buzz/file_transcriber_queue_worker.py‎

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -58,8 +58,6 @@ def _patched_check_output(*args, **kwargs):
5858
subprocess.run = _patched_run
5959
subprocess.check_output = _patched_check_output
6060

61-
from demucs import api as demucsApi
62-
6361
from buzz.locale import _
6462

6563

@@ -189,6 +187,7 @@ def _extract_speech_streaming(conn, file_path: str, speech_path: str, device: st
189187
"""Separate vocals window by window, writing them out as we go."""
190188
import numpy as np
191189
import torch
190+
from demucs import api as demucsApi
192191

193192
has_audio, duration = _probe_audio(file_path)
194193
if not has_audio:
@@ -315,11 +314,7 @@ def _speech_extraction_worker(conn, file_path: str, speech_path: str, device: st
315314

316315
from buzz.model_loader import ModelType
317316
from buzz.transcriber.file_transcriber import FileTranscriber
318-
from buzz.transcriber.openai_whisper_api_file_transcriber import (
319-
OpenAIWhisperAPIFileTranscriber,
320-
)
321317
from buzz.transcriber.transcriber import FileTranscriptionTask, Segment
322-
from buzz.transcriber.whisper_file_transcriber import WhisperFileTranscriber
323318

324319

325320
@dataclass
@@ -529,6 +524,13 @@ def _run_plugins(self) -> bool:
529524
return True
530525

531526
def _create_transcriber(self):
527+
# Imported here rather than at module level: they pull in openai, torch and
528+
# transformers, which would otherwise add seconds to app startup.
529+
from buzz.transcriber.openai_whisper_api_file_transcriber import (
530+
OpenAIWhisperAPIFileTranscriber,
531+
)
532+
from buzz.transcriber.whisper_file_transcriber import WhisperFileTranscriber
533+
532534
model_type = self.current.task.transcription_options.model.model_type
533535
if model_type == ModelType.OPEN_AI_WHISPER_API:
534536
self.current.transcriber = OpenAIWhisperAPIFileTranscriber(

‎buzz/model_loader.py‎

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -42,7 +42,6 @@
4242
_certifi_ca_bundle = None
4343

4444
import requests
45-
import whisper
4645
import huggingface_hub
4746
import zipfile
4847
from dataclasses import dataclass
@@ -764,6 +763,8 @@ def get_whisper_file_path(size: WhisperModelSize) -> str:
764763
if size == WhisperModelSize.CUSTOM:
765764
return os.path.join(root_dir, "custom")
766765

766+
import whisper # imports torch, so keep it out of app startup
767+
767768
url = whisper._MODELS[size.value]
768769
return os.path.join(root_dir, os.path.basename(url))
769770

@@ -1010,6 +1011,8 @@ def _download_whisper_cpp(self) -> None:
10101011
model_path, f"ggml-{model_name}.bin"))
10111012

10121013
def _download_whisper(self) -> None:
1014+
import whisper
1015+
10131016
url = whisper._MODELS[self.model.whisper_model_size.value]
10141017
file_path = get_whisper_file_path(
10151018
size=self.model.whisper_model_size)

‎buzz/transcriber/file_transcriber.py‎

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -10,8 +10,6 @@
1010
from pathlib import Path
1111

1212
from PyQt6.QtCore import QObject, pyqtSignal, pyqtSlot
13-
from yt_dlp import YoutubeDL
14-
from yt_dlp.utils import DownloadCancelled
1513

1614
from buzz import whisper_audio
1715
from buzz.assets import APP_BASE_DIR
@@ -138,6 +136,9 @@ def _identify_speakers(self, segments: List[Segment]) -> List[Segment]:
138136
)
139137

140138
def _download_from_url(self) -> bool:
139+
# Imported here as it is only needed for URL imports and slows app startup
140+
from yt_dlp import YoutubeDL
141+
141142
cookiefile = os.getenv("BUZZ_DOWNLOAD_COOKIEFILE")
142143

143144
temp_dir = Path(tempfile.mkdtemp()).resolve()
@@ -266,6 +267,9 @@ def on_download_progress(self, data: dict):
266267
# yt-dlp calls this hook throughout the download and re-raises
267268
# DownloadCancelled out of extract_info, which aborts it cleanly.
268269
if self._download_canceled.is_set():
270+
# yt-dlp is already loaded here, as it is the one calling this hook
271+
from yt_dlp.utils import DownloadCancelled
272+
269273
raise DownloadCancelled()
270274

271275
if data.get("status") != "downloading":

‎buzz/transcriber/recording_transcriber.py‎

Lines changed: 50 additions & 32 deletions
Original file line numberDiff line numberDiff line change
@@ -15,14 +15,15 @@
1515
# Preload CUDA libraries before importing torch
1616
from buzz import cuda_setup # noqa: F401
1717

18-
import torch
18+
# torch, whisper, faster_whisper, transformers and openai are imported where they
19+
# are used: importing them all takes seconds, and opening the live recording window
20+
# would otherwise freeze the UI while only the selected backend needs its library.
1921
import numpy as np
2022
from buzz.transcriber.cuda_device import cuda_works
2123
import sounddevice
2224
from sounddevice import PortAudioError
2325

2426
from buzz.audio_devices import registered_stream
25-
from openai import OpenAI
2627
from PyQt6.QtCore import QObject, pyqtSignal
2728

2829
from buzz import whisper_audio
@@ -31,12 +32,8 @@
3132
from buzz.model_loader import ModelType, map_language_to_mms
3233
from buzz.settings.settings import Settings
3334
from buzz.transcriber.transcriber import TranscriptionOptions, Task, DEFAULT_WHISPER_TEMPERATURE
34-
from buzz.transformers_whisper import TransformersTranscriber
3535
from buzz.settings.recording_transcriber_mode import RecordingTranscriberMode
3636

37-
import whisper
38-
import faster_whisper
39-
4037

4138
class RecordingTranscriber(QObject):
4239
transcription = pyqtSignal(str)
@@ -171,17 +168,6 @@ def start(self):
171168
def _load_model(self):
172169
model_path = self.model_path
173170

174-
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
175-
use_cuda = cuda_works() and force_cpu == "false"
176-
self.use_cuda = use_cuda
177-
178-
if torch.cuda.is_available():
179-
logging.debug(f"CUDA version detected: {torch.version.cuda}")
180-
181-
if self.transcription_options.model.model_type == ModelType.WHISPER:
182-
device = "cuda" if use_cuda else "cpu"
183-
return whisper.load_model(model_path, device=device)
184-
185171
if self.transcription_options.model.model_type == ModelType.WHISPER_CPP:
186172
self.start_local_whisper_server()
187173
if self.openai_client is None:
@@ -193,7 +179,42 @@ def _load_model(self):
193179
)
194180
return None
195181

182+
if self.transcription_options.model.model_type == ModelType.OPEN_AI_WHISPER_API:
183+
from openai import OpenAI
184+
185+
custom_openai_base_url = self.settings.value(
186+
key=Settings.Key.CUSTOM_OPENAI_BASE_URL, default_value="",
187+
)
188+
self.openai_client = OpenAI(
189+
api_key=self.transcription_options.openai_access_token,
190+
base_url=custom_openai_base_url if custom_openai_base_url else None,
191+
max_retries=0,
192+
)
193+
logging.debug(
194+
"Will use whisper API on %s, %s",
195+
custom_openai_base_url, self.whisper_api_model,
196+
)
197+
return None
198+
199+
# Only the backends that run in this process need torch
200+
import torch
201+
202+
force_cpu = os.getenv("BUZZ_FORCE_CPU", "false")
203+
use_cuda = cuda_works() and force_cpu == "false"
204+
self.use_cuda = use_cuda
205+
206+
if torch.cuda.is_available():
207+
logging.debug(f"CUDA version detected: {torch.version.cuda}")
208+
209+
if self.transcription_options.model.model_type == ModelType.WHISPER:
210+
import whisper
211+
212+
device = "cuda" if use_cuda else "cpu"
213+
return whisper.load_model(model_path, device=device)
214+
196215
if self.transcription_options.model.model_type == ModelType.FASTER_WHISPER:
216+
import faster_whisper
217+
197218
model_root_dir = user_cache_dir("Buzz")
198219
model_root_dir = os.path.join(model_root_dir, "models")
199220
model_root_dir = os.getenv("BUZZ_MODEL_ROOT", model_root_dir)
@@ -224,20 +245,7 @@ def _load_model(self):
224245
cpu_threads=(os.cpu_count() or 8) // 2,
225246
)
226247

227-
if self.transcription_options.model.model_type == ModelType.OPEN_AI_WHISPER_API:
228-
custom_openai_base_url = self.settings.value(
229-
key=Settings.Key.CUSTOM_OPENAI_BASE_URL, default_value="",
230-
)
231-
self.openai_client = OpenAI(
232-
api_key=self.transcription_options.openai_access_token,
233-
base_url=custom_openai_base_url if custom_openai_base_url else None,
234-
max_retries=0,
235-
)
236-
logging.debug(
237-
"Will use whisper API on %s, %s",
238-
custom_openai_base_url, self.whisper_api_model,
239-
)
240-
return None
248+
from buzz.transformers_whisper import TransformersTranscriber
241249

242250
return TransformersTranscriber(model_path)
243251

@@ -260,6 +268,8 @@ def _transcribe(self, samples, model, initial_prompt):
260268
return self._transcribe_via_api(samples, initial_prompt)
261269

262270
def _transcribe_whisper(self, samples, model, initial_prompt):
271+
import whisper
272+
263273
assert isinstance(model, whisper.Whisper)
264274
return model.transcribe(
265275
audio=samples,
@@ -272,6 +282,8 @@ def _transcribe_whisper(self, samples, model, initial_prompt):
272282
)
273283

274284
def _transcribe_faster_whisper(self, samples, model, initial_prompt):
285+
import faster_whisper
286+
275287
assert isinstance(model, faster_whisper.WhisperModel)
276288
segments, _ = model.transcribe(
277289
audio=samples,
@@ -290,6 +302,8 @@ def _transcribe_faster_whisper(self, samples, model, initial_prompt):
290302
return {"text": " ".join(segment.text for segment in segments)}
291303

292304
def _transcribe_hugging_face(self, samples, model):
305+
from buzz.transformers_whisper import TransformersTranscriber
306+
293307
assert isinstance(model, TransformersTranscriber)
294308
if model.is_mms_model:
295309
language = map_language_to_mms(
@@ -371,7 +385,9 @@ def _transcribe_via_api(self, samples, initial_prompt):
371385
def _cleanup_model(self, model):
372386
if model:
373387
del model
374-
if torch.cuda.is_available():
388+
# Nothing to free if no backend imported torch (whisper.cpp and the API)
389+
torch = sys.modules.get("torch")
390+
if torch is not None and torch.cuda.is_available():
375391
torch.cuda.empty_cache()
376392

377393
self.finished.emit()
@@ -558,6 +574,8 @@ def start_local_whisper_server(self):
558574

559575
return
560576

577+
from openai import OpenAI
578+
561579
self.openai_client = OpenAI(
562580
api_key="not-used",
563581
base_url=f"http://127.0.0.1:{self._whisper_server_port}",

‎buzz/translator.py‎

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,6 @@
44
import queue
55

66
from typing import Optional, List, Tuple
7-
from openai import OpenAI
87
from PyQt6.QtCore import QObject, pyqtSignal
98

109
from buzz.settings.settings import Settings
@@ -47,15 +46,26 @@ def __init__(
4746
key=Settings.Key.CUSTOM_OPENAI_BASE_URL, default_value=""
4847
)
4948
)
50-
openai_api_key = os.getenv(
49+
self.openai_api_key = os.getenv(
5150
"BUZZ_TRANSLATION_API_KEY",
5251
get_password(Key.OPENAI_API_KEY)
5352
)
54-
self.openai_client = OpenAI(
55-
api_key=openai_api_key,
56-
base_url=custom_openai_base_url if custom_openai_base_url else None,
57-
max_retries=0
58-
)
53+
self.openai_base_url = custom_openai_base_url if custom_openai_base_url else None
54+
self._openai_client = None
55+
56+
@property
57+
def openai_client(self):
58+
# Created on the first translation, as importing openai takes about
59+
# half a second and every transcription viewer creates a Translator
60+
if self._openai_client is None:
61+
from openai import OpenAI
62+
63+
self._openai_client = OpenAI(
64+
api_key=self.openai_api_key,
65+
base_url=self.openai_base_url,
66+
max_retries=0
67+
)
68+
return self._openai_client
5969

6070
def _translate_single(self, transcript: str, transcript_id: int) -> Tuple[str, int]:
6171
"""Translate a single transcript via the API. Returns (translation, transcript_id)."""

‎buzz/widgets/main_window.py‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -665,14 +665,16 @@ def _maybe_show_cuda_prompt(self):
665665
return
666666
if not cuda_manager.should_offer_cuda_prompt():
667667
return
668+
# Checked before is_cuda_torch_installed(), which imports torch and would
669+
# otherwise add about a second to every launch after the first.
670+
if self.settings.value(Settings.Key.CUDA_PROMPT_SHOWN, False):
671+
return
668672
# The packaged builds all ship CPU-only torch, so this only ever fires
669673
# for a pip or source install where the user set up CUDA themselves.
670674
# Offering them a second, shadowing torch would be worse than useless.
671675
if cuda_manager.is_cuda_torch_installed():
672676
logging.debug("CUDA torch already available; not offering the install")
673677
return
674-
if self.settings.value(Settings.Key.CUDA_PROMPT_SHOWN, False):
675-
return
676678
self.settings.set_value(Settings.Key.CUDA_PROMPT_SHOWN, True)
677679

678680
from PyQt6.QtCore import QTimer

‎buzz/widgets/main_window_toolbar.py‎

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -18,7 +18,6 @@
1818
TRASH_ICON_PATH,
1919
UPDATE_ICON_PATH,
2020
)
21-
from buzz.widgets.recording_transcriber_widget import RecordingTranscriberWidget
2221
from buzz.widgets.toolbar import ToolBar
2322

2423

@@ -114,9 +113,16 @@ def reset_shortcuts(self):
114113
)
115114

116115
def on_record_action_triggered(self):
116+
# Imported on demand: the recording transcriber loads torch, whisper and
117+
# transformers, which would otherwise add seconds to app startup.
118+
from buzz.widgets.recording_transcriber_widget import RecordingTranscriberWidget
119+
117120
recording_transcriber_window = RecordingTranscriberWidget(
118121
self, flags=Qt.WindowType.Window
119122
)
123+
# Otherwise every closed window stays alive until the app exits, and
124+
# tearing it down then crashes in PyQt's exit cleanup
125+
recording_transcriber_window.setAttribute(Qt.WidgetAttribute.WA_DeleteOnClose)
120126
recording_transcriber_window.show()
121127

122128
def set_stop_transcription_action_enabled(self, enabled: bool):

‎buzz/widgets/menu_bar.py‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import platform
22
import webbrowser
3-
from typing import Optional
3+
from typing import Optional, TYPE_CHECKING
44

55
from PyQt6.QtCore import pyqtSignal
66
from PyQt6.QtGui import QAction, QKeySequence
@@ -12,9 +12,9 @@
1212
from buzz.settings.shortcuts import Shortcuts
1313
from buzz.widgets.about_dialog import AboutDialog
1414
from buzz.widgets.preferences_dialog.models.preferences import Preferences
15-
from buzz.widgets.preferences_dialog.preferences_dialog import (
16-
PreferencesDialog,
17-
)
15+
16+
if TYPE_CHECKING:
17+
from buzz.widgets.preferences_dialog.preferences_dialog import PreferencesDialog
1818

1919

2020
class MenuBar(QMenuBar):
@@ -24,7 +24,7 @@ class MenuBar(QMenuBar):
2424
shortcuts_changed = pyqtSignal()
2525
openai_api_key_changed = pyqtSignal(str)
2626
preferences_changed = pyqtSignal(Preferences)
27-
preferences_dialog: Optional[PreferencesDialog] = None
27+
preferences_dialog: Optional["PreferencesDialog"] = None
2828

2929
def __init__(
3030
self,
@@ -83,6 +83,9 @@ def on_about_action_triggered(self):
8383
about_dialog.open()
8484

8585
def on_preferences_action_triggered(self):
86+
# Imported on demand: the dialog pulls in openai, which slows app startup.
87+
from buzz.widgets.preferences_dialog.preferences_dialog import PreferencesDialog
88+
8689
preferences_dialog = PreferencesDialog(
8790
shortcuts=self.shortcuts,
8891
preferences=self.preferences,

0 commit comments

Comments
 (0)