Spaces:
Sleeping
Sleeping
| # Copyright 2024 The YourMT3 Authors. | |
| # | |
| # Licensed under the Apache License, Version 2.0 (the "License"); | |
| # you may not use this file except in compliance with the License. | |
| # You may obtain a copy of the License at | |
| # | |
| # http://www.apache.org/licenses/LICENSE-2.0 | |
| # | |
| # Please see the details in the LICENSE file. | |
| from typing import Dict, Tuple | |
| from copy import deepcopy | |
| import soundfile as sf | |
| import torch | |
| from utils.data_modules import AMTDataModule | |
| from config.data_presets import data_preset_single_cfg, data_preset_multi_cfg | |
| from utils.augment import intra_stem_augment_processor | |
| def get_ds(data_preset_multi: Dict, train_num_samples_per_epoch: int = 90000): | |
| dm = AMTDataModule(data_preset_multi=data_preset_multi, train_num_samples_per_epoch=train_num_samples_per_epoch) | |
| dm.setup('fit') | |
| dl = dm.train_dataloader() | |
| ds = dl.flattened[0].dataset | |
| return ds | |
| def debug_func(num_segments: int = 10): | |
| sampled_data, sampled_ids = ds._get_rand_segments_from_cache(num_segments) | |
| ux_sampled_data, _ = ds._get_rand_segments_from_cache(ux_count_sum, False, sampled_ids) | |
| s = deepcopy(sampled_data) | |
| intra_stem_augment_processor(sampled_data, submix_audio=False) | |
| def gen_audio(index: int = 0): | |
| # audio_arr: (b, 1, nframe), note_token_arr: (b, l), task_token_arr: (b, task_l) | |
| audio_arr, note_token_arr, task_token_arr = ds.__getitem__(index) | |
| # merge all the segments into one audio file | |
| audio = audio_arr.permute(0, 2, 1).reshape(-1).squeeze().numpy() | |
| # save the audio file | |
| sf.write('xaug_demo_audio.wav', audio, 16000, subtype='PCM_16') | |
| data_preset_multi = data_preset_multi_cfg["all_cross_rebal5"] | |
| ds = get_ds(data_preset_multi) | |
| ds.random_amp_range = [0.8, 1.1] | |
| ds.stem_xaug_policy = { | |
| "max_k": 5, | |
| "tau": 0.3, | |
| "alpha": 1.0, | |
| "max_subunit_stems": 12, | |
| "no_instr_overlap": True, | |
| "no_drum_overlap": True, | |
| "uhat_intra_stem_augment": True, | |
| } | |
| gen_audio(3) | |
| # for k in ds.cache.keys(): | |
| # arr = ds.cache[k]['audio_array'] | |
| # arr = np.sum(arr, axis=1).reshape(-1) | |
| # # sf.write(f'xxx/{k}.wav', arr, 16000, subtype='PCM_16') | |
| # if np.min(arr) > -0.5: | |
| # print(k) | |
| # arr = ds.cache[52]['audio_array'] | |
| # for i in range(arr.shape[1]): | |
| # a = arr[:, i, :].reshape(-1) | |
| # sf.write(f'xxx52/52_{i}.wav', a, 16000, subtype='PCM_16') | |