Update app.py

app.py

@@ -13,6 +13,8 @@ from pathlib import Path
 import matplotlib.pyplot as plt
 from io import BytesIO
 from PIL import Image
+import zipfile
+import datetime
 
 # === Helper Functions ===
 def audiosegment_to_array(audio):
@@ -98,6 +100,22 @@ def apply_vocal_isolation(audio_path):
     save_track(out_path, vocal_track, model.samplerate)
     return out_path
 
+# === Stem Splitting (Drums, Bass, Other, Vocals) ===
+def stem_split(audio_path):
+    model = pretrained.get_model(name='htdemucs')
+    wav = load_track_local(audio_path, model.samplerate, channels=2)
+    sources = apply_model(model, wav[None])[0]
+
+    output_dir = tempfile.mkdtemp()
+    stem_paths = []
+
+    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))
+
+    return stem_paths
+
 # === Preset Loader with Fallback ===
 def load_presets():
     try:
@@ -110,10 +128,8 @@ def load_presets():
                     data = json.load(infile)
                     if "name" in data and "effects" in data:
                         presets[data["name"]] = data["effects"]
-                    else:
-                        print(f"Invalid format in {f}")
             except json.JSONDecodeError:
-                print(f"Failed to parse {f} — invalid JSON")
+                print(f"Invalid JSON: {f}")
         return presets
     except FileNotFoundError:
         print("Presets folder not found")
@@ -146,6 +162,17 @@ def show_waveform(audio_file):
     except Exception as e:
         return None
 
+# === Session Info Export ===
+def generate_session_log(audio_path, effects, isolate_vocals, export_format):
+    log = {
+        "timestamp": str(datetime.datetime.now()),
+        "filename": os.path.basename(audio_path),
+        "effects_applied": effects,
+        "isolate_vocals": isolate_vocals,
+        "export_format": export_format
+    }
+    return json.dumps(log, indent=2)
+
 # === Main Processing Function ===
 def process_audio(audio_file, selected_effects, isolate_vocals, preset_name, export_format):
     audio = AudioSegment.from_file(audio_file)
@@ -178,27 +205,46 @@ def process_audio(audio_file, selected_effects, isolate_vocals, preset_name, exp
 
         output_path = f.name
         final_audio.export(output_path, format=export_format.lower())
-        return output_path, show_waveform(output_path)
 
-# === Batch Processing Function ===
-import zipfile
+        waveform_image = show_waveform(output_path)
+        session_log = generate_session_log(audio_file, effects_to_apply, isolate_vocals, export_format)
+
+        return output_path, waveform_image, session_log
 
+# === Batch Processing Function ===
 def batch_process_audio(files, selected_effects, isolate_vocals, preset_name, export_format):
     output_dir = tempfile.mkdtemp()
     results = []
+    session_logs = []
 
     for file in files:
-        processed_path = process_audio(file.name, selected_effects, isolate_vocals, preset_name, export_format)
-        results.append(processed_path[0])  # Take just the file path
+        processed_path, _, log = process_audio(file.name, selected_effects, isolate_vocals, preset_name, export_format)
+        results.append(processed_path)
+        session_logs.append(log)
 
     zip_path = os.path.join(output_dir, "batch_output.zip")
     with zipfile.ZipFile(zip_path, 'w') as zipf:
         for i, res in enumerate(results):
-            zipf.write(res, f"processed_{i}.{export_format.lower()}")
+            filename = f"processed_{i}.{export_format.lower()}"
+            zipf.write(res, filename)
+            zipf.writestr(f"session_info_{i}.json", session_logs[i])
 
     return zip_path
 
-# === Effect List ===
+# === 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 ===
 effect_options = [
     "Noise Reduction",
     "Compress Dynamic Range",
@@ -215,6 +261,27 @@ effect_options = [
 with gr.Blocks(title="AI Audio Studio") as demo:
     gr.Markdown("## 🎧 AI Audio Studio\nUpload, edit, export — all powered by AI")
 
+    # ----- Single File Studio Tab -----
+    with gr.Tab("🎵 Single File Studio"):
+        gr.Interface(
+            fn=process_audio,
+            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.Audio(label="Processed Audio", type="filepath"),
+                gr.Image(label="Waveform Preview"),
+                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"
+        )
+
     # ----- Batch Processing Tab -----
     with gr.Tab("🔊 Batch Processing"):
         gr.Interface(
@@ -234,24 +301,32 @@ with gr.Blocks(title="AI Audio Studio") as demo:
             clear_btn=False
         )
 
-    # ----- Single File Studio Tab -----
-    with gr.Tab("🎵 Single File Studio"):
+    # ----- 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]
+
         gr.Interface(
-            fn=process_audio,
-            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")
-            ],
+            fn=remix_mode,
+            inputs=gr.Audio(label="Upload Music Track", type="filepath"),
             outputs=[
-                gr.Audio(label="Processed Audio", type="filepath"),
-                gr.Image(label="Waveform Preview")
+                gr.File(label="Stem Files (Vocals, Drums, etc.)"),
+                gr.Textbox(label="Stem Names")
             ],
-            title="Edit One File at a Time",
-            description="Apply effects, preview waveform, and export as MP3 or WAV",
-            allow_flagging="never"
+            title="Split Into Drums, Bass, Vocals",
+            description="Use AI to separate musical elements like vocals, drums, and bass."
+        )
+
+    # ----- Preset Manager Tab -----
+    with gr.Tab("⚙️ Preset Manager"):
+        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."
         )
 
 demo.launch()