Update app.py

app.py

@@ -1,17 +1,33 @@
+import gradio as gr
+import torch
+import numpy as np
+import random
+import os
+import subprocess
+import scipy.io.wavfile as wavfile
+from transformers import MusicgenForConditionalGeneration, AutoProcessor
+from pydub import AudioSegment, effects
+from pedalboard import Pedalboard, Compressor, Gain, HighpassFilter, LowShelfFilter
+from pedalboard.io import AudioFile
+
+# 1. BASH SETUP (Runs once on startup)
+if os.path.exists("setup.sh"):
+    subprocess.run(["sh", "setup.sh"])
+
+# 2. LOAD MODEL
+device = "cuda" if torch.cuda.is_available() else "cpu"
+model = MusicgenForConditionalGeneration.from_pretrained("facebook/musicgen-small").to(device)
+processor = AutoProcessor.from_pretrained("facebook/musicgen-small")
+
 def apply_audacity_fixes(sampling_rate, audio_data, bass_boost_db, fade_sec):
-    """
-    1. Saves Raw AI Audio
-    2. Masters with Pedalboard (Audacity/Studio Style)
-    3. Exports with Bash/FFmpeg (Insane Quality)
-    """
     temp_raw = "raw_output.wav"
     temp_mastered = "mastered_temp.wav"
     
-    # --- STEP 1: Save Raw AI Output ---
+    # STEP 1: Save Raw AI Output
     audio_norm = np.clip(audio_data, -1.0, 1.0)
     wavfile.write(temp_raw, sampling_rate, (audio_norm * 32767).astype(np.int16))
     
-    # --- STEP 2: Audacity-Style Mastering (Pedalboard) ---
+    # STEP 2: Audacity-Style Mastering (Pedalboard)
     with AudioFile(temp_raw) as f:
         audio_signal = f.read(f.frames)
         sr = f.sample_rate
@@ -25,24 +41,62 @@ def apply_audacity_fixes(sampling_rate, audio_data, bass_boost_db, fade_sec):
 
     mastered_signal = board(audio_signal, sr)
     
-    # Save the mastered signal to a temporary wav
     with AudioFile(temp_mastered, 'w', sr, mastered_signal.shape[0]) as f:
         f.write(mastered_signal)
 
-    # --- STEP 3: Apply Fades (Pydub) ---
+    # STEP 3: Apply Fades (Pydub)
     audio_segment = AudioSegment.from_wav(temp_mastered)
     fade_ms = int(fade_sec * 1000)
     final_pydub = audio_segment.fade_in(fade_ms).fade_out(fade_ms)
     final_pydub.export("final_wav_stage.wav", format="wav")
 
-    # --- STEP 4: BASH EXPORT (FFmpeg Insane Quality) ---
-    # -qscale:a 0 is the highest quality MP3 setting available
+    # STEP 4: BASH EXPORT (Highest Quality MP3)
+    # This creates the 'studio_master.mp3' file
     os.system("ffmpeg -y -i final_wav_stage.wav -codec:a libmp3lame -qscale:a 0 studio_master.mp3")
     
-    # Cleanup
+    # Cleanup temporary files
     for file in [temp_raw, temp_mastered, "final_wav_stage.wav"]:
         if os.path.exists(file): os.remove(file)
 
-    return "studio_master.mp3"
+    return "studio_master.mp3" # ONLY ONE RETURN HERE
+
+def generate_music(prompt, duration, instruments, energy, bass_boost_db, fade_sec):
+    if not prompt: return None, "Enter a prompt!"
+    
+    inst_text = ", ".join(instruments)
+    full_prompt = f"{prompt} with {inst_text}. {energy} energy, studio quality."
+    
+    inputs = processor(text=[full_prompt], padding=True, return_tensors="pt").to(device)
     
-    return "final_output.mp3"
+    with torch.no_grad():
+        audio_values = model.generate(**inputs, max_new_tokens=int(duration * 50), do_sample=True, guidance_scale=4.0)
+    
+    sampling_rate = model.config.audio_encoder.sampling_rate
+    audio_data = audio_values[0, 0].cpu().numpy()
+    
+    # Call the mastering function
+    final_mp3 = apply_audacity_fixes(sampling_rate, audio_data, bass_boost_db, fade_sec)
+    
+    return final_mp3, f"✅ Mastered {inst_text} Track"
+
+# 3. UI LAYOUT
+with gr.Blocks(theme=gr.themes.Soft()) as demo:
+    gr.Markdown("# 🎹 NEURAL STUDIO PRO")
+    with gr.Row():
+        with gr.Column():
+            txt = gr.Textbox(label="Song Description")
+            ins = gr.CheckboxGroup(["Piano", "Drums", "Guitar", "Synth", "Violin"], label="Instruments", value=["Piano"])
+            en = gr.Radio(["Low", "Medium", "High"], label="Energy", value="Medium")
+            dur = gr.Slider(5, 30, value=10, label="Seconds")
+            bass = gr.Slider(0, 10, value=3, label="Bass Boost")
+            fade = gr.Slider(0, 5, value=2, label="Fade")
+            btn = gr.Button("Generate", variant="primary")
+        with gr.Column():
+            aud = gr.Audio(label="Result", type="filepath")
+            out_txt = gr.Textbox(label="Status")
+
+    btn.click(generate_music, [txt, dur, ins, en, bass, fade], [aud, out_txt])
+
+# 4. START (Crucial for Hugging Face)
+if __name__ == "__main__":
+    demo.queue().launch()