Update app.py

app.py

@@ -1,4 +1,4 @@
-# MusicGen + Gradio + GPT Demo App (Optimized for Hugging Face Spaces)
+# MusicGen + Gradio + GPT Demo App (CPU-Optimized for Hugging Face Spaces)
 
 import gradio as gr
 import os
@@ -8,15 +8,18 @@ from transformers import AutoProcessor, MusicgenForConditionalGeneration
 from openai import OpenAI
 import scipy.io.wavfile
 
-# Load MusicGen model
+# Force CPU device (no GPU required)
+device = torch.device("cpu")
+
+# Load MusicGen model onto CPU
 model_name = "facebook/musicgen-small"
-model = MusicgenForConditionalGeneration.from_pretrained(model_name)
+model = MusicgenForConditionalGeneration.from_pretrained(model_name).to(device)
 processor = AutoProcessor.from_pretrained(model_name)
 
-# Initialize OpenAI client with Hugging Face environment variable
+# Initialize OpenAI client (set OPENAI_API_KEY in HF Spaces Secrets)
 client = OpenAI(api_key=os.getenv("OPENAI_API_KEY"))
 
-# Function to enhance user input prompt
+# Refine user prompt via GPT
 def refine_prompt(user_input):
     completion = client.chat.completions.create(
         model="gpt-4",
@@ -27,10 +30,11 @@ def refine_prompt(user_input):
     )
     return completion.choices[0].message.content.strip()
 
-# Function to generate music
-def generate_music(prompt):
-    inputs = processor(text=[prompt], return_tensors="pt")
-    audio_values = model.generate(**inputs, max_new_tokens=256)
+# Generate music (shorter tokens for CPU speed)
+def generate_music(prompt, max_new_tokens: int = 128):
+    inputs = processor(text=[prompt], return_tensors="pt").to(device)
+    # Warning: Generation on CPU may be slow
+    audio_values = model.generate(**inputs, max_new_tokens=max_new_tokens)
     sampling_rate = model.config.audio_encoder.sampling_rate
     audio = audio_values[0].cpu().numpy()
 
@@ -38,7 +42,7 @@ def generate_music(prompt):
     audio = audio / np.max(np.abs(audio))
     audio = audio.astype(np.float32)
 
-    # Save as .wav file for optional download (not used in UI here)
+    # Save as .wav file (in /tmp for Spaces)
     int_audio = (audio * 32767).astype(np.int16)
     scipy.io.wavfile.write("/tmp/output.wav", sampling_rate, int_audio)
 
@@ -50,21 +54,24 @@ def main(user_input):
     sampling_rate, audio = generate_music(detailed_prompt)
     return detailed_prompt, (sampling_rate, audio)
 
-# Gradio interface
-def build_ui():
-    with gr.Blocks() as demo:
-        gr.Markdown("""# 🎵 AI Music Generator\nEnter a music idea or mood and get a short AI-generated track.""")
-
-        user_input = gr.Textbox(label="Describe the mood or style of music")
-        generate_btn = gr.Button("Generate Music")
+# Build Gradio UI
+with gr.Blocks() as demo:
+    gr.Markdown("""# 🎵 AI Music Generator  
+Enter a music idea or mood and get a short AI-generated track. (CPU mode)""")
 
-        refined_output = gr.Textbox(label="Enhanced Prompt by GPT")
-        audio_output = gr.Audio(label="Generated Audio", type="numpy")
+    user_input = gr.Textbox(label="Describe the mood or style of music")
+    max_tokens = gr.Slider(32, 256, value=128, step=32, label="Length (tokens) for CPU")
+    generate_btn = gr.Button("Generate Music")
 
-        generate_btn.click(main, inputs=user_input, outputs=[refined_output, audio_output])
+    refined_output = gr.Textbox(label="Enhanced Prompt by GPT")
+    audio_output = gr.Audio(label="Generated Audio", type="numpy")
+    download_wav = gr.File(label="Download .wav file", file_name="generated.wav")
 
-    return demo
+    generate_btn.click(
+        lambda inp, tok: (main(inp)[0], (main(inp)[1][0], main(inp)[1][1]), "/tmp/output.wav"),
+        inputs=[user_input, max_tokens],
+        outputs=[refined_output, audio_output, download_wav]
+    )
 
-# Launch app in SSR mode (better for Spaces)
-demo = build_ui()
-demo.launch()
+# Launch in SSR mode
+demo.launch(server_name="0.0.0.0", server_port=7860, share=False, enable_queue=True)