softwareweaver/MusicGen

1# Copyright (c) Meta Platforms, Inc. and affiliates.2# All rights reserved.3#4# This source code is licensed under the license found in the5# LICENSE file in the root directory of this source tree.6"""Resampling script.7"""8import argparse9from pathlib import Path10import shutil11import typing as tp12 13import submitit14import tqdm15 16from audiocraft.data.audio import audio_read, audio_write17from audiocraft.data.audio_dataset import load_audio_meta, find_audio_files18from audiocraft.data.audio_utils import convert_audio19from audiocraft.environment import AudioCraftEnvironment20 21 22def read_txt_files(path: tp.Union[str, Path]):23    with open(args.files_path) as f:24        lines = [line.rstrip() for line in f]25        print(f"Read {len(lines)} in .txt")26        lines = [line for line in lines if Path(line).suffix not in ['.json', '.txt', '.csv']]27        print(f"Filtered and keep {len(lines)} from .txt")28        return lines29 30 31def read_egs_files(path: tp.Union[str, Path]):32    path = Path(path)33    if path.is_dir():34        if (path / 'data.jsonl').exists():35            path = path / 'data.jsonl'36        elif (path / 'data.jsonl.gz').exists():37            path = path / 'data.jsonl.gz'38        else:39            raise ValueError("Don't know where to read metadata from in the dir. "40                             "Expecting either a data.jsonl or data.jsonl.gz file but none found.")41    meta = load_audio_meta(path)42    return [m.path for m in meta]43 44 45def process_dataset(args, n_shards: int, node_index: int, task_index: tp.Optional[int] = None):46    if task_index is None:47        env = submitit.JobEnvironment()48        task_index = env.global_rank49    shard_index = node_index * args.tasks_per_node + task_index50 51    if args.files_path is None:52        lines = [m.path for m in find_audio_files(args.root_path, resolve=False, progress=True, workers=8)]53    else:54        files_path = Path(args.files_path)55        if files_path.suffix == '.txt':56            print(f"Reading file list from .txt file: {args.files_path}")57            lines = read_txt_files(args.files_path)58        else:59            print(f"Reading file list from egs: {args.files_path}")60            lines = read_egs_files(args.files_path)61 62    total_files = len(lines)63    print(64        f"Total of {total_files} processed with {n_shards} shards. " +65        f"Current idx = {shard_index} -> {total_files // n_shards} files to process"66    )67    for idx, line in tqdm.tqdm(enumerate(lines)):68 69        # skip if not part of this shard70        if idx % n_shards != shard_index:71            continue72 73        path = str(AudioCraftEnvironment.apply_dataset_mappers(line))74        root_path = str(args.root_path)75        if not root_path.endswith('/'):76            root_path += '/'77        assert path.startswith(str(root_path)), \78            f"Mismatch between path and provided root: {path} VS {root_path}"79 80        try:81            metadata_path = Path(path).with_suffix('.json')82            out_path = args.out_path / path[len(root_path):]83            out_metadata_path = out_path.with_suffix('.json')84            out_done_token = out_path.with_suffix('.done')85 86            # don't reprocess existing files87            if out_done_token.exists():88                continue89 90            print(idx, out_path, path)91            mix, sr = audio_read(path)92            mix_channels = args.channels if args.channels is not None and args.channels > 0 else mix.size(0)93            # enforce simple stereo94            out_channels = mix_channels95            if out_channels > 2:96                print(f"Mix has more than two channels: {out_channels}, enforcing 2 channels")97                out_channels = 298            out_sr = args.sample_rate if args.sample_rate is not None else sr99            out_wav = convert_audio(mix, sr, out_sr, out_channels)100            audio_write(out_path.with_suffix(''), out_wav, sample_rate=out_sr,101                        format=args.format, normalize=False, strategy='clip')102            if metadata_path.exists():103                shutil.copy(metadata_path, out_metadata_path)104            else:105                print(f"No metadata found at {str(metadata_path)}")106            out_done_token.touch()107        except Exception as e:108            print(f"Error processing file line: {line}, {e}")109 110 111if __name__ == '__main__':112    parser = argparse.ArgumentParser(description="Resample dataset with SLURM.")113    parser.add_argument(114        "--log_root",115        type=Path,116        default=Path.home() / 'tmp' / 'resample_logs',117    )118    parser.add_argument(119        "--files_path",120        type=Path,121        help="List of files to process, either .txt (one file per line) or a jsonl[.gz].",122    )123    parser.add_argument(124        "--root_path",125        type=Path,126        required=True,127        help="When rewriting paths, this will be the prefix to remove.",128    )129    parser.add_argument(130        "--out_path",131        type=Path,132        required=True,133        help="When rewriting paths, `root_path` will be replaced by this.",134    )135    parser.add_argument("--xp_name", type=str, default="shutterstock")136    parser.add_argument(137        "--nodes",138        type=int,139        default=4,140    )141    parser.add_argument(142        "--tasks_per_node",143        type=int,144        default=20,145    )146    parser.add_argument(147        "--cpus_per_task",148        type=int,149        default=4,150    )151    parser.add_argument(152        "--memory_gb",153        type=int,154        help="Memory in GB."155    )156    parser.add_argument(157        "--format",158        type=str,159        default="wav",160    )161    parser.add_argument(162        "--sample_rate",163        type=int,164        default=32000,165    )166    parser.add_argument(167        "--channels",168        type=int,169    )170    parser.add_argument(171        "--partition",172        default='learnfair',173    )174    parser.add_argument("--qos")175    parser.add_argument("--account")176    parser.add_argument("--timeout", type=int, default=4320)177    parser.add_argument('--debug', action='store_true', help='debug mode (local run)')178    args = parser.parse_args()179    n_shards = args.tasks_per_node * args.nodes180    if args.files_path is None:181        print("Warning: --files_path not provided, not recommended when processing more than 10k files.")182    if args.debug:183        print("Debugging mode")184        process_dataset(args, n_shards=n_shards, node_index=0, task_index=0)185    else:186 187        log_folder = Path(args.log_root) / args.xp_name / '%j'188        print(f"Logging to: {log_folder}")189        log_folder.parent.mkdir(parents=True, exist_ok=True)190        executor = submitit.AutoExecutor(folder=str(log_folder))191        if args.qos:192            executor.update_parameters(slurm_partition=args.partition, slurm_qos=args.qos, slurm_account=args.account)193        else:194            executor.update_parameters(slurm_partition=args.partition)195        executor.update_parameters(196            slurm_job_name=args.xp_name, timeout_min=args.timeout,197            cpus_per_task=args.cpus_per_task, tasks_per_node=args.tasks_per_node, nodes=1)198        if args.memory_gb:199            executor.update_parameters(mem=f'{args.memory_gb}GB')200        jobs = []201        with executor.batch():202            for node_index in range(args.nodes):203                job = executor.submit(process_dataset, args, n_shards=n_shards, node_index=node_index)204                jobs.append(job)205        for job in jobs:206            print(f"Waiting on job {job.job_id}")207            job.results()208