lib/audio_generation.py

import os
import tempfile
import typing

import requests

class AudioGeneration:
    def __init__(self):
        self.session = requests.session()
    
    def request_generation(self, prompt: typing.Union[str, typing.List[str]]):
        response = self.session.post(self.endpoint + "generate", json={
            "caption": prompt
        })

        file_path = response.json().get("file_path")
        return file_path
    
    def request_download(self, file_path: str):
        response = self.session.post(self.endpoint + "download", json={
            "file_path": file_path
        })

        audio_file_path = None
        with tempfile.NamedTemporaryFile(delete=False) as f:
            audio_file_path = f.name
            f.write(response.content)
        
        return audio_file_path
    
    def generate(self, prompt: typing.Union[str, typing.List[str]], endpoint: str):
        self.endpoint = endpoint
        stored_file_path = self.request_generation(prompt)
        
        audio_file = self.request_download(stored_file_path)
        return audio_file