Update app.py

app.py

@@ -322,15 +322,30 @@ def time_post_request(api_url, headers=None, payload=None, timeout=None):
 
     return response
 
-def generate_audio(genre, energy_level, tempo, description, duration):
-    # Create placeholders for dynamic updates
-    status_placeholder = st.empty()
-    progress_placeholder = st.empty()
-    subscription_info_placeholder = st.empty()
-    server_message_placeholder = st.empty()
-    stage_message_placeholder = st.empty()
+def check_server_status():
+    """Check if the server is warm by making a minimal request"""
+    warmup_payload = {"inputs": {"prompt": "warmup", "duration": 1}}
+    try:
+        response = time_post_request(
+            API_URL,
+            headers=get_api_headers(),
+            payload=warmup_payload,
+            timeout=10  # Short timeout for quick check
+        )
+        return response is not None and response.status_code == 200
+    except:
+        return False
 
-    # Define the subscription tiers information
+def wait_for_server_warmup(status_placeholder, progress_placeholder, subscription_info_placeholder, stage_placeholder):
+    """Handle server warmup with detailed progress stages"""
+    stages = [
+        {"time": 0, "message": "🚀 Initializing server resources..."},
+        {"time": 60, "message": "⚙️ Loading AI models into memory..."},
+        {"time": 180, "message": "🔧 Optimizing model parameters..."},
+        {"time": 300, "message": "📊 Configuring generation settings..."},
+        {"time": 360, "message": "✨ Finalizing server preparation..."}
+    ]
+    
     subscription_info = """
     ### 📋 **Subscription Tiers**
 
@@ -355,173 +370,131 @@ def generate_audio(genre, energy_level, tempo, description, duration):
     significantly reducing wait times from 7 minutes to just 30 seconds!
     """
 
-    # Start server warmup with an initial request
-    warmup_payload = {"inputs": {"prompt": "warmup", "duration": 1}}
-    api_headers = get_api_headers()
-    
-    # Create a thread for the initial warmup request
-    def make_warmup_request():
-        try:
-            return time_post_request(API_URL, headers=api_headers, payload=warmup_payload, timeout=600)
-        except:
-            return None
-
-    warmup_thread = threading.Thread(target=make_warmup_request)
-    warmup_thread.start()
-
-    # Start the countdown while server is warming up
     start_time = time.time()
-    total_warmup_time = 420  # 7 minutes in seconds
-    
-    # Define server initialization stages
-    stages = [
-        {"time": 0, "message": "🚀 Initializing server resources..."},
-        {"time": 60, "message": "⚙️ Loading AI models into memory..."},
-        {"time": 180, "message": "🔧 Optimizing model parameters..."},
-        {"time": 300, "message": "📊 Configuring generation settings..."},
-        {"time": 360, "message": "✨ Finalizing server preparation..."}
-    ]
+    total_warmup_time = 420  # 7 minutes
     current_stage = 0
-
-    while time.time() - start_time < total_warmup_time and warmup_thread.is_alive():
+    
+    while time.time() - start_time < total_warmup_time:
         elapsed = time.time() - start_time
         remaining = total_warmup_time - elapsed
         progress = elapsed / total_warmup_time
         
-        # Update current stage based on elapsed time
+        # Update current stage
         while current_stage < len(stages) - 1 and elapsed > stages[current_stage + 1]["time"]:
             current_stage += 1
