Update main.py

main.py

@@ -1,13 +1,14 @@
 from fastapi import FastAPI
 from pydantic import BaseModel
-from huggingface_hub import InferenceClient
+import requests
 import base64
 import uvicorn
+import os
 
 app = FastAPI()
 
-# Initialize the InferenceClient
-client = InferenceClient("facebook/musicgen-small")
+# Get the Hugging Face API token from environment variable
+API_TOKEN = os.environ.get("HF_API_TOKEN")
 
 class Item(BaseModel):
     prompt: str
@@ -16,15 +17,25 @@ class Item(BaseModel):
 @app.post("/generate/")
 async def generate_music(item: Item):
     try:
-        # Call the Hugging Face Inference API
-        audio = client.audio_generation(
-            item.prompt,
-            max_new_tokens=256,
-            duration=item.duration
-        )
+        # Prepare the API request
+        API_URL = "https://api-inference.huggingface.co/models/facebook/musicgen-small"
+        headers = {"Authorization": f"Bearer {API_TOKEN}"}
+        payload = {
+            "inputs": item.prompt,
+            "parameters": {"duration": item.duration}
+        }
+
+        # Make the API call
+        response = requests.post(API_URL, headers=headers, json=payload)
+        
+        if response.status_code != 200:
+            return {"error": f"API request failed with status code {response.status_code}: {response.text}"}
+
+        # The response content is the audio file
+        audio_content = response.content
         
         # Convert audio to base64
-        audio_base64 = base64.b64encode(audio).decode('utf-8')
+        audio_base64 = base64.b64encode(audio_content).decode('utf-8')
         
         return {"audio": audio_base64}
     except Exception as e: