File size: 1,776 Bytes
4949a8d
 
e1b46cf
f8bd524
4949a8d
 
7814ee2
f8bd524
 
a008552
7814ee2
5029ee2
 
f8bd524
 
7814ee2
 
 
f26b7a5
671b217
7814ee2
 
 
 
 
 
 
 
 
f8bd524
4949a8d
f8bd524
7814ee2
4949a8d
7814ee2
 
 
 
671b217
 
7814ee2
 
 
 
 
 
 
 
 
f8bd524
4949a8d
f8bd524
7814ee2
4949a8d
7814ee2
 
 
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
import time

import spaces
import torch
from audiocraft.data.audio import audio_write
from audiocraft.models import AudioGen, MusicGen

# Load the pretrained models and move them to GPU if available
device = "cuda" if torch.cuda.is_available() else "cpu"
print("Using device:", device)

sound_model = AudioGen.get_pretrained('facebook/audiogen-medium')
music_model = MusicGen.get_pretrained('facebook/musicgen-small')

# Set generation parameters (for example, audio duration of 5 seconds)
sound_model.set_generation_params(duration=5)
music_model.set_generation_params(duration=5)

@spaces.GPU(duration=20)
def generate_sound(prompt: str):
    """
    Generate sound using Audiocraft based on the given prompt.
    
    Args:
    - prompt (str): The description of the sound/music to generate.
    
    Returns:
    - str: The path to the saved audio file.
    """
    descriptions = [prompt]  
    timestamp = str(time.time()).replace(".", "")
    wav = sound_model.generate(descriptions)  # Generate audio
    
    output_path = f'{prompt}_{timestamp}'
    audio_write(output_path, wav[0].cpu(), sound_model.sample_rate, strategy="loudness")
    
    return f"{output_path}.wav"

@spaces.GPU(duration=15)
def generate_music(prompt: str):
    """
    Generate music using Audiocraft based on the given prompt.
    
    Args:
    - prompt (str): The description of the music to generate.
    
    Returns:
    - str: The path to the saved audio file.
    """
    descriptions = [prompt]  
    timestamp = str(time.time()).replace(".", "")
    wav = music_model.generate(descriptions)  # Generate music
    
    output_path = f'{prompt}_{timestamp}'
    audio_write(output_path, wav[0].cpu(), music_model.sample_rate, strategy="loudness")
    
    return f"{output_path}.wav"