-
+            
         minutes_remaining = int(remaining // 60)
         seconds_remaining = int(remaining % 60)
         
+        # Update UI
         status_placeholder.markdown("### 🔄 AI Server Initialization in Progress")
         progress_placeholder.progress(progress)
-        
-        stage_message_placeholder.info(stages[current_stage]["message"])
-        
-        server_message_placeholder.warning(
+        stage_placeholder.info(
+            f"{stages[current_stage]['message']}\n\n"
             f"⏳ Estimated time remaining: {minutes_remaining}m {seconds_remaining}s\n\n"
-            f"🎵 Free tier server initialization takes about 7 minutes\n\n"
             "💡 Premium users enjoy dramatically faster startup times (30s or less)!"
         )
-        
         subscription_info_placeholder.markdown(subscription_info)
+        
+        # Check if server is ready every 30 seconds
+        if elapsed % 30 == 0:
+            if check_server_status():
+                return True
+                
         time.sleep(0.1)
-
-    # Clear initialization messages
-    status_placeholder.empty()
-    progress_placeholder.empty()
-    stage_message_placeholder.empty()
-    server_message_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
-
-    # Headers for authenticated requests to your WordPress site
-    auth_headers = get_auth_headers()
-
-    # 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:
-        st.error("Failed to retrieve your user information.")
-        return
-
-    user_data = user_response.json()
-    user_id = user_data["user_id"]
-
-    # Get subscription status
-    subscription_url = "https://songlabai.com/wp-json/custom-api/v1/subscription"
-    subscription_response = requests.get(subscription_url, headers=auth_headers)
-    if subscription_response.status_code != 200:
-        st.error("Failed to retrieve your subscription status.")
-        return
-
-    subscription_data = subscription_response.json()
-    subscription_plan_id = subscription_data["subscription_plan_id"]
-    has_exceeded_generation_limit = subscription_data["has_exceeded_generation_limit"]
-
-    # Check subscription plan and limits
-    if has_exceeded_generation_limit:
-        st.error(
-            "You have reached your generation limit. "
-            "Please upgrade your subscription for additional access."
-            " [Upgrade Now](https://songlabai.com/subscribe/)"
-        )
-        return
-
-    # Prepare the prompt and payload for the music generation API
-    prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description}"
-    payload = {"inputs": {"prompt": prompt, "duration": duration}}
-
-    # Enhanced retry logic with better feedback
-    max_retries = 3
-    retry_delay = 30  # increased initial delay
-    for attempt in range(max_retries):
-        current_delay = retry_delay * (attempt + 1)  # Progressive delay
         
-        with st.spinner(f"Generating audio... Attempt {attempt+1}/{max_retries}"):
-            try:
-                response = time_post_request(API_URL, headers=api_headers, payload=payload, timeout=600)
-
-                if response is None:
-                    server_message_placeholder.warning(
-                        "🔄 Finalizing server preparation. "
-                        f"Retrying in {current_delay} seconds... \n\n"
-                        "💡 Premium subscribers enjoy faster access with dedicated servers!"
-                    )
-                    time.sleep(current_delay)
-                    continue
-
-                if response.status_code == 200:
-                    placeholder1 = st.empty()
-                    placeholder1.success("✔️ Audio Generated, Loading...")
-                    load_and_play_generated_audio(response)
-                    placeholder1.empty()
-                    
-                    # Update generation count
-                    update_generation_url = "https://songlabai.com/wp-json/custom-api/v1/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.")
-                    
-                    # Check if the user has now exceeded the limit
-                    subscription_status_url = "https://songlabai.com/wp-json/custom-api/v1/subscription"
-                    subscription_response = requests.get(subscription_status_url, headers=auth_headers)
-                    if subscription_response.status_code == 200:
-                        subscription_data = subscription_response.json()
-                        if subscription_data["has_exceeded_generation_limit"]:
-                            st.error(
-                                "You have reached your generation limit. "
-                                "Please upgrade your subscription for additional access."
-                                " [Upgrade Now](https://songlabai.com/subscribe/)"
-                            )
-                    return
-                else:
-                    server_message_placeholder.warning(
-                        f"⚠️ Attempt {attempt+1} failed with status code {response.status_code}. "
-                        f"Retrying in {current_delay} seconds...\n\n"
-                        "💡 Premium users experience fewer delays and faster generation times!"
-                    )
-                    time.sleep(current_delay)
+    return False
 
