diff --git a/starter_ai_agents/ai_music_generator_agent/music_generator_agent.py b/starter_ai_agents/ai_music_generator_agent/music_generator_agent.py
index 9687958..10c1e30 100644
--- a/starter_ai_agents/ai_music_generator_agent/music_generator_agent.py
+++ b/starter_ai_agents/ai_music_generator_agent/music_generator_agent.py
@@ -1,6 +1,5 @@
 import os
 from uuid import uuid4
-
 import requests
 from agno.agent import Agent, RunResponse
 from agno.models.openai import OpenAIChat
@@ -23,9 +22,9 @@ if openai_api_key and models_lab_api_key:
     agent = Agent(
         name="ModelsLab Music Agent",
         agent_id="ml_music_agent",
-        model=OpenAIChat(id="gpt-4o", api_key=openai_api_key),  # Pass OpenAI API key here
+        model=OpenAIChat(id="gpt-4o", api_key=openai_api_key),
         show_tool_calls=True,
-        tools=[ModelsLabTools(api_key=models_lab_api_key, wait_for_completion=True, file_type=FileType.MP3)],  # Pass ModelsLab API key here
+        tools=[ModelsLabTools(api_key=models_lab_api_key, wait_for_completion=True, file_type=FileType.MP3)],
         description="You are an AI agent that can generate music using the ModelsLabs API.",
         instructions=[
             "When generating music, use the `generate_media` tool with detailed prompts that specify:",
@@ -40,10 +39,7 @@ if openai_api_key and models_lab_api_key:
         debug_mode=True,
     )
 
-    # Disable the button if either API key is missing
-    disable_button = not openai_api_key or not models_lab_api_key
-
-    if st.button("Generate Music", disabled=disable_button):
+    if st.button("Generate Music"):
         if prompt.strip() == "":
             st.warning("Please enter a prompt first.")
         else:
@@ -55,19 +51,31 @@ if openai_api_key and models_lab_api_key:
                         save_dir = "audio_generations"
                         os.makedirs(save_dir, exist_ok=True)
 
-                        # Download the generated audio
                         url = music.audio[0].url
                         response = requests.get(url)
+
+                        # 🛡️ Validate response
+                        if not response.ok:
+                            st.error(f"Failed to download audio. Status code: {response.status_code}")
+                            st.stop()
+
+                        content_type = response.headers.get("Content-Type", "")
+                        if "audio" not in content_type:
+                            st.error(f"Invalid file type returned: {content_type}")
+                            st.write("🔍 Debug: Downloaded content was not an audio file.")
+                            st.write("🔗 URL:", url)
+                            st.stop()
+
+                        # ✅ Save audio
                         filename = f"{save_dir}/music_{uuid4()}.mp3"
                         with open(filename, "wb") as f:
                             f.write(response.content)
 
-                        # Streamlit audio player
+                        # 🎧 Play audio
                         st.success("Music generated successfully! 🎶")
                         audio_bytes = open(filename, "rb").read()
                         st.audio(audio_bytes, format="audio/mp3")
 
-                        # Optional download button
                         st.download_button(
                             label="Download Music",
                             data=audio_bytes,
@@ -80,6 +88,6 @@ if openai_api_key and models_lab_api_key:
                 except Exception as e:
                     st.error(f"An error occurred: {e}")
                     logger.error(f"Streamlit app error: {e}")
+
 else:
-    # Show warning if keys are not entered
     st.sidebar.warning("Please enter both the OpenAI and ModelsLab API keys to use the app.")