Politrees commited on
Commit
d642c92
·
verified ·
1 Parent(s): 88c2f54

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +166 -139
app.py CHANGED
@@ -105,166 +105,193 @@ def rename_stems(input_file, output_dir, stems, output_format):
105
  for i, stem in enumerate(stems):
106
  new_name = f"{base_name}_(Stem{i+1}).{output_format}"
107
  new_path = os.path.join(output_dir, new_name)
108
- os.rename(os.path.join(output_dir, stem), new_path)
109
- renamed_stems.append(new_path)
 
 
 
110
  return renamed_stems
111
 
112
  def prepare_output_dir(input_file, output_dir):
113
  """Create a directory for the output files and clean it if it already exists."""
114
  base_name = os.path.splitext(os.path.basename(input_file))[0]
115
  out_dir = os.path.join(output_dir, base_name)
116
- if os.path.exists(out_dir):
117
- shutil.rmtree(out_dir)
118
- os.makedirs(out_dir)
 
 
 
 
119
  return out_dir
120
 
121
  def roformer_separator(audio, model_key, seg_size, overlap, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
122
  """Separate audio using Roformer model."""
123
  model = ROFORMER_MODELS[model_key]
124
- out_dir = prepare_output_dir(audio, out_dir)
125
- separator = Separator(
126
- log_level=logging.WARNING,
127
- model_file_dir=model_dir,
128
- output_dir=out_dir,
129
- output_format=out_format,
130
- normalization_threshold=norm_thresh,
131
- amplification_threshold=amp_thresh,
132
- mdxc_params={
133
- "batch_size": 1,
134
- "segment_size": seg_size,
135
- "overlap": overlap,
136
- }
137
- )
138
-
139
- progress(0.2, desc="Model loaded")
140
- separator.load_model(model_filename=model)
141
-
142
- progress(0.7, desc="Audio separated")
143
- separation = separator.separate(audio)
144
-
145
- progress(0.9, desc="Stems renamed")
146
- stems = rename_stems(audio, out_dir, separation, out_format)
147
-
148
- return stems[0], stems[1]
 
 
 
 
149
 
150
  def mdx23c_separator(audio, model, seg_size, overlap, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
151
  """Separate audio using MDX23C model."""
152
- out_dir = prepare_output_dir(audio, out_dir)
153
- separator = Separator(
154
- log_level=logging.WARNING,
155
- model_file_dir=model_dir,
156
- output_dir=out_dir,
157
- output_format=out_format,
158
- normalization_threshold=norm_thresh,
159
- amplification_threshold=amp_thresh,
160
- mdxc_params={
161
- "batch_size": 1,
162
- "segment_size": seg_size,
163
- "overlap": overlap,
164
- }
165
- )
166
-
167
- progress(0.2, desc="Model loaded")
168
- separator.load_model(model_filename=model)
169
-
170
- progress(0.7, desc="Audio separated")
171
- separation = separator.separate(audio)
172
-
173
- progress(0.9, desc="Stems renamed")
174
- stems = rename_stems(audio, out_dir, separation, out_format)
175
-
176
- return stems[0], stems[1]
 
 
 
 
177
 
178
  def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
179
  """Separate audio using MDX-NET model."""
180
- out_dir = prepare_output_dir(audio, out_dir)
181
- separator = Separator(
182
- log_level=logging.WARNING,
183
- model_file_dir=model_dir,
184
- output_dir=out_dir,
185
- output_format=out_format,
186
- normalization_threshold=norm_thresh,
187
- amplification_threshold=amp_thresh,
188
- mdx_params={
189
- "batch_size": 1,
190
- "hop_length": hop_length,
191
- "segment_size": seg_size,
192
- "overlap": overlap,
193
- "enable_denoise": denoise,
194
- }
195
- )
196
-
197
- progress(0.2, desc="Model loaded")
198
- separator.load_model(model_filename=model)
199
-
200
- progress(0.7, desc="Audio separated")
201
- separation = separator.separate(audio)
202
-
203
- progress(0.9, desc="Stems renamed")
204
- stems = rename_stems(audio, out_dir, separation, out_format)
205
-
206
- return stems[0], stems[1]
 
 
 
 
207
 
208
  def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
209
  """Separate audio using VR ARCH model."""
210
- out_dir = prepare_output_dir(audio, out_dir)
211
- separator = Separator(
212
- log_level=logging.WARNING,
213
- model_file_dir=model_dir,
214
- output_dir=out_dir,
215
- output_format=out_format,
216
- normalization_threshold=norm_thresh,
217
- amplification_threshold=amp_thresh,
218
- vr_params={
219
- "batch_size": 1,
220
- "window_size": window_size,
221
- "aggression": aggression,
222
- "enable_tta": tta,
223
- "enable_post_process": post_process,
224
- "post_process_threshold": post_process_threshold,
225
- "high_end_process": high_end_process,
226
- }
227
- )
228
-
229
- progress(0.2, desc="Model loaded")
230
- separator.load_model(model_filename=model)
231
-
232
- progress(0.7, desc="Audio separated")
233
- separation = separator.separate(audio)
234
-
235
- progress(0.9, desc="Stems renamed")
236
- stems = rename_stems(audio, out_dir, separation, out_format)
237
-
238
- return stems[0], stems[1]
 
 
 
 
239
 
240
  def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
241
  """Separate audio using Demucs model."""
242
- out_dir = prepare_output_dir(audio, out_dir)
243
- separator = Separator(
244
- log_level=logging.WARNING,
245
- model_file_dir=model_dir,
246
- output_dir=out_dir,
247
- output_format=out_format,
248
- normalization_threshold=norm_thresh,
249
- amplification_threshold=amp_thresh,
250
- demucs_params={
251
- "segment_size": seg_size,
252
- "shifts": shifts,
253
- "overlap": overlap,
254
- "segments_enabled": segments_enabled,
255
- }
256
- )
257
-
258
- progress(0.2, desc="Model loaded")
259
- separator.load_model(model_filename=model)
260
-
261
- progress(0.7, desc="Audio separated")
262
- separation = separator.separate(audio)
263
-
264
- progress(0.9, desc="Stems renamed")
265
- stems = rename_stems(audio, out_dir, separation, out_format)
266
-
267
- return stems[0], stems[1], stems[2], stems[3]
 
 
 
 
268
 
269
  with gr.Blocks(
270
  title="🎵 PolUVR - Politrees 🎵",
@@ -285,8 +312,8 @@ with gr.Blocks(
285
  output_dir = gr.Textbox(value="output", label="File output directory", placeholder="output", interactive=False)
286
  output_format = gr.Dropdown(value="wav", choices=["wav", "flac", "mp3"], label="Output Format")
287
  with gr.Row():
288
- norm_threshold = gr.Slider(value=0.9, step=0.1, minimum=0, maximum=1, label="Normalization", info="max peak amplitude to normalize input and output audio.")
289
- amp_threshold = gr.Slider(value=0.6, step=0.1, minimum=0, maximum=1, label="Amplification", info="min peak amplitude to amplify input and output audio.")
290
 
291
  with gr.Tab("Roformer"):
292
  with gr.Group():
 
105
  for i, stem in enumerate(stems):
106
  new_name = f"{base_name}_(Stem{i+1}).{output_format}"
107
  new_path = os.path.join(output_dir, new_name)
108
+ try:
109
+ os.rename(os.path.join(output_dir, stem), new_path)
110
+ renamed_stems.append(new_path)
111
+ except Exception as e:
112
+ logging.error(f"Failed to rename stem {stem}: {e}")
113
  return renamed_stems
114
 
115
  def prepare_output_dir(input_file, output_dir):
116
  """Create a directory for the output files and clean it if it already exists."""
117
  base_name = os.path.splitext(os.path.basename(input_file))[0]
118
  out_dir = os.path.join(output_dir, base_name)
119
+ try:
120
+ if os.path.exists(out_dir):
121
+ shutil.rmtree(out_dir)
122
+ os.makedirs(out_dir)
123
+ except Exception as e:
124
+ logging.error(f"Failed to prepare output directory {out_dir}: {e}")
125
+ raise
126
  return out_dir
127
 
128
  def roformer_separator(audio, model_key, seg_size, overlap, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
129
  """Separate audio using Roformer model."""
130
  model = ROFORMER_MODELS[model_key]
131
+ try:
132
+ out_dir = prepare_output_dir(audio, out_dir)
133
+ separator = Separator(
134
+ log_level=logging.WARNING,
135
+ model_file_dir=model_dir,
136
+ output_dir=out_dir,
137
+ output_format=out_format,
138
+ normalization_threshold=norm_thresh,
139
+ amplification_threshold=amp_thresh,
140
+ mdxc_params={
141
+ "batch_size": 1,
142
+ "segment_size": seg_size,
143
+ "overlap": overlap,
144
+ }
145
+ )
146
+
147
+ progress(0.2, desc="Model loaded")
148
+ separator.load_model(model_filename=model)
149
+
150
+ progress(0.7, desc="Audio separated")
151
+ separation = separator.separate(audio)
152
+
153
+ progress(0.9, desc="Stems renamed")
154
+ stems = rename_stems(audio, out_dir, separation, out_format)
155
+
156
+ return stems[0], stems[1]
157
+ except Exception as e:
158
+ logging.error(f"Roformer separation failed: {e}")
159
+ return None, None
160
 
161
  def mdx23c_separator(audio, model, seg_size, overlap, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
162
  """Separate audio using MDX23C model."""
163
+ try:
164
+ out_dir = prepare_output_dir(audio, out_dir)
165
+ separator = Separator(
166
+ log_level=logging.WARNING,
167
+ model_file_dir=model_dir,
168
+ output_dir=out_dir,
169
+ output_format=out_format,
170
+ normalization_threshold=norm_thresh,
171
+ amplification_threshold=amp_thresh,
172
+ mdxc_params={
173
+ "batch_size": 1,
174
+ "segment_size": seg_size,
175
+ "overlap": overlap,
176
+ }
177
+ )
178
+
179
+ progress(0.2, desc="Model loaded")
180
+ separator.load_model(model_filename=model)
181
+
182
+ progress(0.7, desc="Audio separated")
183
+ separation = separator.separate(audio)
184
+
185
+ progress(0.9, desc="Stems renamed")
186
+ stems = rename_stems(audio, out_dir, separation, out_format)
187
+
188
+ return stems[0], stems[1]
189
+ except Exception as e:
190
+ logging.error(f"MDX23C separation failed: {e}")
191
+ return None, None
192
 
193
  def mdx_separator(audio, model, hop_length, seg_size, overlap, denoise, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
194
  """Separate audio using MDX-NET model."""
195
+ try:
196
+ out_dir = prepare_output_dir(audio, out_dir)
197
+ separator = Separator(
198
+ log_level=logging.WARNING,
199
+ model_file_dir=model_dir,
200
+ output_dir=out_dir,
201
+ output_format=out_format,
202
+ normalization_threshold=norm_thresh,
203
+ amplification_threshold=amp_thresh,
204
+ mdx_params={
205
+ "batch_size": 1,
206
+ "hop_length": hop_length,
207
+ "segment_size": seg_size,
208
+ "overlap": overlap,
209
+ "enable_denoise": denoise,
210
+ }
211
+ )
212
+
213
+ progress(0.2, desc="Model loaded")
214
+ separator.load_model(model_filename=model)
215
+
216
+ progress(0.7, desc="Audio separated")
217
+ separation = separator.separate(audio)
218
+
219
+ progress(0.9, desc="Stems renamed")
220
+ stems = rename_stems(audio, out_dir, separation, out_format)
221
+
222
+ return stems[0], stems[1]
223
+ except Exception as e:
224
+ logging.error(f"MDX-NET separation failed: {e}")
225
+ return None, None
226
 
227
  def vr_separator(audio, model, window_size, aggression, tta, post_process, post_process_threshold, high_end_process, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
228
  """Separate audio using VR ARCH model."""
229
+ try:
230
+ out_dir = prepare_output_dir(audio, out_dir)
231
+ separator = Separator(
232
+ log_level=logging.WARNING,
233
+ model_file_dir=model_dir,
234
+ output_dir=out_dir,
235
+ output_format=out_format,
236
+ normalization_threshold=norm_thresh,
237
+ amplification_threshold=amp_thresh,
238
+ vr_params={
239
+ "batch_size": 1,
240
+ "window_size": window_size,
241
+ "aggression": aggression,
242
+ "enable_tta": tta,
243
+ "enable_post_process": post_process,
244
+ "post_process_threshold": post_process_threshold,
245
+ "high_end_process": high_end_process,
246
+ }
247
+ )
248
+
249
+ progress(0.2, desc="Model loaded")
250
+ separator.load_model(model_filename=model)
251
+
252
+ progress(0.7, desc="Audio separated")
253
+ separation = separator.separate(audio)
254
+
255
+ progress(0.9, desc="Stems renamed")
256
+ stems = rename_stems(audio, out_dir, separation, out_format)
257
+
258
+ return stems[0], stems[1]
259
+ except Exception as e:
260
+ logging.error(f"VR ARCH separation failed: {e}")
261
+ return None, None
262
 
263
  def demucs_separator(audio, model, seg_size, shifts, overlap, segments_enabled, model_dir, out_dir, out_format, norm_thresh, amp_thresh, progress=gr.Progress()):
264
  """Separate audio using Demucs model."""
265
+ try:
266
+ out_dir = prepare_output_dir(audio, out_dir)
267
+ separator = Separator(
268
+ log_level=logging.WARNING,
269
+ model_file_dir=model_dir,
270
+ output_dir=out_dir,
271
+ output_format=out_format,
272
+ normalization_threshold=norm_thresh,
273
+ amplification_threshold=amp_thresh,
274
+ demucs_params={
275
+ "segment_size": seg_size,
276
+ "shifts": shifts,
277
+ "overlap": overlap,
278
+ "segments_enabled": segments_enabled,
279
+ }
280
+ )
281
+
282
+ progress(0.2, desc="Model loaded")
283
+ separator.load_model(model_filename=model)
284
+
285
+ progress(0.7, desc="Audio separated")
286
+ separation = separator.separate(audio)
287
+
288
+ progress(0.9, desc="Stems renamed")
289
+ stems = rename_stems(audio, out_dir, separation, out_format)
290
+
291
+ return stems[0], stems[1], stems[2], stems[3]
292
+ except Exception as e:
293
+ logging.error(f"Demucs separation failed: {e}")
294
+ return None, None, None, None
295
 
296
  with gr.Blocks(
297
  title="🎵 PolUVR - Politrees 🎵",
 
312
  output_dir = gr.Textbox(value="output", label="File output directory", placeholder="output", interactive=False)
313
  output_format = gr.Dropdown(value="wav", choices=["wav", "flac", "mp3"], label="Output Format")
314
  with gr.Row():
315
+ norm_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.9, label="Normalization", info="max peak amplitude to normalize input and output audio.")
316
+ amp_threshold = gr.Slider(minimum=0.1, maximum=1, step=0.1, value=0.6, label="Amplification", info="min peak amplitude to amplify input and output audio.")
317
 
318
  with gr.Tab("Roformer"):
319
  with gr.Group():