-            except requests.exceptions.Timeout:
-                server_message_placeholder.warning(
-                    "⏱️ Request timed out while server is initializing.\n\n"
-                    "💡 Premium subscribers get priority access to ready servers!"
+def generate_audio(genre, energy_level, tempo, description, duration):
+    # Create placeholders
+    status_placeholder = st.empty()
+    progress_placeholder = st.empty()
+    subscription_placeholder = st.empty()
+    stage_placeholder = st.empty()
+    
+    try:
+        # First, verify user is logged in and has permissions
+        if st.session_state.get("jwt_token") is None:
+            st.error("You must be logged in to generate audio.")
+            return
+            
+        auth_headers = get_auth_headers()
+        
+        # Check subscription status
+        subscription_url = "https://songlabai.com/wp-json/custom-api/v1/subscription"
+        subscription_response = requests.get(subscription_url, headers=auth_headers)
+        if subscription_response.status_code != 200:
+            st.error("Failed to retrieve your subscription status.")
+            return
+            
+        subscription_data = subscription_response.json()
+        if subscription_data["has_exceeded_generation_limit"]:
+            st.error(
+                "You have reached your generation limit. "
+                "Please upgrade your subscription for additional access."
+                " [Upgrade Now](https://songlabai.com/subscribe/)"
+            )
+            return
+            
+        # Check if server is cold
+        if not check_server_status():
+            status_placeholder.warning(
+                "🔄 Server needs to warm up. This may take up to 7 minutes for free users.\n\n"
+                "💡 Premium users enjoy much faster startup times!"
+            )
+            
+            # Wait for server to warm up with progress display
+            server_ready = wait_for_server_warmup(
+                status_placeholder,
+                progress_placeholder,
+                subscription_placeholder,
+                stage_placeholder
+            )
+            
+            if not server_ready:
+                st.error(
+                    "❌ Server initialization timed out.\n\n"
+                    "Consider upgrading to premium for reliable, fast access.\n\n"
+                    "[Upgrade Now](https://songlabai.com/subscribe/)"
                 )
-                time.sleep(current_delay)
-                continue
-
-    # If all retries failed
-    st.error(
-        "❌ Failed to generate audio after multiple attempts.\n\n"
-        "Consider upgrading to premium for better server availability and faster generation times.\n\n"
-        "[Upgrade Now](https://songlabai.com/subscribe/)"
-    )
+                return
+                
+        # Clear warmup messages
+        for placeholder in [status_placeholder, progress_placeholder, subscription_placeholder, stage_placeholder]:
+            placeholder.empty()
+            
+        # Prepare generation request
+        prompt = f"Genre: {genre}, Energy Level: {energy_level}, Tempo: {tempo}, Description: {description}"
+        payload = {"inputs": {"prompt": prompt, "duration": duration}}
+        api_headers = get_api_headers()
+        
+        # Make the generation request
+        with st.spinner("🎵 Generating your music..."):
+            response = time_post_request(API_URL, headers=api_headers, payload=payload, timeout=600)
+            
+            if response and response.status_code == 200:
+                st.success("✨ Music generated successfully!")
+                load_and_play_generated_audio(response)
+                
+                # Update generation count
+                update_generation_url = "https://songlabai.com/wp-json/custom-api/v1/update-generation-count"
+                requests.post(update_generation_url, headers=auth_headers)
+            else:
+                st.error(
+                    "❌ Failed to generate audio.\n\n"
+                    "This might be due to high server load. Premium users experience fewer failures.\n\n"
+                    "[Upgrade Now](https://songlabai.com/subscribe/)"
+                )
+                
+    except Exception as e:
+        st.error(
+            f"An error occurred: {str(e)}\n\n"
+            "Please try again or consider upgrading for more reliable service."
+        )
+    finally:
+        # Clean up any remaining placeholders
+        for placeholder in [status_placeholder, progress_placeholder, subscription_placeholder, stage_placeholder]:
+            placeholder.empty()
 
 def download_audio():
     if st_state.audio_pydub is None: