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

# Load the pretrained model (you can choose "small", "medium", or "large")
sound_model = AudioGen.get_pretrained('facebook/audiogen-medium')
music_model = MusicGen.get_pretrained('facebook/musicgen-small')

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

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.
    """
    # Generate the audio for the provided prompt
    descriptions = [prompt]  # We use the prompt as a description for the model
    wav = sound_model.generate(descriptions)  # Generates 2 samples
    
    # Save the generated audio file with loudness normalization
    output_path = 'generated_audio'
    audio_write(output_path, wav[0].cpu(), sound_model.sample_rate, strategy="loudness")
    
    return f"{output_path}.wav"


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.
    """
    # Generate the music for the provided prompt
    descriptions = [prompt]  # We use the prompt as a description for the model
    wav = music_model.generate(descriptions)  # Generates 2 samples
    
    # Save the generated audio file with loudness normalization
    output_path = 'generated_audio'
    audio_write(output_path, wav[0].cpu(), music_model.sample_rate, strategy="loudness")
    
    return f"{output_path}.wav"