Repository: ClinicianFOCUS/FreeScribe
Pull request #451: Feat: Implement low memory mode for STT and llm model (mac-dev)
Author: login: ItsSimko
Created: 2025-03-07T21:06:51Z | Merged: 2025-03-13T17:51:09Z

## Summary by Sourcery

Implements a low memory mode that unloads the STT model when not in use to reduce memory consumption. The model is loaded and unloaded based on the low memory mode setting and when recording is toggled or a whole audio process is cancelled.

New Features:
- Introduces a low memory mode to reduce memory consumption.

Enhancements:
- Adds functionality to load and unload the STT model based on the low memory mode setting.
- Improves memory management by unloading the STT model when it is not actively being used.


Diff:
diff --git a/client_requirements.txt b/client_requirements.txt
index 5c9de08d..850ee88c 100644
--- a/client_requirements.txt
+++ b/client_requirements.txt
@@ -59,6 +59,7 @@ tqdm==4.66.5
 typing_extensions==4.12.2
 tzdata==2024.2
 tzlocal==5.2
+transformers==4.49.0
 urllib3==2.2.3
 docker==7.1.0
 markdown==3.7
diff --git a/client_requirements_nvidia.txt b/client_requirements_nvidia.txt
index eee154fb..6fd27642 100644
--- a/client_requirements_nvidia.txt
+++ b/client_requirements_nvidia.txt
@@ -59,6 +59,7 @@ tqdm==4.66.5
 typing_extensions==4.12.2
 tzdata==2024.2
 tzlocal==5.2
+transformers==4.49.0
 urllib3==2.2.3
 docker==7.1.0
 markdown==3.7
diff --git a/src/FreeScribe.client/Model.py b/src/FreeScribe.client/Model.py
index 1e6d9c8a..c22c939e 100644
--- a/src/FreeScribe.client/Model.py
+++ b/src/FreeScribe.client/Model.py
@@ -303,3 +303,5 @@ def unload_model():
             ModelManager.local_model.model.close()
             del ModelManager.local_model
             ModelManager.local_model = None
+
+        print("*** Unloaded LLM Model")
diff --git a/src/FreeScribe.client/UI/SettingsWindow.py b/src/FreeScribe.client/UI/SettingsWindow.py
index 50d235b0..cf223d12 100644
--- a/src/FreeScribe.client/UI/SettingsWindow.py
+++ b/src/FreeScribe.client/UI/SettingsWindow.py
@@ -537,7 +537,7 @@ def set_main_window(self, window):
         self.main_window = window
 
     def load_or_unload_model(self, old_model, new_model, old_use_local_llm, new_use_local_llm, old_architecture, new_architecture,
-                             old_context_window, new_context_window):
+                             old_context_window, new_context_window, old_low_mem, new_low_mem):
         """
         Determine if the model needs to be loaded or unloaded based on settings changes.
 
@@ -566,6 +566,14 @@ def load_or_unload_model(self, old_model, new_model, old_use_local_llm, new_use_
                  - reload_flag: True if new model should be loaded
         :rtype: tuple(bool, bool)
         """
+        # unload models if low mem is now on
+        if old_low_mem != new_low_mem and new_low_mem == 1:
+            return True, True
+
+        # load models if low mem is turned off
+        if old_low_mem != new_low_mem and new_low_mem == 0:
+            return True, False
+
         # Check if old model and new model are different if they are reload and make sure new model is checked.
         if old_model != new_model and new_use_local_llm == 1:
             return True, True
@@ -638,6 +646,12 @@ def update_whisper_model(self):
         old_cpu_count = self.editable_settings[SettingsKeys.WHISPER_CPU_COUNT.value]
         old_compute_type = self.editable_settings[SettingsKeys.WHISPER_COMPUTE_TYPE.value]
 
+        new_low_mem = self.editable_settings_entries[SettingsKeys.USE_LOW_MEM_MODE.value].get()
+
+        # IF unchecked then we need to load the model
+        if not new_low_mem:
+            return True
+
         # loading the model after the window is closed to prevent the window from freezing
         # if Local Whisper is selected, compare the old model with the new model and reload the model if it has changed
         # if switched from remote to local whisper
diff --git a/src/FreeScribe.client/UI/SettingsWindowUI.py b/src/FreeScribe.client/UI/SettingsWindowUI.py
index f12e03cc..c4ebaa24 100644
--- a/src/FreeScribe.client/UI/SettingsWindowUI.py
+++ b/src/FreeScribe.client/UI/SettingsWindowUI.py
@@ -31,6 +31,7 @@
 from UI.SettingsConstant import SettingsKeys, Architectures, FeatureToggle
 from UI.Widgets.PopupBox import PopupBox
 import utils.whisper.Constants
