-
Notifications
You must be signed in to change notification settings - Fork 1
/
ainuralbeat.py
91 lines (74 loc) · 3.34 KB
/
ainuralbeat.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
"""
AinuralBeat class
Creates a type of binural beat with music gen AI
Using the melody model
"""
import os
import torchaudio
from audiocraft.models import MusicGen
from audiocraft.data.audio import audio_write
import uuid
class AinuralBeat:
def __init__(self, beat_type, duration):
self.valid_types = ["sleep", "meditate", "relax"]
data_path = os.path.abspath("data/")
# check if there is a data folder and if not make one
if not os.path.exists(data_path):
os.makedirs(data_path)
if beat_type not in self.valid_types:
print(f"{beat_type} not a valid type")
raise AttributeError
else:
self.beat_type = beat_type
self.duration = duration
self.model = None
self.descriptions = {
"relax": ["Relax, unwind, down tempo, loop, quiet, binaural beats, 1.8hz range, high quality sound, deep, low bpm, heart beat, low energy, chill"],
"meditate": ["Meditation, loop, focused, low tempo, soft, singing bowls, introspective, thinking, slow, high quality sound, deep, low bpm, heart beat"],
"sleep": ["Sleep, rest, night time, loop, down tempo, quiet, binaural beats, 1.8hz range, high quality sound, deep, low bpm, heart beat, low energy, chill"]
}
asset_path = os.path.abspath("assets/")
# check if there is a asset folder and if fail as assets from repo are needed or your own
if not os.path.exists(asset_path):
print("assets folder is required")
raise AttributeError
self.beat_examples = {
"sleep": f"{asset_path}/sleeptoo.mp3",
"meditate": f"{asset_path}/bowlmeditate.mp3",
"relax": f"{asset_path}/relaxcut.mp3"
}
file_id = str(uuid.uuid4()).replace("-", "")
self.output_file = f"{data_path}/{beat_type}_{file_id}.wav"
def generate_beat(self):
"""
Generate music of valid beat type with MusicGen
"""
print(f"Generating beat of type {self.beat_type}...")
try:
self.model = MusicGen.get_pretrained('melody')
self.model.set_generation_params(duration=60)
wav = self.model.generate(self.descriptions[self.beat_type])
waveform, sample_rate = torchaudio.load(
self.beat_examples[self.beat_type]
)
# expand depends on the number of descriptions and melody matching up
# waveform[None].expand(1, -1, -1)
# 1 description = 1 in the x or i of the tensor
wav = self.model.generate_with_chroma(
self.descriptions[self.beat_type],
waveform[None].expand(1, -1, -1),
sample_rate
)
for one_wav in wav:
# Will save under {idx}.wav, with loudness normalization at -14 db LUFS.
audio_write(
self.output_file,
one_wav.cpu(),
self.model.sample_rate,
strategy="loudness",
loudness_compressor=True
)
print(f"Beat of type {self.beat_type} has been generated @ {self.output_file}")
except Exception as err:
print(f"Error generating beat: {err}")
raise err