Update app.py

app.py

@@ -13,19 +13,9 @@ import tempfile
 from dotenv import load_dotenv
 import spaces
 
-# Load environment variables
 load_dotenv()
 hf_token = os.getenv("HF_TOKEN")
 
-# Check and enable Xformers for memory-efficient attention
-if torch.cuda.is_available():
-    try:
-        from xformers.ops import memory_efficient_attention
-        os.environ["XFORMERS_ATTENTION"] = "1"
-        print("Xformers is enabled for memory-efficient attention.")
-    except ImportError:
-        print("Xformers is not installed or could not be imported.")
-
 # ---------------------------------------------------------------------
 # Load Llama 3 Pipeline with Zero GPU (Encapsulated)
 # ---------------------------------------------------------------------
@@ -53,6 +43,7 @@ def generate_script(user_prompt: str, model_id: str, token: str):
     except Exception as e:
         return f"Error generating script: {e}"
 
+
 # ---------------------------------------------------------------------
 # Load MusicGen Model (Encapsulated)
 # ---------------------------------------------------------------------
@@ -62,42 +53,49 @@ def generate_audio(prompt: str, audio_length: int):
         musicgen_model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
         musicgen_processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
 
-        musicgen_model.to("cuda")
-        inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt")
+        # Ensure everything is on the same device (GPU or CPU)
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+        musicgen_model.to(device)
+
+        inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt").to(device)
         outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
-        musicgen_model.to("cpu")  # Return the model to CPU
 
-        sr = musicgen_model.config.audio_encoder.sampling_rate
+        # Move outputs to CPU for further processing
         audio_data = outputs[0, 0].cpu().numpy()
-        normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
 
+        # Normalize and save the audio file
+        normalized_audio = (audio_data / max(abs(audio_data)) * 32767).astype("int16")
         output_path = f"{tempfile.gettempdir()}/generated_audio.wav"
-        write(output_path, sr, normalized_audio)
+        write(output_path, musicgen_model.config.audio_encoder.sampling_rate, normalized_audio)
 
         return output_path
     except Exception as e:
         return f"Error generating audio: {e}"
 
+
 # ---------------------------------------------------------------------
 # Gradio Interface Functions
 # ---------------------------------------------------------------------
 def interface_generate_script(user_prompt, llama_model_id):
     return generate_script(user_prompt, llama_model_id, hf_token)
 
+
 def interface_generate_audio(script, audio_length):
     return generate_audio(script, audio_length)
 
+
 # ---------------------------------------------------------------------
 # Interface
 # ---------------------------------------------------------------------
 with gr.Blocks() as demo:
     # Header
-    gr.Markdown("""
-        # 🎙️ AI-Powered Radio Imaging Studio 🚀  
+    gr.Markdown(
+        """
+        # 🎧 AI-Powered Radio Imaging Studio 🚀  
         ### Create stunning **radio promos** with **Llama 3** and **MusicGen**  
         🔥 **Zero GPU** integration for efficiency and ease!  
-        🙌 Thanks to the Hugging Face community for supporting this space.
-    """)
+        """
+    )
 
     # Script Generation Section
     gr.Markdown("## ✍️ Step 1: Generate Your Promo Script")
@@ -109,43 +107,46 @@ with gr.Blocks() as demo:
             info="Describe your promo idea clearly to generate a creative script."
         )
         llama_model_id = gr.Textbox(
-            label="🎛️ Llama 3 Model ID",
+            label="🎿 Llama 3 Model ID",
             value="meta-llama/Meta-Llama-3-8B-Instruct",
             info="Enter the Hugging Face model ID for Llama 3."
         )
     generate_script_button = gr.Button("Generate Script ✨")
     script_output = gr.Textbox(
-        label="📜 Generated Promo Script",
+        label="🖌️ Generated Promo Script",
         lines=4,
         interactive=False,
         info="Your generated promo script will appear here."
     )
 
     # Audio Generation Section
-    gr.Markdown("## 🎧 Step 2: Generate Audio from Your Script")
+    gr.Markdown("## 🎵 Step 2: Generate Audio from Your Script")
     with gr.Row():
         audio_length = gr.Slider(
-            label="🎵 Audio Length (tokens)",
+            label="🎶 Audio Length (tokens)",
             minimum=128,
             maximum=1024,
             step=64,
             value=512,
-            info="Select the desired audio token length."
+            info="Select the desired audio token length."  
         )
     generate_audio_button = gr.Button("Generate Audio 🎶")
     audio_output = gr.Audio(
-        label="🎶 Generated Audio File",
+        label="🎵 Generated Audio File",
         type="filepath",
         interactive=False
     )
 
     # Footer
-    gr.Markdown("""
+    gr.Markdown(
+        """
         <br><hr>
         <p style="text-align: center; font-size: 0.9em;">
             Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
         </p>
-    """, elem_id="footer")
+        """,
+        elem_id="footer"
+    )
 
     # Button Actions
     generate_script_button.click(