Update TextGen/router.py

TextGen/router.py

@@ -1,4 +1,5 @@
 import os
+import time
 from pydantic import BaseModel
 from fastapi import FastAPI, HTTPException, Query
 from fastapi.responses import FileResponse
@@ -6,6 +7,7 @@ from fastapi.middleware.cors import CORSMiddleware
 
 from langchain.chains import LLMChain
 from langchain.prompts import PromptTemplate
+from suno import custom_generate_audio, get_audio_information
 from langchain_google_genai import (
     ChatGoogleGenerativeAI,
     HarmBlockThreshold,
@@ -14,6 +16,7 @@ from langchain_google_genai import (
 from TextGen import app
 from gradio_client import Client
 
+song_base_api=os.environ["VERCEL_API"]
 
 my_hf_token=os.environ["HF_TOKEN"]
 
@@ -87,4 +90,26 @@ async def generate_wav(text: str, language: str = "en"):
         return FileResponse(wav_file_path, media_type="audio/wav", filename="output.wav")
 
     except Exception as e:
-        raise HTTPException(status_code=500, detail=str(e))
+        raise HTTPException(status_code=500, detail=str(e))
+
+@app.get("/generate_song")
+async def generate_song(text: str):
+    try:
+        data = custom_generate_audio({
+            "prompt": f"{text}",
+            "make_instrumental": False,
+            "wait_audio": False
+        })
+        ids = f"{data[0]['id']},{data[1]['id']}"
+        print(f"ids: {ids}")
+
+        for _ in range(60):
+            data = get_audio_information(ids)
+            if data[0]["status"] == 'streaming':
+                print(f"{data[0]['id']} ==> {data[0]['audio_url']}")
+                print(f"{data[1]['id']} ==> {data[1]['audio_url']}")
+                break
+            # sleep 5s
+            time.sleep(5)
+    except:
+        print("Error")