Update app.py

app.py

@@ -15,6 +15,8 @@ from io import BytesIO
 from PIL import Image
 import zipfile
 import datetime
+import librosa
+import librosa.display
 
 # === Helper Functions ===
 def audiosegment_to_array(audio):
@@ -95,7 +97,7 @@ def apply_vocal_isolation(audio_path):
     sources = apply_model(model, wav[None])[0]
     wav += ref[:, None]
 
-    vocal_track = sources[3].cpu()  # index 3 = vocals
+    vocal_track = sources[3].cpu()
     out_path = os.path.join(tempfile.gettempdir(), "vocals.wav")
     save_track(out_path, vocal_track, model.samplerate)
     return out_path
@@ -112,7 +114,7 @@ def stem_split(audio_path):
     for i, name in enumerate(['drums', 'bass', 'other', 'vocals']):
         path = os.path.join(output_dir, f"{name}.wav")
         save_track(path, sources[i].cpu(), model.samplerate)
-        stem_paths.append((path, name))
+        stem_paths.append(path)
 
     return stem_paths
 
@@ -146,21 +148,33 @@ if not preset_choices:
 
 preset_names = list(preset_choices.keys())
 
-# === Waveform Generator ===
+# === Waveform + Spectrogram Generator ===
 def show_waveform(audio_file):
-    try:
-        audio = AudioSegment.from_file(audio_file)
-        samples = np.array(audio.get_array_of_samples())
-        plt.figure(figsize=(10, 2))
-        plt.plot(samples[:10000], color="blue")
-        plt.axis("off")
-        buf = BytesIO()
-        plt.savefig(buf, format="png", bbox_inches="tight", dpi=100)
-        plt.close()
-        buf.seek(0)
-        return Image.open(buf)
-    except Exception as e:
-        return None
+    audio = AudioSegment.from_file(audio_file)
+    samples = np.array(audio.get_array_of_samples())
+    plt.figure(figsize=(10, 2))
+    plt.plot(samples[:10000], color="blue")
+    plt.axis("off")
+    buf = BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", dpi=100)
+    plt.close()
+    buf.seek(0)
+    return Image.open(buf)
+
+def show_spectrogram(audio_file):
+    y, sr = torchaudio.load(audio_file)
+    y_np = y.numpy().flatten()
+    S = librosa.feature.melspectrogram(y=y_np, sr=sr)
+    plt.figure(figsize=(10, 2))
+    librosa.display.specshow(librosa.power_to_db(S, ref=np.max), sr=sr, x_axis='time', y_axis='mel')
+    plt.colorbar(format='%+2.0f dB')
+    plt.title('Spectrogram')
+    plt.tight_layout()
+    buf = BytesIO()
+    plt.savefig(buf, format="png", bbox_inches="tight", dpi=100)
+    plt.close()
+    buf.seek(0)
+    return Image.open(buf)
 
 # === Session Info Export ===
 def generate_session_log(audio_path, effects, isolate_vocals, export_format):
@@ -207,9 +221,10 @@ def process_audio(audio_file, selected_effects, isolate_vocals, preset_name, exp
         final_audio.export(output_path, format=export_format.lower())
 
         waveform_image = show_waveform(output_path)
+        spectrogram_image = show_spectrogram(output_path)
         session_log = generate_session_log(audio_file, effects_to_apply, isolate_vocals, export_format)
 
-        return output_path, waveform_image, session_log
+        return output_path, waveform_image, spectrogram_image, session_log
 
 # === Batch Processing Function ===
 def batch_process_audio(files, selected_effects, isolate_vocals, preset_name, export_format):
@@ -218,7 +233,7 @@ def batch_process_audio(files, selected_effects, isolate_vocals, preset_name, ex
     session_logs = []
 
     for file in files:
-        processed_path, _, log = process_audio(file.name, selected_effects, isolate_vocals, preset_name, export_format)
+        processed_path, _, _, log = process_audio(file.name, selected_effects, isolate_vocals, preset_name, export_format)
         results.append(processed_path)
         session_logs.append(log)
 
@@ -231,20 +246,7 @@ def batch_process_audio(files, selected_effects, isolate_vocals, preset_name, ex
 
     return zip_path
 
-# === Custom Preset Upload Handler ===
-def upload_preset(preset_file):
-    try:
-        with open(preset_file.name, "r") as f:
-            data = json.load(f)
-        if "name" in data and "effects" in data:
-            preset_choices[data["name"]] = data["effects"]
-            return f"✅ Loaded custom preset: {data['name']}"
-        else:
-            return "❌ Invalid preset file"
-    except Exception as e:
-        return f"⚠️ Error loading preset: {str(e)}"
-
-# === Gradio Interface ===
+# === Gradio Interface Setup ===
 effect_options = [
     "Noise Reduction",
     "Compress Dynamic Range",
@@ -258,8 +260,11 @@ effect_options = [
 ]
 
 # === Multi-Tab UI ===
-with gr.Blocks(title="AI Audio Studio") as demo:
-    gr.Markdown("## 🎧 AI Audio Studio\nUpload, edit, export — all powered by AI")
+with gr.Blocks(title="AI Audio Studio", css="style.css") as demo:
+    gr.Markdown("""
+    # 🎧 AI Audio Studio – Powered by Hugging Face & Demucs
+    Upload, edit, and export audio with AI-powered tools.
+    """)
 
     # ----- Single File Studio Tab -----
     with gr.Tab("🎵 Single File Studio"):
@@ -275,11 +280,14 @@ with gr.Blocks(title="AI Audio Studio") as demo:
             outputs=[
                 gr.Audio(label="Processed Audio", type="filepath"),
                 gr.Image(label="Waveform Preview"),
+                gr.Image(label="Spectrogram View"),
                 gr.Textbox(label="Session Log (JSON)", lines=5)
             ],
             title="Edit One File at a Time",
-            description="Apply effects, preview waveform, and export as MP3 or WAV",
-            allow_flagging="never"
+            description="Apply effects, preview waveform and spectrogram, and get full session log.",
+            flagging_mode="never",
+            submit_btn="Process Audio",
+            clear_btn=None
         )
 
     # ----- Batch Processing Tab -----
@@ -296,37 +304,51 @@ with gr.Blocks(title="AI Audio Studio") as demo:
             outputs=gr.File(label="Download ZIP of All Processed Files"),
             title="Batch Audio Processor",
             description="Upload multiple files, apply effects in bulk, and download all results in a single ZIP.",
-            flagging_mode="never",        # ✅ New way to disable flagging
+            flagging_mode="never",
             submit_btn="Process All Files",
-            clear_btn=None                # ✅ Instead of False
+            clear_btn=None
         )
 
     # ----- Remix Mode Tab -----
     with gr.Tab("🎛 Remix Mode (Split Stems)"):
         def remix_mode(audio_file):
-            stem_paths = stem_split(audio_file.name)
-            return [path for path, _ in stem_paths], \
-                   [name for _, name in stem_paths]
+            stems = stem_split(audio_file.name)
+            return [gr.File(value=stem) for stem in stems]
 
         gr.Interface(
             fn=remix_mode,
             inputs=gr.Audio(label="Upload Music Track", type="filepath"),
             outputs=[
-                gr.File(label="Stem Files (Vocals, Drums, etc.)"),
-                gr.Textbox(label="Stem Names")
+                gr.File(label="Vocals"),
+                gr.File(label="Drums"),
+                gr.File(label="Bass"),
+                gr.File(label="Other")
             ],
-            title="Split Into Drums, Bass, Vocals",
-            description="Use AI to separate musical elements like vocals, drums, and bass."
+            title="Split Into Drums, Bass, Vocals, and More",
+            description="Use AI to separate musical elements like vocals, drums, and bass.",
+            flagging_mode="never",
+            clear_btn=None
         )
 
-    # ----- Preset Manager Tab -----
-    with gr.Tab("⚙️ Preset Manager"):
+    # ----- Session Info Tab -----
+    with gr.Tab("📝 Session Info"):
+        def get_session_info(audio_file, selected_effects, isolate_vocals, preset_name, export_format):
+            return generate_session_log(audio_file, selected_effects, isolate_vocals, export_format)
+
         gr.Interface(
-            fn=upload_preset,
-            inputs=gr.File(label="Upload Your Own Preset (.json)"),
-            outputs=gr.Textbox(label="Preset Status"),
-            title="Load Custom Presets",
-            description="Upload your own `.json` preset to customize effect chains."
+            fn=get_session_info,
+            inputs=[
+                gr.Audio(label="Upload Audio", type="filepath"),
+                gr.CheckboxGroup(choices=effect_options, label="Apply Effects in Order"),
+                gr.Checkbox(label="Isolate Vocals After Effects"),
+                gr.Dropdown(choices=preset_names, label="Select Preset", value=preset_names[0] if preset_names else None),
+                gr.Dropdown(choices=["MP3", "WAV"], label="Export Format", value="MP3")
+            ],
+            outputs=gr.Textbox(label="Your Session Info (Copy or Save This)", lines=10),
+            title="Save Your Session Settings",
+            description="Get a full log of what was done to your track.",
+            flagging_mode="never",
+            clear_btn=None
         )
 
 demo.launch()