Update app.py

app.py

@@ -160,11 +160,13 @@ headers = {
 
 # Streamlit app title
 st.title("Songlab AI")
-cookies = cookie_manager.get_all(key='get_all_initial')
+cookies = cookie_manager.get_all(key='init_cookies')
 
 
 # Initialize session state for JWT token
 if "jwt_token" not in st.session_state:
+    # Synchronize cookies
+    cookies = cookie_manager.get_all(key='get_cookies')
     # Check if jwt_token is stored in cookies
     jwt_token = cookies.get('jwt_token')
     if jwt_token:
@@ -172,16 +174,14 @@ if "jwt_token" not in st.session_state:
     else:
         st.session_state["jwt_token"] = None
 
-
-
-if st.session_state["jwt_token"] is None:
+if st.session_state.get("jwt_token") is None:
     st.header("Please log in to continue")
-    username = st.text_input("Username")
-    password = st.text_input("Password", type="password")
+    username = st.text_input("Username", key="username_input")
+    password = st.text_input("Password", type="password", key="password_input")
     # Create two columns for the buttons
     col1, col2 = st.columns([1, 1])
     with col1:
-        if st.button("Log in"):
+        if st.button("Log in", key="login_button"):
             # Send POST request to WordPress JWT authentication endpoint
             login_url = "https://songlabai.com/wp-json/jwt-auth/v1/token"
             data = {
@@ -194,35 +194,35 @@ if st.session_state["jwt_token"] is None:
                 jwt_token = result["token"]
                 # Set jwt_token in session state and cookie
                 st.session_state["jwt_token"] = jwt_token
-                cookie_manager.set('jwt_token', jwt_token)
+                cookie_manager.set('jwt_token', jwt_token, key='set_jwt_token')
                 st.success("Logged in successfully!")
                 st.experimental_rerun()
             else:
                 st.error("Invalid username or password")
                 st.stop()
     with col2:
-        if st.button("Don't have an account?"):
+        if st.button("Don't have an account?", key="register_button"):
             js = "window.location.href = 'https://songlabai.com/register/';"
             html = '<script>{}</script>'.format(js)
             st.components.v1.html(html)
     st.stop()
 else:
     # User is logged in, continue with the app
-    if st.button("Log out"):
-        # Clear jwt_token from session state
-        st.session_state["jwt_token"] = None
+    if st.button("Log out", key="logout_button"):
+        # Clear all session state variables
+        for key in list(st.session_state.keys()):
+            del st.session_state[key]
         # Synchronize cookies
-        cookies = cookie_manager.get_all(key='get_all_logout')
-        # Attempt to delete the cookie
-        try:
+        cookies = cookie_manager.get_all(key='get_cookies_logout')
+        # Delete the cookie if it exists
+        if 'jwt_token' in cookies:
             cookie_manager.delete('jwt_token', key='delete_jwt_token')
-        except KeyError:
-            print("Cookie 'jwt_token' does not exist.")
         st.success("Logged out successfully!")
         st.experimental_rerun()
 
 
 
+
 # Initialize session state variables
 if "vocal_audio" not in st_state:
     st_state.vocal_audio = None
@@ -329,94 +329,100 @@ def time_post_request(api_url, headers=None, payload=None):
 
     return response
 
-def generate_audio(genre, energy_level, tempo, description, duration):
-    # Start the 300-second countdown
-    count_down = 10
-    count_down_placeholder = st.empty()
-    for seconds in range(count_down):
-        count_down_placeholder.info(
-            f"Preparing the server... Please wait ⏳ {count_down - seconds} seconds."
-        )
-        time.sleep(1)
-    count_down_placeholder.empty()
-
-    # Check if the user is logged in
-    if st.session_state.get("jwt_token") is None:
-        st.error("You must be logged in to generate audio.")
-        return
-
-    jwt_token = st.session_state["jwt_token"]
-
-    # Headers for authenticated requests
-    auth_headers = {
-        "Authorization": f"Bearer {jwt_token}"
-    }
-
-    # Get user ID
-    user_id_url = "https://songlabai.com/wp-json/custom-api/v1/current-user-id"
-    user_response = requests.get(user_id_url, headers=auth_headers)
-    if user_response.status_code == 200:
-        user_data = user_response.json()
-        user_id = user_data["user_id"]
-        st.write("User Data:", user_data)  # Add this line to display the user data
-        st.write("user_id:", user_id)  # Add this line to display the user data
-
-
-
-        # Get subscription status
-        subscription_url = "https://songlabai.com/wp-json/custom-api/subscription"
-        subscription_response = requests.get(subscription_url, headers=auth_headers)
-        if subscription_response.status_code == 200:
-            subscription_data = subscription_response.json()
-            subscription_plan_id = subscription_data["subscription_plan_id"]
-            st.write("subscription_plan_id:", subscription_plan_id)  # Add this line to display the user data
-
-            has_exceeded_daily_limit = subscription_data["has_exceeded_daily_limit"]
-            st.write("has_exceeded_daily_limit:", has_exceeded_daily_limit)  # Add this line to display the user data
+    def generate_audio(genre, energy_level, tempo, description, duration):
+        # Start the countdown
+        count_down = 10
+        count_down_placeholder = st.empty()
+        for seconds in range(count_down):
+            count_down_placeholder.info(
+                f"Preparing the server... Please wait ⏳ {count_down - seconds} seconds."
+            )
+            time.sleep(1)
+        count_down_placeholder.empty()
 
+        # Check if the user is logged in
+        if st.session_state.get("jwt_token") is None:
+            st.error("You must be logged in to generate audio.")
+            return
 
-            # Check subscription plan and limits
-            if subscription_plan_id == "576":
-                if has_exceeded_daily_limit:
-                    st.error("You have exceeded your daily limit for generating music.")
-                    return
-                else:
-                    # Update user's daily limit
-                    update_limit_url = "https://songlabai.com/wp-json/custom-api/update-limit"
-                    update_response = requests.post(update_limit_url, headers=auth_headers)
-                    if update_response.status_code != 200:
-                        st.error("Failed to update your daily limit.")
+        jwt_token = st.session_state["jwt_token"]
+
+        # Headers for authenticated requests
+        auth_headers = {
+            "Authorization": f"Bearer {jwt_token}"
+        }
+
+        # Get user ID
+        user_id_url = "https://songlabai.com/wp-json/custom-api/v1/current-user-id"
+        user_response = requests.get(user_id_url, headers=auth_headers)
+        if user_response.status_code == 200:
+            user_data = user_response.json()
+            user_id = user_data["user_id"]
+            st.write("User Data:", user_data)
+            st.write("user_id:", user_id)
+
+            # Get subscription status
+            subscription_url = "https://songlabai.com/wp-json/custom-api/subscription"
+            subscription_response = requests.get(subscription_url, headers=auth_headers)
+            if subscription_response.status_code == 200:
+                subscription_data = subscription_response.json()
+                subscription_plan_id = subscription_data["subscription_plan_id"]
+                st.write("subscription_plan_id:", subscription_plan_id)
+                has_exceeded_generation_limit = subscription_data["has_exceeded_generation_limit"]
+                st.write("has_exceeded_generation_limit:", has_exceeded_generation_limit)
+
+                # Check subscription plan and limits
+                if subscription_plan_id in [None, ""]:  # Free user
+                    if has_exceeded_generation_limit:
+                        st.error("You have reached your generation limit.")
+                        return
+                    else:
+                        # Update user's generation count
+                        update_generation_url = "https://songlabai.com/wp-json/custom-api/update-generation-count"
+                        update_response = requests.post(update_generation_url, headers=auth_headers)
+                        if update_response.status_code != 200:
+                            st.error("Failed to update your generation count.")
+                            return
+                        # Proceed to generate audio
+                elif subscription_plan_id in ["304", "305", "306"]:
+                    if has_exceeded_generation_limit:
+                        st.error("You have reached your daily generation limit.")
                         return
-                    # Proceed to generate audio
-            elif subscription_plan_id in ["304", "305", "306"]:
-                # User has unlimited generation
-                pass
+                    else:
+                        # Update user's generation count
+                        update_generation_url = "https://songlabai.com/wp-json/custom-api/update-generation-count"
+                        update_response = requests.post(update_generation_url, headers=auth_headers)
+                        if update_response.status_code != 200:
+                            st.error("Failed to update your generation count.")
+                            return
+                        # Proceed to generate audio
+                else:
+                    st.error("Your subscription does not allow you to generate music.")
+                    return
             else:
-                st.error("Your subscription does not allow you to generate music.")
+                st.error("Failed to retrieve your subscription status.")
                 return
         else:
-            st.error("Failed to retrieve your subscription status.")
+            st.error("Failed to retrieve your user information.")
             return
-    else:
-        st.error("Failed to retrieve your user information.")
-        return
 
-    # Proceed to generate audio after checks
-    prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description},"
-    payload = {"inputs": {"prompt": prompt, "duration": duration}}
+        # Proceed to generate audio after checks
+        prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description},"
+        payload = {"inputs": {"prompt": prompt, "duration": duration}}
 