+from utils.whisper.WhisperModel import unload_stt_model
 
 LONG_ENTRY_WIDTH = 30
 SHORT_ENTRY_WIDTH = 20
@@ -653,6 +654,8 @@ def save_settings(self, close_window=True):
             self.architecture_dropdown.get(),
             self.settings.editable_settings[SettingsKeys.LOCAL_LLM_CONTEXT_WINDOW.value],
             self.settings.editable_settings_entries[SettingsKeys.LOCAL_LLM_CONTEXT_WINDOW.value].get(),
+            self.settings.editable_settings_entries[SettingsKeys.USE_LOW_MEM_MODE.value].get(),
+            self.settings.editable_settings[SettingsKeys.USE_LOW_MEM_MODE.value]
         )
 
         if self.get_selected_model() not in ["Loading models...", "Failed to load models"]:
@@ -685,9 +688,14 @@ def save_settings(self, close_window=True):
         # unload / reload model after the settings are saved
         if local_model_unload_flag:
             ModelManager.unload_model()
-        if local_model_reload_flag and not self.is_low_mem_mode():
+        if local_model_reload_flag:
             ModelManager.start_model_threaded(self.settings, self.main_window.root)
 
+        # check if we should unload the model
+        # unload models if low mem is now on
+        if self.settings.editable_settings_entries[SettingsKeys.USE_LOW_MEM_MODE.value].get():
+            unload_stt_model()
+
         if self.settings.editable_settings["Use Docker Status Bar"] and self.main_window.docker_status_bar is None:
             self.main_window.create_docker_status_bar()
         elif not self.settings.editable_settings["Use Docker Status Bar"] and self.main_window.docker_status_bar is not None:
diff --git a/src/FreeScribe.client/client.py b/src/FreeScribe.client/client.py
index a0bf4388..0a8effe2 100644
--- a/src/FreeScribe.client/client.py
+++ b/src/FreeScribe.client/client.py
@@ -59,7 +59,7 @@
 from UI.Widgets.TimestampListbox import TimestampListbox
 from UI.ScrubWindow import ScrubWindow
 from Model import ModelStatus
-from utils.whisper.WhisperModel import load_stt_model, faster_whisper_transcribe, is_whisper_valid, is_whisper_lock, load_model_with_loading_screen
+from utils.whisper.WhisperModel import load_stt_model, faster_whisper_transcribe, is_whisper_valid, is_whisper_lock, load_model_with_loading_screen, unload_stt_model, get_model_from_settings
 
 
 if os.environ.get("FREESCRIBE_DEBUG"):
@@ -588,6 +588,10 @@ def realtime_text():
                         # close buffer. we dont need it anymore
                         buffer.close()
             audio_queue.task_done()
+
+        # unload thestt model on low mem mode
+        if app_settings.is_low_mem_mode():
+            unload_stt_model()  
     else:
         is_realtimeactive = False
 
@@ -628,6 +632,12 @@ def toggle_recording():
     realtime_thread = threaded_realtime_text()
 
     if not is_recording:
+        #load the stt model for transcription
+        if not is_whisper_valid() and app_settings.is_low_mem_mode():
+            loading_screen = LoadingWindow(root, "Loading Speech to Text model", "Loading Speech to Text model. Please wait.")
+            load_stt_model(app_settings=app_settings)
+            loading_screen.destroy()
+            
         disable_recording_ui_elements()
         REALTIME_TRANSCRIBE_THREAD_ID = realtime_thread.ident
         user_input.scrolled_text.configure(state='normal')
@@ -911,11 +921,24 @@ def cancel_whole_audio_process(thread_id):
             delete_file = False if uploaded_file_path else True
             uploaded_file_path = None
 
+            # load stt model for transcription
+            if not is_whisper_valid() and app_settings.is_low_mem_mode():
+                model_id = get_model_from_settings(app_settings=app_settings)
+                model_load_window = LoadingWindow(root, 
+                title = "Speech to Text model", 
+                initial_text = f"Loading Speech to Text model({model_id}). Please wait.")
+                load_thread = load_stt_model(app_settings=app_settings)
+                load_thread.join()
+                model_load_window.destroy()
+
             # Transcribe the audio file using the loaded model
             try:
                 result = faster_whisper_transcribe(file_to_send, app_settings=app_settings)
             except Exception as e:
                 result = f"An error occurred ({type(e).__name__}): {e}"
+            finally:
+                if app_settings.is_low_mem_mode():
+                    unload_stt_model()
 
             transcribed_text = result
 
@@ -1458,7 +1481,6 @@ def cancel_note_generation(thread_id, screen_thread):
 
     
     loading_window.destroy()
-    loading_window = LoadingWindow(root, "Generating Note.", "Generating Note. Please wait.", on_cancel=lambda: (cancel_note_generation(GENERATION_THREAD_ID, screen_thread)))
 
     loading_window = LoadingWindow(
         root,


Reviews:
user: login: sourcery-ai[bot]
body: Hey @ItsSimko - I've reviewed your changes - here's some feedback:

**Overall Comments**:
- Consider adding logging to track when the STT model is loaded and unloaded for debugging purposes.
- It might be good to add a check to prevent loading the model multiple times simultaneously.

<details>
<summary>Here's what I looked at during the review</summary>

- 🟡 **General issues**: 1 issue found
- 🟢 **Security**: all looks good
- 🟢 **Testing**: all looks good
- 🟢 **Complexity**: all looks good
- 🟢 **Documentation**: all looks good
</details>

***

<details>
<summary>Sourcery is free for open source - if you like our reviews please consider sharing them ✨</summary>

- [X](https://twitter.com/intent/tweet?text=I%20just%20got%20an%20instant%20code%20review%20from%20%40SourceryAI%2C%20and%20it%20was%20brilliant%21%20It%27s%20free%20for%20open%20source%20and%20has%20a%20free%20trial%20for%20private%20code.%20Check%20it%20out%20https%3A//sourcery.ai)
- [Mastodon](https://mastodon.social/share?text=I%20just%20got%20an%20instant%20code%20review%20from%20%40SourceryAI%2C%20and%20it%20was%20brilliant%21%20It%27s%20free%20for%20open%20source%20and%20has%20a%20free%20trial%20for%20private%20code.%20Check%20it%20out%20https%3A//sourcery.ai)
- [LinkedIn](https://www.linkedin.com/sharing/share-offsite/?url=https://sourcery.ai)
- [Facebook](https://www.facebook.com/sharer/sharer.php?u=https://sourcery.ai)

</details>

<sub>
Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.
</sub>
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-07T21:09:44Z
user: login: ItsSimko
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-07T21:11:15Z

Review comments:
diff_hunk: @@ -537,7 +537,7 @@ def set_main_window(self, window):
         self.main_window = window
 
     def load_or_unload_model(self, old_model, new_model, old_use_local_llm, new_use_local_llm, old_architecture, new_architecture,
-                             old_context_window, new_context_window):
+                             old_context_window, new_context_window, old_low_mem, new_low_mem):
path: src/FreeScribe.client/UI/SettingsWindow.py
user: login: sourcery-ai[bot]
body: **suggestion:** Standardize low memory flag type usage.

The code uses comparisons like 'if new_low_mem == 1:' and elsewhere checks 'if not new_low_mem:'. To avoid potential confusion or bugs, consider converting the low memory flag to a consistent type (e.g. boolean) throughout the code.
created_at: 2025-03-07T21:09:44Z
author_association: NONE
diff_hunk: @@ -537,7 +537,7 @@ def set_main_window(self, window):
         self.main_window = window
 
     def load_or_unload_model(self, old_model, new_model, old_use_local_llm, new_use_local_llm, old_architecture, new_architecture,
-                             old_context_window, new_context_window):
+                             old_context_window, new_context_window, old_low_mem, new_low_mem):
path: src/FreeScribe.client/UI/SettingsWindow.py
user: login: ItsSimko
body: Its the same comparison. 1 is true. The settings.txt save bools as 1 and 0 so this is to prevent any errors. Xun fixed this in a different pr in the dev branch
created_at: 2025-03-07T21:11:15Z
author_association: COLLABORATOR

PR comments:
user: login: sourcery-ai[bot]
created_at: 2025-03-07T21:06:55Z
body: <!-- Generated by sourcery-ai[bot]: start review_guide -->

## Reviewer's Guide by Sourcery

This pull request introduces a low memory mode that reduces memory consumption by unloading the STT model when it is not in use. The model is loaded and unloaded based on the low memory mode setting and when recording is toggled or a whole audio process is cancelled.

#### Sequence diagram for toggling recording with low memory mode enabled

```mermaid
sequenceDiagram
  participant User
  participant Client
  participant AppSettings
  participant WhisperModel

  User->>Client: Toggles Recording
  Client->>AppSettings: is_low_mem_mode()
  AppSettings-->>Client: True
  Client->>WhisperModel: is_whisper_valid()
  WhisperModel-->>Client: False
  Client->>Client: Show Loading Screen
  Client->>WhisperModel: load_stt_model(app_settings)
  WhisperModel-->>Client: Model Loaded
  Client->>Client: Destroy Loading Screen
  Client->>Client: Start Realtime Transcription
```

#### Sequence diagram for realtime text processing with low memory mode enabled

```mermaid
sequenceDiagram
  participant Client
  participant AppSettings
  participant WhisperModel

  Client->>Client: Process Realtime Audio
  Client->>AppSettings: is_low_mem_mode()
  AppSettings-->>Client: True
  Client->>WhisperModel: unload_stt_model()
  WhisperModel-->>Client: Model Unloaded
```

#### Updated class diagram for ModelManager

```mermaid
classDiagram
  class ModelManager {
    +local_model
    +start_model_threaded()
    +unload_model()
  }

  note for ModelManager "ModelManager now unloads the LLM model and prints a message when unload_model() is called."
```

### File-Level Changes

| Change | Details | Files |
| ------ | ------- | ----- |
| Implements low memory mode by loading and unloading the STT model based on settings and recording state. | <ul><li>Added `unload_stt_model` to unload the STT model.</li><li>Modified `realtime_text` to unload the STT model when low memory mode is enabled and the audio process is complete.</li><li>Modified `toggle_recording` to load the STT model when low memory mode is enabled and recording starts.</li><li>Modified `cancel_whole_audio_process` to load the STT model when low memory mode is enabled and the audio process is cancelled, and unload the STT model after transcription if low memory mode is enabled.</li><li>Modified `save_settings` to unload the STT model when low memory mode is enabled.</li><li>Modified `load_or_unload_model` to unload the STT model when low memory mode is enabled and load the STT model when low memory mode is disabled.</li></ul> | `src/FreeScribe.client/client.py`<br/>`src/FreeScribe.client/UI/SettingsWindow.py`<br/>`src/FreeScribe.client/UI/SettingsWindowUI.py`<br/>`src/FreeScribe.client/Model.py` |

---

<details>
<summary>Tips and commands</summary>

#### Interacting with Sourcery

- **Trigger a new review:** Comment `@sourcery-ai review` on the pull request.
- **Continue discussions:** Reply directly to Sourcery's review comments.
- **Generate a GitHub issue from a review comment:** Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with `@sourcery-ai issue` to create an issue from it.
- **Generate a pull request title:** Write `@sourcery-ai` anywhere in the pull
  request title to generate a title at any time. You can also comment
  `@sourcery-ai title` on the pull request to (re-)generate the title at any time.
- **Generate a pull request summary:** Write `@sourcery-ai summary` anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment `@sourcery-ai summary` on the pull request to
  (re-)generate the summary at any time.
- **Generate reviewer's guide:** Comment `@sourcery-ai guide` on the pull
  request to (re-)generate the reviewer's guide at any time.
- **Resolve all Sourcery comments:** Comment `@sourcery-ai resolve` on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
- **Dismiss all Sourcery reviews:** Comment `@sourcery-ai dismiss` on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  `@sourcery-ai review` to trigger a new review!
- **Generate a plan of action for an issue:** Comment `@sourcery-ai plan` on
  an issue to generate a plan of action for it.

#### Customizing Your Experience

Access your [dashboard](https://app.sourcery.ai) to:
- Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
- Change the review language.
- Add, remove or edit custom review instructions.
- Adjust other review settings.

#### Getting Help

- [Contact our support team](mailto:support@sourcery.ai) for questions or feedback.
- Visit our [documentation](https://docs.sourcery.ai) for detailed guides and information.
- Keep in touch with the Sourcery team by following us on [X/Twitter](https://x.com/SourceryAI), [LinkedIn](https://www.linkedin.com/company/sourcery-ai/) or [GitHub](https://github.com/sourcery-ai).

</details>

<!-- Generated by sourcery-ai[bot]: end review_guide -->
author_association: NONE
user: login: ItsSimko
created_at: 2025-03-07T21:09:36Z
body: @sourcery-ai review
author_association: COLLABORATOR