ASesYusuf1/SESA_Audio_Separation
14
1# apollo_processing.py
2import os
3import subprocess
4import librosa
5import soundfile as sf
6import numpy as np
7from helpers import clamp_percentage, sanitize_filename
8
9def process_with_apollo(
10 output_files,
11 output_dir,
12 apollo_chunk_size,
13 apollo_overlap,
14 apollo_method,
15 apollo_normal_model,
16 apollo_midside_model,
17 output_format,
18 progress=None,
19 total_progress_start=80,
20 total_progress_end=100
21):
22 """
23 Process audio files with Apollo enhancement.
24
25 Args:
26 output_files: List of input audio file paths to process.
27 output_dir: Directory to store enhanced output files.
28 apollo_chunk_size: Chunk size for Apollo processing.
29 apollo_overlap: Overlap for Apollo processing.
30 apollo_method: Apollo processing method ('normal_method' or 'mid_side_method').
31 apollo_normal_model: Apollo model for normal method.
32 apollo_midside_model: Apollo model for mid-side method.
33 output_format: Output audio format (e.g., 'wav').
34 progress: Gradio progress object for UI updates.
35 total_progress_start: Starting progress percentage (default: 80).
36 total_progress_end: Ending progress percentage (default: 100).
37
38 Returns:
39 List of enhanced file paths or original files if processing fails.
40 """
41 try:
42 apollo_script = "/content/Apollo/inference.py"
43 print(f"Apollo parameters - chunk_size: {apollo_chunk_size}, overlap: {apollo_overlap}, method: {apollo_method}, normal_model: {apollo_normal_model}, midside_model: {apollo_midside_model}")
44
45 # Select checkpoint and config based on method and model
46 if apollo_method == "mid_side_method":
47 if apollo_midside_model == "MP3 Enhancer":
48 ckpt = "/content/Apollo/model/pytorch_model.bin"
49 config = "/content/Apollo/configs/apollo.yaml"
50 elif apollo_midside_model == "Lew Vocal Enhancer":
51 ckpt = "/content/Apollo/model/apollo_model.ckpt"
52 config = "/content/Apollo/configs/apollo.yaml"
53 elif apollo_midside_model == "Lew Vocal Enhancer v2 (beta)":
54 ckpt = "/content/Apollo/model/apollo_model_v2.ckpt"
55 config = "/content/Apollo/configs/config_apollo_vocal.yaml"
56 else:
57 ckpt = "/content/Apollo/model/apollo_universal_model.ckpt"
58 config = "/content/Apollo/configs/config_apollo.yaml"
59 else:
60 if apollo_normal_model == "MP3 Enhancer":
61 ckpt = "/content/Apollo/model/pytorch_model.bin"
62 config = "/content/Apollo/configs/apollo.yaml"
63 elif apollo_normal_model == "Lew Vocal Enhancer":
64 ckpt = "/content/Apollo/model/apollo_model.ckpt"
65 config = "/content/Apollo/configs/apollo.yaml"
66 elif apollo_normal_model == "Lew Vocal Enhancer v2 (beta)":
67 ckpt = "/content/Apollo/model/apollo_model_v2.ckpt"
68 config = "/content/Apollo/configs/config_apollo_vocal.yaml"
69 else:
70 ckpt = "/content/Apollo/model/apollo_universal_model.ckpt"
71 config = "/content/Apollo/configs/config_apollo.yaml"
72
73 if not os.path.exists(ckpt):
74 raise FileNotFoundError(f"Apollo checkpoint file not found: {ckpt}")
75 if not os.path.exists(config):
76 raise FileNotFoundError(f"Apollo configuration file not found: {config}")
77
78 enhanced_files = []
79 total_files = len([f for f in output_files if f and os.path.exists(f)])
80 progress_per_file = (total_progress_end - total_progress_start) / total_files if total_files > 0 else (total_progress_end - total_progress_start)
81
82 for idx, output_file in enumerate(output_files):
83 if output_file and os.path.exists(output_file):
84 original_file_name = sanitize_filename(os.path.splitext(os.path.basename(output_file))[0])
85 enhancement_suffix = "_Mid_Side_Enhanced" if apollo_method == "mid_side_method" else "_Enhanced"
86 enhanced_output = os.path.join(output_dir, f"{original_file_name}{enhancement_suffix}.{output_format}")
87
88 try:
89 # Progress update
90 if progress is not None and callable(getattr(progress, '__call__', None)):
91 current_progress = total_progress_start + (idx * progress_per_file)
92 current_progress = clamp_percentage(current_progress)
93 progress(current_progress, desc=f"Enhancing with Apollo... ({idx+1}/{total_files})")
94 else:
95 print(f"Progress is not callable or None, skipping Apollo progress update: file {idx+1}/{total_files}")
96
97 if apollo_method == "mid_side_method":
98 audio, sr = librosa.load(output_file, mono=False, sr=None)
99 if audio.ndim == 1:
100 audio = np.array([audio, audio])
101
102 mid = (audio[0] + audio[1]) * 0.5
103 side = (audio[0] - audio[1]) * 0.5
104
105 mid_file = os.path.join(output_dir, f"{original_file_name}_mid_temp.wav")
106 side_file = os.path.join(output_dir, f"{original_file_name}_side_temp.wav")
107 sf.write(mid_file, mid, sr)
108 sf.write(side_file, side, sr)
109
110 mid_output = os.path.join(output_dir, f"{original_file_name}_mid_enhanced.{output_format}")
111 command_mid = [
112 "python", apollo_script,
113 "--in_wav", mid_file,
114 "--out_wav", mid_output,
115 "--chunk_size", str(int(apollo_chunk_size)),
116 "--overlap", str(int(apollo_overlap)),
117 "--ckpt", ckpt,
118 "--config", config
119 ]
120 print(f"Running Apollo Mid command: {' '.join(command_mid)}")
121 result_mid = subprocess.run(command_mid, capture_output=True, text=True)
122 if result_mid.returncode != 0:
123 print(f"Apollo Mid processing failed: {result_mid.stderr}")
124 enhanced_files.append(output_file)
125 continue
126
127 side_output = os.path.join(output_dir, f"{original_file_name}_side_enhanced.{output_format}")
128 command_side = [
129 "python", apollo_script,
130 "--in_wav", side_file,
131 "--out_wav", side_output,
132 "--chunk_size", str(int(apollo_chunk_size)),
133 "--overlap", str(int(apollo_overlap)),
134 "--ckpt", ckpt,
135 "--config", config
136 ]
137 print(f"Running Apollo Side command: {' '.join(command_side)}")
138 result_side = subprocess.run(command_side, capture_output=True, text=True)
139 if result_side.returncode != 0:
140 print(f"Apollo Side processing failed: {result_side.stderr}")
141 enhanced_files.append(output_file)
142 continue
143
144 if not (os.path.exists(mid_output) and os.path.exists(side_output)):
145 print(f"Apollo outputs missing: mid={mid_output}, side={side_output}")
146 enhanced_files.append(output_file)
147 continue
148
149 mid_audio, _ = librosa.load(mid_output, sr=sr, mono=True)
150 side_audio, _ = librosa.load(side_output, sr=sr, mono=True)
151 left = mid_audio + side_audio
152 right = mid_audio - side_audio
153 combined = np.array([left, right])
154
155 os.makedirs(os.path.dirname(enhanced_output), exist_ok=True)
156 sf.write(enhanced_output, combined.T, sr)
157
158 temp_files = [mid_file, side_file, mid_output, side_output]
159 for temp_file in temp_files:
160 try:
161 if os.path.exists(temp_file):
162 os.remove(temp_file)
163 except Exception as e:
164 print(f"Could not delete temporary file {temp_file}: {str(e)}")
165
166 enhanced_files.append(enhanced_output)
167 else:
168 command = [
169 "python", apollo_script,
170 "--in_wav", output_file,
171 "--out_wav", enhanced_output,
172 "--chunk_size", str(int(apollo_chunk_size)),
173 "--overlap", str(int(apollo_overlap)),
174 "--ckpt", ckpt,
175 "--config", config
176 ]
177 print(f"Running Apollo Normal command: {' '.join(command)}")
178 apollo_process = subprocess.run(
179 command,
180 capture_output=True,
181 text=True
182 )
183 if apollo_process.returncode != 0:
184 print(f"Apollo processing failed: {output_file}: {apollo_process.stderr}")
185 enhanced_files.append(output_file)
186 continue
187
188 if not os.path.exists(enhanced_output):
189 print(f"Apollo output missing: {enhanced_output}")
190 enhanced_files.append(output_file)
191 continue
192
193 enhanced_files.append(enhanced_output)
194
195 # Progress update after each file
196 if progress is not None and callable(getattr(progress, '__call__', None)):
197 current_progress = total_progress_start + ((idx + 1) * progress_per_file)
198 current_progress = clamp_percentage(current_progress)
199 progress(current_progress, desc=f"Enhancing with Apollo... ({idx+1}/{total_files})")
200
201 except Exception as e:
202 print(f"Error during Apollo processing: {output_file}: {str(e)}")
203 enhanced_files.append(output_file)
204 continue
205 else:
206 enhanced_files.append(output_file)
207
208 # Final progress update
209 if progress is not None and callable(getattr(progress, '__call__', None)):
210 progress(total_progress_end, desc="Apollo enhancement complete")
211
212 return enhanced_files
213
214 except Exception as e:
215 print(f"Apollo processing error: {str(e)}")
216 return [f for f in output_files]