-    # Use the original headers for the API request
-    with st.spinner("Generating audio ..."):
-        response = time_post_request(API_URL, headers, payload)
+        # Use the original headers for the API request
+        with st.spinner("Generating audio ..."):
+            response = time_post_request(API_URL, headers, payload)
+
+        placeholder1 = st.empty()
+        if response.status_code != 200:
+            st.error("Failed to generate audio.")
+            return
+        else:
+            placeholder1.success(f"✔️ Audio Generated, Loading...")
+            load_and_play_generated_audio(response)
+            placeholder1.empty()
 
-    placeholder1 = st.empty()
-    if response.status_code != 200:
-        st.error("Failed to generate audio.")
-        return
-    else:
-        placeholder1.success(f"✔️ Audio Generated, Loading...")
-        load_and_play_generated_audio(response)
-        placeholder1.empty()
 
 
 
@@ -461,12 +467,12 @@ def load_and_play_generated_audio(response):
         st.write(f"To download use the following product code: {product_code}")
 
 
-if st.button("Generate Audio"):
-    if genre and energy_level and description and tempo:
-        generate_audio(genre, energy_level, tempo, description, duration)
-        
-    if description == "":
-        st.info("Description field is required.")
+    if st.button("Generate Audio", key="generate_audio_button"):
+        if genre and energy_level and description and tempo:
+            generate_audio(genre, energy_level, tempo, description, duration)
+        else:
+            st.info("Description field is required.")
+
 
 # Post-processing options
 st.header("Post-processing Options")