Update app.py

app.py

@@ -13,13 +13,23 @@ 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)
 # ---------------------------------------------------------------------
-@spaces.GPU(duration=300)
+@spaces.GPU(duration=300)  # GPU allocation for 300 seconds
 def generate_script(user_prompt: str, model_id: str, token: str):
     try:
         tokenizer = AutoTokenizer.from_pretrained(model_id, use_auth_token=token)
@@ -43,7 +53,6 @@ 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)
 # ---------------------------------------------------------------------
@@ -56,7 +65,7 @@ def generate_audio(prompt: str, audio_length: int):
         musicgen_model.to("cuda")
         inputs = musicgen_processor(text=[prompt], padding=True, return_tensors="pt")
         outputs = musicgen_model.generate(**inputs, max_new_tokens=audio_length)
-        musicgen_model.to("cpu")
+        musicgen_model.to("cpu")  # Return the model to CPU
 
         sr = musicgen_model.config.audio_encoder.sampling_rate
         audio_data = outputs[0, 0].cpu().numpy()
@@ -69,18 +78,15 @@ def generate_audio(prompt: str, audio_length: int):
     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
 # ---------------------------------------------------------------------
@@ -90,7 +96,7 @@ with gr.Blocks() as demo:
         # 🎙️ AI-Powered Radio Imaging Studio 🚀  
         ### Create stunning **radio promos** with **Llama 3** and **MusicGen**  
         🔥 **Zero GPU** integration for efficiency and ease!  
-        ❤️ A huge thanks to the **Hugging Face community** for making this possible.
+        🙌 Thanks to the Hugging Face community for supporting this innovation.
     """)
 
     # Script Generation Section
@@ -126,19 +132,18 @@ with gr.Blocks() as demo:
             value=512,
             info="Select the desired audio token length."
         )
-        generate_audio_button = gr.Button("Generate Audio 🎶")
-        audio_output = gr.Audio(
-            label="🎶 Generated Audio File",
-            type="filepath",
-            interactive=False
-        )
+    generate_audio_button = gr.Button("Generate Audio 🎶")
+    audio_output = gr.Audio(
+        label="🎶 Generated Audio File",
+        type="filepath",
+        interactive=False
+    )
 
     # Footer
     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>  
-            Special thanks to the <strong>Hugging Face community</strong> for their incredible support and tools!
+            Created with ❤️ by <a href="https://bilsimaging.com" target="_blank">bilsimaging.com</a>
         </p>
     """, elem_id="footer")