Update app.py

app.py

@@ -7,7 +7,6 @@ from datetime import datetime
 from io import BytesIO
 from tempfile import NamedTemporaryFile
 from xmlrpc.client import Binary
-
 import numpy as np
 import requests
 import scipy
@@ -22,6 +21,15 @@ from wordpress_xmlrpc import Client
 from wordpress_xmlrpc.compat import xmlrpc_client
 from wordpress_xmlrpc.methods import media
 
+# Function to extract nonce from URL parameters
+def get_nonce_from_url():
+    params = st.experimental_get_query_params()
+    if 'nonce' in params:
+        return params['nonce'][0]  # Extract the nonce from the query parameters
+    return None
+
+nonce = get_nonce_from_url()
+
 # API URLs and Tokens
 API_URL = os.getenv("API_URL", None)
 BEARER_TOKEN = os.getenv("BEARER_TOKEN", None)
@@ -34,10 +42,11 @@ if not BEARER_TOKEN:
     st.error("BEARER_TOKEN environment variable is not set.")
     st.stop()
 
-# Define the headers for requests
+# Define the headers for requests and include nonce
 headers = {
     "Authorization": f"Bearer {BEARER_TOKEN}",
     "Content-Type": "application/json",
+    "X-WP-Nonce": nonce  # Include nonce in the headers
 }
 
 # Add background image for Streamlit app
@@ -51,42 +60,28 @@ background-size: cover;
 '''
 st.markdown(page_bg_img, unsafe_allow_html=True)
 
-# Function to extract nonce from URL parameters
-def get_nonce_from_url():
-    params = st.experimental_get_query_params()
-    if 'nonce' in params:
-        return params['nonce'][0]  # Extract the nonce from the query parameters
-    return None
-
-nonce = get_nonce_from_url()
-
-# Use the nonce in your API requests
+# Function to check subscription status via WordPress API
 def check_subscription():
     if nonce:
-        headers = {
-            "X-WP-Nonce": nonce,  # Pass the nonce in the request headers
-        }
-        response = requests.get("songlabai.com/wp-json/custom/v1/subscription-status", headers=headers)
+        response = requests.get("https://songlabai.com/wp-json/custom/v1/subscription-status", headers=headers)
         if response.status_code == 200:
             return response.json()
         else:
-            st.error("Failed to fetch subscription status.")
             return None
     else:
-        st.error("Nonce not found.")
+        st.error("Nonce not found in the URL.")
         return None
 
-# Call the subscription status function
-subscription_status = check_subscription()
-st.write("Subscription Status:", subscription_status)
-
 # Function to update generation timestamp for free users
 def update_generation_time():
-    response = requests.post("https://songlabai.com/wp-json/custom/v1/update-generation-time")
-    if response.status_code == 200:
-        st.success("Generation time updated.")
+    if nonce:
+        response = requests.post("https://songlabai.com/wp-json/custom/v1/update-generation-time", headers=headers)
+        if response.status_code == 200:
+            st.success("Generation time updated.")
+        else:
+            st.error("Failed to update generation time.")
     else:
-        st.error("Failed to update generation time.")
+        st.error("Nonce not found in the URL.")
 
 # Display title
 st.title("Songlab AI")
@@ -147,7 +142,7 @@ def generate_audio(genre, energy_level, tempo, description, duration):
     prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description}"
     payload = {"inputs": {"prompt": prompt, "duration": duration}}
     with st.spinner("Generating audio ..."):
-        response = time_post_request(API_URL, headers, payload)
+        response = time_post_request(API_URL, headers=headers, payload=payload)
     if response.status_code == 200:
         st.success("✔️ Audio Generated, Loading...")
         load_and_play_generated_audio(response)
@@ -176,4 +171,4 @@ if st.button("Generate Audio"):
     if description:
         generate_audio(genre, energy_level, tempo, description, duration)
     else:
-        st.info("Description field is required.")
+        st.info("Description field is required.")