Merge pull request #291 from jhj0517/feature/add-offload

app.py

@@ -138,6 +138,8 @@ class App:
                                             choices=self.whisper_inf.music_separator.available_models)
             nb_uvr_segment_size = gr.Number(label="Segment Size", value=uvr_params["segment_size"], precision=0)
             cb_uvr_save_file = gr.Checkbox(label="Save separated files to output", value=uvr_params["save_file"])
+            cb_uvr_enable_offload = gr.Checkbox(label="Offload UVR model after separation",
+                                                value=uvr_params["enable_offload"])
 
         with gr.Accordion("VAD", open=False):
             cb_vad_filter = gr.Checkbox(label="Enable Silero VAD Filter", value=vad_params["vad_filter"],
@@ -187,7 +189,7 @@ class App:
                 language_detection_segments=nb_language_detection_segments,
                 prompt_reset_on_temperature=sld_prompt_reset_on_temperature, is_bgm_separate=cb_bgm_separation,
                 uvr_device=dd_uvr_device, uvr_model_size=dd_uvr_model_size, uvr_segment_size=nb_uvr_segment_size,
-                uvr_save_file=cb_uvr_save_file
+                uvr_save_file=cb_uvr_save_file, uvr_enable_offload=cb_uvr_enable_offload
             ),
             dd_file_format,
             cb_timestamp

configs/default_parameters.yaml

@@ -48,6 +48,7 @@ bgm_separation:
   model_size: "UVR-MDX-NET-Inst_HQ_4"
   segment_size: 256
   save_file: false
+  enable_offload: true
 
 translation:
   deepl:

modules/whisper/whisper_parameter.py

@@ -51,6 +51,7 @@ class WhisperParameters:
     uvr_device: gr.Dropdown
     uvr_segment_size: gr.Number
     uvr_save_file: gr.Checkbox
+    uvr_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.
@@ -218,6 +219,10 @@ class WhisperParameters:
         
     uvr_save_file: gr.Checkbox
         This parameter is related to UVR. Boolean value that determines whether to save the file or not.
+        
+    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.
     """
 
     def as_list(self) -> list:
@@ -292,6 +297,7 @@ class WhisperValues:
     uvr_device: str = "cuda"
     uvr_segment_size: int = 256
     uvr_save_file: bool = False
+    uvr_enable_offload: bool = True
     """
     A data class to use Whisper parameters.
     """
@@ -347,6 +353,7 @@ class WhisperValues:
                 "model_size": self.uvr_model_size,
                 "segment_size": self.uvr_segment_size,
                 "save_file": self.uvr_save_file,
+                "enable_offload": self.uvr_enable_offload
             },
         }
         return data