Update app.py

app.py

@@ -2,71 +2,48 @@ import gradio as gr
 import tempfile
 import torch
 import scipy.io.wavfile as wavfile
-import os
 from transformers import AutoProcessor, MusicgenForConditionalGeneration
 
-# --- CPU optimization setup ---
-torch.set_num_threads(os.cpu_count())          # use all CPU cores
-torch.set_num_interop_threads(4)               # reasonable inter-op threads
-torch.backends.quantized.engine = 'qnnpack'    # enable quantized ops where possible
-
-# --- Model Loading ---
-print("🧠 Loading model (optimized for CPU)...")
+# Load model
 processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model.to(device)
 
-try:
-    # ✅ Try quantized load (if bitsandbytes available)
-    from transformers import BitsAndBytesConfig
-    bnb_config = BitsAndBytesConfig(load_in_8bit=True)
-    model = MusicgenForConditionalGeneration.from_pretrained(
-        "facebook/musicgen-small",
-        quantization_config=bnb_config,
-        device_map="cpu"
-    )
-    print("✅ Using 8-bit quantized model")
-except Exception:
-    # fallback
-    model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small")
-    model = model.to("cpu").to(torch.float16)
-    print("⚙️ Using standard float16 CPU model")
-
-device = "cpu"
-MAX_DURATION = 30  # hard cap for CPU — can do 30s comfortably
-
-# --- Generation Function ---
 def generate_music(prompt, duration):
     if not prompt.strip():
-        return None, "❌ Please enter a music prompt."
-    if duration > MAX_DURATION:
-        return None, f"⚠️ Duration too long for CPU — max allowed is {MAX_DURATION} seconds."
-
-    # Dynamic token scaling (smaller = faster)
-    max_new_tokens = int(128 * (duration / 8))
-    max_new_tokens = min(max_new_tokens, 1024)  # cap for stability
+        return None, "Please enter a prompt."
+    if duration > 60:
+        return None, "❌ Duration too long — max allowed is 60 seconds."
 
+    # Prepare inputs
     inputs = processor(text=[prompt], return_tensors="pt").to(device)
 
+    # Scale tokens with duration (MusicGen ≈ 256 tokens ≈ 8 seconds)
+    max_new_tokens = int(256 * (duration / 8))
+    max_new_tokens = min(max_new_tokens, 2048)  # Safety cap
+
+    # Generate audio
     with torch.no_grad():
         audio = model.generate(**inputs, max_new_tokens=max_new_tokens)
 
     sr = model.config.audio_encoder.sampling_rate
     audio_arr = audio[0, 0].cpu().numpy()
 
+    # Save temp file
     tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
     wavfile.write(tmp.name, rate=sr, data=audio_arr)
 
-    return tmp.name, f"✅ Generated {duration}s of audio on CPU!"
+    return tmp.name, f"✅ Generated {duration}s of audio!"
 
-# --- Gradio UI ---
-with gr.Blocks(title="🎶 MusicGen — CPU Optimized") as demo:
-    gr.Markdown("# 🎵 MusicGen — Text-to-Music (CPU Fast Mode)")
+with gr.Blocks(title="MusicGen 🎶") as demo:
+    gr.Markdown("# 🎵 MusicGen — Text-to-Music Generator (Extended 60 s Version)")
     with gr.Row():
         prompt = gr.Textbox(label="🎼 Describe your music", placeholder="e.g. dreamy lo-fi with soft piano")
-        duration = gr.Slider(4, MAX_DURATION, value=10, step=1, label="Duration (seconds)")
-    btn = gr.Button("🎧 Generate Music")
+        duration = gr.Slider(4, 60, value=15, step=1, label="Duration (seconds)")
+    btn = gr.Button("Generate 🎧")
     audio_out = gr.Audio(label="🎶 Output", type="filepath")
     msg = gr.Textbox(label="Status", interactive=False)
     btn.click(generate_music, inputs=[prompt, duration], outputs=[audio_out, msg])
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=int(os.getenv("PORT", 7860)))
+demo.launch(share=True)