optimize gpu memory usage

preprocess/tools/vocal_separation/model.py

@@ -54,7 +54,7 @@ def build_model(args):
     return model, config
 
 
-def build_models(dict_args):
+def build_models(dict_args, use_der: bool = True):
     args = parse_args_inference(dict_args)
 
     ########## load model ##########
@@ -65,13 +65,13 @@ def build_models(dict_args):
 
     sep_model, sep_config = build_model(args)
 
+    if not use_der:
+        return sep_model, sep_config, None, None, args
+
     args.config_path = args.der_config_path
     args.start_check_point = args.der_start_check_point
-    
-    dereverb_model, dereverb_config = build_model(args)
 
-    sep_model = sep_model
-    dereverb_model = dereverb_model
+    dereverb_model, dereverb_config = build_model(args)
 
     return sep_model, sep_config, dereverb_model, dereverb_config, args
 
@@ -83,7 +83,10 @@ def main(args, sep_model=None, sep_config=None, dereverb_model=None, dereverb_co
 
     mix, _ = librosa.load(path, sr=sample_rate, mono=False)
     vocals = process(mix, sep_model, args, sep_config, device)
-    dereverbed_vocals = process(vocals.mean(0), dereverb_model, args, dereverb_config, device)
+    if dereverb_model is not None and dereverb_config is not None:
+        dereverbed_vocals = process(vocals.mean(0), dereverb_model, args, dereverb_config, device)
+    else:
+        dereverbed_vocals = vocals
     accompaniment = mix - dereverbed_vocals
 
     return mix, vocals, dereverbed_vocals, accompaniment, sample_rate
@@ -113,6 +116,7 @@ class VocalSeparator:
         der_model_path: str,
         der_config_path: str,
         *,
+        use_der: bool = False,
         model_type: str = "mel_band_roformer",
         disable_detailed_pbar: bool = True,
         device: str = "cuda",
@@ -127,6 +131,7 @@ class VocalSeparator:
             sep_start_check_point: Checkpoint path for separation model.
             der_config_path: Config path for dereverb model.
             der_start_check_point: Checkpoint path for dereverb model.
+            use_der: If False, do not load or run dereverb; vocals_dereverbed will equal vocals.
             disable_detailed_pbar: Disable detailed progress bars in underlying utils.
             verbose: Whether to print verbose logs.
         """
@@ -144,10 +149,11 @@ class VocalSeparator:
         if verbose:
             print("[vocal extraction] init: start")
 
-        sep_model, sep_config, dereverb_model, dereverb_config, args = build_models(args_dict)
+        sep_model, sep_config, dereverb_model, dereverb_config, args = build_models(args_dict, use_der=use_der)
 
         sep_model = sep_model.to(device)
-        dereverb_model = dereverb_model.to(device)
+        if dereverb_model is not None:
+            dereverb_model = dereverb_model.to(device)
 
         self.sep_model = sep_model
         self.sep_config = sep_config
@@ -159,7 +165,7 @@ class VocalSeparator:
 
         if verbose:
             print(
-                "[vocal extraction] init success: sep=loaded, dereverb=loaded, device=",
+                "[vocal extraction] init success: sep=loaded, dereverb=" + ("loaded" if use_der else "skipped") + ", device=",
                 device,
             )
 

webui_svc.py

@@ -244,17 +244,11 @@ APP_STATE = AppState()
 
 
 @spaces.GPU
-def _start_svc(
-    prompt_audio, 
-    target_audio, 
-    prompt_vocal_sep=False, 
-    target_vocal_sep=True, 
-    auto_shift=True, 
-    auto_mix_acc=True, 
-    pitch_shift=0, 
-    n_step=32, 
-    cfg=1.0, 
-    seed=42
+def _run_svc_preprocess(
+    prompt_audio,
+    target_audio,
+    prompt_vocal_sep=False,
+    target_vocal_sep=True,
 ):
 	try:
 		prompt_audio = _normalize_audio_input(prompt_audio)
@@ -288,11 +282,43 @@ def _start_svc(
 			print(target_msg, file=sys.stderr, flush=True)
 			return None
 
+		return (
+			str(session_base),
+			str(prompt_wav),
+			str(prompt_f0),
+			str(target_wav),
+			str(target_f0),
+		)
+	except Exception:
+		_print_exception("_run_svc_preprocess")
+		return None
+	finally:
+		gc.collect()
+		if torch.cuda.is_available():
+			torch.cuda.empty_cache()
+
+
+@spaces.GPU
+def _run_svc_convert(
+    preprocess_state,
+    auto_shift=True,
+    auto_mix_acc=True,
+    pitch_shift=0,
+    n_step=32,
+    cfg=1.0,
+    seed=42,
+):
+	try:
+		if preprocess_state is None or not isinstance(preprocess_state, (tuple, list)) or len(preprocess_state) != 5:
+			return None
+		session_base_str, prompt_wav, prompt_f0, target_wav, target_f0 = preprocess_state
+		session_base = Path(session_base_str)
+
 		ok, msg, generated = APP_STATE.run_svc(
-			prompt_wav_path=prompt_wav,
-			target_wav_path=target_wav,
-			prompt_f0_path=prompt_f0,
-			target_f0_path=target_f0,
+			prompt_wav_path=Path(prompt_wav),
+			target_wav_path=Path(target_wav),
+			prompt_f0_path=Path(prompt_f0),
+			target_f0_path=Path(target_f0),
 			session_base=session_base,
 			auto_shift=bool(auto_shift),
 			auto_mix_acc=bool(auto_mix_acc),
@@ -306,7 +332,7 @@ def _start_svc(
 			return None
 		return str(generated)
 	except Exception:
-		_print_exception("_start_svc")
+		_print_exception("_run_svc_convert")
 		return None
 	finally:
 		gc.collect()
@@ -314,8 +340,26 @@ def _start_svc(
 			torch.cuda.empty_cache()
 
 
+@spaces.GPU
+def _start_svc(
+    prompt_audio,
+    target_audio,
+    prompt_vocal_sep=False,
+    target_vocal_sep=True,
+    auto_shift=True,
+    auto_mix_acc=True,
+    pitch_shift=0,
+    n_step=32,
+    cfg=1.0,
+    seed=42,
+):
+	state = _run_svc_preprocess(prompt_audio, target_audio, prompt_vocal_sep, target_vocal_sep)
+	if state is None:
+		return None
+	return _run_svc_convert(state, auto_shift, auto_mix_acc, pitch_shift, n_step, cfg, seed)
+
+
 def render_tab_content() -> None:
-    """Render SVC tab content (for embedding in app.py). Same UI style as webui: two columns, no title."""
     with gr.Row(equal_height=False):
         # ── Left column: inputs & controls ──
         with gr.Column(scale=1):
@@ -351,6 +395,7 @@ def render_tab_content() -> None:
         # ── Right column: output ──
         with gr.Column(scale=1):
             output_audio = gr.Audio(label="Generated audio", type="filepath", interactive=False)
+            svc_state = gr.State(value=None)
             gr.Examples(
                 examples=EXAMPLE_LIST,
                 inputs=[prompt_audio, target_audio],
@@ -361,19 +406,12 @@ def render_tab_content() -> None:
             )
 
     run_btn.click(
-        fn=_start_svc,
-        inputs=[
-            prompt_audio,
-            target_audio,
-            prompt_vocal_sep,
-            target_vocal_sep,
-            auto_shift,
-            auto_mix_acc,
-            pitch_shift,
-            n_step,
-            cfg,
-            seed_input,
-        ],
+        fn=_run_svc_preprocess,
+        inputs=[prompt_audio, target_audio, prompt_vocal_sep, target_vocal_sep],
+        outputs=[svc_state],
+    ).then(
+        fn=_run_svc_convert,
+        inputs=[svc_state, auto_shift, auto_mix_acc, pitch_shift, n_step, cfg, seed_input],
         outputs=[output_audio],
     )