Spaces:
Running on Zero
Running on Zero
| """ | |
| Requires installing beat_this from https://github.com/CPJKU/beat_this | |
| """ | |
| import os | |
| import glob | |
| import multiprocessing as mp | |
| from tqdm import tqdm | |
| from beat_this.inference import File2Beats | |
| from beat_this.utils import save_beat_tsv | |
| worker_model = None | |
| out_dir = '/data/beats' | |
| def init_worker(): | |
| """ | |
| Initializer function called once per worker process. | |
| Loads the model into VRAM for this specific worker. | |
| """ | |
| global worker_model | |
| worker_model = File2Beats(checkpoint_path="final0", device="cuda", dbn=False) | |
| def process_audio(audio_path): | |
| """ | |
| The target function executed by the worker pool for each file. | |
| """ | |
| name = os.path.basename(audio_path) | |
| outpath = os.path.join(out_dir, name) | |
| if os.path.exists(outpath): | |
| return True | |
| try: | |
| beats, downbeats = worker_model(audio_path) | |
| save_beat_tsv(beats, downbeats, outpath) | |
| return True | |
| except Exception as e: | |
| return f"Error on {name}: {e}" | |
| if __name__ == '__main__': | |
| mp.set_start_method('spawn', force=True) | |
| paths = glob.glob('/data/wavs/*') | |
| os.makedirs(out_dir, exist_ok=True) | |
| NUM_WORKERS = 20 | |
| print(f"Processing {len(paths)} files with {NUM_WORKERS} parallel workers...") | |
| with mp.Pool(processes=NUM_WORKERS, initializer=init_worker) as pool: | |
| results = list(tqdm(pool.imap_unordered(process_audio, paths), total=len(paths))) | |
| errors = [r for r in results if isinstance(r, str)] | |
| if errors: | |
| print(f"\nCompleted with {len(errors)} errors:") | |
| for err in errors[:10]: | |
| print(err) |