Update modules/whisper/whisper_parameter.py

modules/whisper/whisper_parameter.py

@@ -52,6 +52,8 @@ class WhisperParameters:
     uvr_segment_size: gr.Number
     uvr_save_file: gr.Checkbox
     uvr_enable_offload: gr.Checkbox
+    whisper_enable_offload: gr.Checkbox
+    diarization_enable_offload: gr.Checkbox
     """
     A data class for Gradio components of the Whisper Parameters. Use "before" Gradio pre-processing.
     This data class is used to mitigate the key-value problem between Gradio components and function parameters.
@@ -223,6 +225,12 @@ class WhisperParameters:
     uvr_enable_offload: gr.Checkbox
         This parameter is related to UVR. Boolean value that determines whether to offload the UVR model or not
         after each transcription.
+
+    whisper_enable_offload: gr.Checkbox
+        This parameter is related to Whisper. Boolean value that determines whether to offload the Whisper model.
+
+    diarization_enable_offload: gr.Checkbox
+        This parameter is related to diarization. Boolean value that determines whether to offload the diarization model.
     """
 
     def as_list(self) -> list:
@@ -298,6 +306,8 @@ class WhisperValues:
     uvr_segment_size: int = 256
     uvr_save_file: bool = False
     uvr_enable_offload: bool = True
+    whisper_enable_offload: bool = True
+    diarization_enable_offload: bool = True
     """
     A data class to use Whisper parameters.
     """
@@ -335,6 +345,7 @@ class WhisperValues:
                 "hotwords": None if not self.hotwords else self.hotwords,
                 "language_detection_threshold": self.language_detection_threshold,
                 "language_detection_segments": self.language_detection_segments,
+                "enable_offload": self.whisper_enable_offload
             },
             "vad": {
                 "vad_filter": self.vad_filter,
@@ -347,6 +358,7 @@ class WhisperValues:
             "diarization": {
                 "is_diarize": self.is_diarize,
                 "hf_token": self.hf_token
+                "enable_offload": self.diarization_enable_offload
             },
             "bgm_separation": {
                 "is_separate_bgm": self.is_bgm_separate,