Update app.py

app.py

@@ -4,6 +4,8 @@ import torch.nn.functional as F
 import librosa
 import soundfile as sf
 import numpy as np
+import zipfile
+import os
 from huggingface_hub import hf_hub_download
 
 # ================= MODEL (unchanged) =================
@@ -41,9 +43,10 @@ def load_model():
 
 model = load_model()
 
+# ================= SEPARATION + ZIP =================
 def separate_audio(path):
     if not path:
-        return [None]*4
+        return [None]*5
 
     mix, sr = librosa.load(path, sr=44100, mono=False)
     if mix.ndim == 1:
@@ -69,14 +72,28 @@ def separate_audio(path):
             cnt[:,:,:,s:e] += 1
 
     stems = (out / cnt.clamp(min=1)).cpu().numpy()[0]
-    files=[]
-    for i in range(4):
-        f=f"stem_{i}.wav"
-        sf.write(f, stems[i].T, sr)
-        files.append(f)
-    return files
-
-# ================= UI CSS (minimal enhancement) =================
+
+    stem_files = []
+    names = [
+        "lead_vocals.wav",
+        "mridangam_percussion.wav",
+        "tanpura_drone.wav",
+        "violin_accompaniment.wav"
+    ]
+
+    for i, name in enumerate(names):
+        sf.write(name, stems[i].T, sr)
+        stem_files.append(name)
+
+    # Create ZIP
+    zip_path = "stems.zip"
+    with zipfile.ZipFile(zip_path, "w", zipfile.ZIP_DEFLATED) as z:
+        for f in stem_files:
+            z.write(f)
+
+    return [stem_files[0], stem_files[1], stem_files[2], stem_files[3], zip_path]
+
+# ================= UI (UNCHANGED VISUALLY) =================
 css = """
 @import url('https://fonts.googleapis.com/css2?family=Anton&family=Poppins:wght@400;500;600&display=swap');
 
@@ -94,7 +111,6 @@ css = """
 html, body, .gradio-container {
   height: 100%;
   margin: 0;
-  padding: 0;
   background: linear-gradient(180deg, var(--bg1), var(--bg2)) !important;
   color: var(--ink);
   font-family: Poppins, sans-serif;
@@ -112,7 +128,7 @@ html, body, .gradio-container {
 
 /* brand */
 .brand { display:flex; flex-direction:column; gap:6px; }
-.logo { font-family: Anton, sans-serif; font-size:46px; letter-spacing:1px; }
+.logo { font-family: Anton, sans-serif; font-size:46px; }
 .tagline { font-size:14px; color:var(--accent); opacity:0.9; }
 
 /* layout */
@@ -130,7 +146,6 @@ html, body, .gradio-container {
 .left .gradio-audio {
   background: var(--panel) !important;
   border-radius: var(--radius);
-  border: 1px solid rgba(255,255,255,0.04);
   min-height: 260px;
   display:flex;
   align-items:center;
@@ -171,34 +186,25 @@ html, body, .gradio-container {
   font-size:13px;
   font-weight:600;
   color: var(--accent);
-  letter-spacing: 0.3px;
 }
 
 .stem-info {
   font-size:11px;
   color:var(--muted);
-  line-height:1.35;
   opacity:0.85;
 }
 
 .stem-surface .gradio-audio label { display:none !important; }
 .stem-surface audio { width:92%; max-height:36px; }
-
-@media (max-width: 980px) {
-  .main { grid-template-columns: 1fr; }
-  .stems { grid-template-columns: 1fr 1fr; }
-}
 """
 
 with gr.Blocks() as demo:
     with gr.Column(elem_classes="app"):
 
-        # Brand
         with gr.Row(elem_classes="brand"):
             gr.HTML('<div class="logo">SWARA STUDIO</div>')
             gr.HTML('<div class="tagline">Separating Music Into Its elements</div>')
 
-        # Main
         with gr.Row(elem_classes="main"):
             # LEFT
             with gr.Column(elem_classes="left"):
@@ -208,6 +214,7 @@ with gr.Blocks() as demo:
                 """)
                 input_audio = gr.Audio(type="filepath")
                 run_btn = gr.Button("Separate", elem_classes="button-primary")
+                zip_out = gr.File(label="Download all stems", interactive=False)
 
             # RIGHT
             with gr.Column():
@@ -241,7 +248,11 @@ with gr.Blocks() as demo:
                         """)
                         out_other = gr.Audio(interactive=False)
 
-    run_btn.click(separate_audio, input_audio, [out_vocals, out_drums, out_bass, out_other])
+    run_btn.click(
+        separate_audio,
+        input_audio,
+        [out_vocals, out_drums, out_bass, out_other, zip_out]
+    )
 
 if __name__ == "__main__":
     demo.launch(css=css, theme=gr.themes.Base())