api_handler.py

import json
import time
import functools
import requests

class APIHandler:
    def __init__(self, api_key: str):
        self.models = {
            "tencent/SongGeneration": "https://api-inference.huggingface.co/models/tencent/SongGeneration",
            "daydreamlive/DreamVAE": "https://api-inference.huggingface.co/models/daydreamlive/DreamVAE"
        }
        self.api_key = api_key

    # -------------------------------
    # SONG GENERATION
    # -------------------------------
    @functools.lru_cache(maxsize=128)
    def call_song_gen(self, model, prompt, lyrics, voice, duration):
        if model not in self.models:
            return {"status": "error", "message": "Invalid model"}

        payload = {
            "prompt": prompt,
            "lyrics": lyrics,
            "voice": voice,
            "duration": duration
        }

        headers = {
            "Authorization": f"Bearer {self.api_key}",
            "Content-Type": "application/json"
        }

        try:
            response = requests.post(
                self.models[model],
                headers=headers,
                data=json.dumps(payload),
                timeout=60
            )

            if response.status_code != 200:
                return {
                    "status": "error",
                    "code": response.status_code,
                    "details": response.text
                }

            return {
                "status": "success",
                "audio_data": response.content,
                "meta": {
                    "duration": duration,
                    "voice": voice
                }
            }

        except requests.exceptions.Timeout:
            return {"status": "error", "message": "Request timed out"}

        except Exception as e:
            return {"status": "error", "message": str(e)}

    # -------------------------------
    # AUDIO EDITING
    # -------------------------------
    def call_audio_edit(self, model, audio_bytes, action, param=None):
        if model not in self.models:
            return {"status": "error", "message": "Invalid model"}

        if not audio_bytes:
            return {"status": "error", "message": "No audio provided"}

        headers = {
            "Authorization": f"Bearer {self.api_key}"
        }

        files = {
            "audio": ("input.wav", audio_bytes, "audio/wav")
        }

        data = {
            "action": action,
            "param": json.dumps(param) if param else None
        }

        try:
            response = requests.post(
                self.models[model],
                headers=headers,
                files=files,
                data=data,
                timeout=60
            )

            if response.status_code != 200:
                return {
                    "status": "error",
                    "code": response.status_code,
                    "details": response.text
                }

            return {
                "status": "success",
                "edited_audio": response.content
            }

        except requests.exceptions.Timeout:
            return {"status": "error", "message": "Audio edit request timed out"}

        except Exception as e:
            return {"status": "error", "message": str(e)}