diff --git a/buzz/model_loader.py b/buzz/model_loader.py index 224bd6d3..af7facdb 100644 --- a/buzz/model_loader.py +++ b/buzz/model_loader.py @@ -482,8 +482,7 @@ class HuggingfaceDownloadMonitor: def __init__(self, model_root: str, progress: pyqtSignal(tuple), total_file_size: int): self.model_root = model_root self.progress = progress - # To keep dialog open even if it reports 100% - self.total_file_size = round(total_file_size * 1.1) + self.total_file_size = total_file_size self.incomplete_download_root = None self.stop_event = threading.Event() self.monitor_thread = None diff --git a/buzz/widgets/model_download_progress_dialog.py b/buzz/widgets/model_download_progress_dialog.py index 3c2bba80..0b1daa79 100644 --- a/buzz/widgets/model_download_progress_dialog.py +++ b/buzz/widgets/model_download_progress_dialog.py @@ -8,6 +8,12 @@ from PyQt6.QtWidgets import QProgressDialog, QWidget, QPushButton from buzz.locale import _ from buzz.model_loader import ModelType +NO_PROGRESS_MODEL_TYPES = { + ModelType.HUGGING_FACE, + ModelType.FASTER_WHISPER, + ModelType.WHISPER_CPP, +} + class ModelDownloadProgressDialog(QProgressDialog): def __init__( @@ -22,17 +28,24 @@ class ModelDownloadProgressDialog(QProgressDialog): self.cancelable = ( model_type == ModelType.WHISPER ) + self.has_no_progress = model_type in NO_PROGRESS_MODEL_TYPES self.start_time = datetime.now() - self.setRange(0, 100) self.setMinimumDuration(0) self.setWindowModality(modality) - self.update_label_text(0) cancel_button = QPushButton(_("Cancel"), self) self.setCancelButton(cancel_button) if not self.cancelable: cancel_button.setEnabled(False) + if self.has_no_progress: + self.setRange(0, 0) + self.setLabelText(_("Downloading model")) + self.show() + else: + self.setRange(0, 100) + self.update_label_text(0) + def update_label_text(self, fraction_completed: float): downloading_text = _("Downloading model") remaining_text = _("remaining") @@ -48,6 +61,8 @@ class ModelDownloadProgressDialog(QProgressDialog): def set_value(self, fraction_completed: float): if self.wasCanceled(): return + if self.has_no_progress: + return self.setValue(int(fraction_completed * self.maximum())) self.update_label_text(fraction_completed) diff --git a/buzz/widgets/preferences_dialog/models_preferences_widget.py b/buzz/widgets/preferences_dialog/models_preferences_widget.py index 1835453f..a6466d20 100644 --- a/buzz/widgets/preferences_dialog/models_preferences_widget.py +++ b/buzz/widgets/preferences_dialog/models_preferences_widget.py @@ -261,6 +261,7 @@ class ModelsPreferencesWidget(QWidget): self.model.open_file_location() def on_download_completed(self, _: str): + self.progress_dialog.close() self.progress_dialog = None self.download_button.setEnabled(True) self.reset() diff --git a/tests/widgets/model_download_progress_dialog.py b/tests/widgets/model_download_progress_dialog.py index e9738ddc..1fb7f67d 100644 --- a/tests/widgets/model_download_progress_dialog.py +++ b/tests/widgets/model_download_progress_dialog.py @@ -9,7 +9,7 @@ class TestModelDownloadProgressDialog: def test_should_show_dialog(self, qtbot): dialog = ModelDownloadProgressDialog(model_type=ModelType.WHISPER, parent=None) qtbot.add_widget(dialog) - assert dialog.labelText() == f"{_('Downloading model')} (0%)" + assert dialog.labelText() == f"{_('Downloading model')} ()" def test_should_update_label_on_progress(self, qtbot): dialog = ModelDownloadProgressDialog(model_type=ModelType.WHISPER, parent=None) @@ -17,10 +17,12 @@ class TestModelDownloadProgressDialog: dialog.set_value(0.0) dialog.set_value(0.01) - assert dialog.labelText().startswith(f"{_('Downloading model')} (1%") + assert dialog.labelText().startswith(f"{_('Downloading model')} (") + assert dialog.labelText() != f"{_('Downloading model')} ()" dialog.set_value(0.1) - assert dialog.labelText().startswith(f"{_('Downloading model')} (10%") + assert dialog.labelText().startswith(f"{_('Downloading model')} (") + assert dialog.labelText() != f"{_('Downloading model')} ()" # Other windows should not be processing while models are being downloaded def test_should_be_an_application_modal(self, qtbot):