Update app.py

app.py

@@ -1,154 +1,291 @@
 import gradio as gr
+import torch
+import torch.nn.functional as F
+import librosa
+import soundfile as sf
+import numpy as np
+from huggingface_hub import hf_hub_download
 
-# -----------------------------
-# CSS (Gradio 6 compatible)
-# -----------------------------
+# ================= MODEL (unchanged) =================
+try:
+    from bs_roformer import BSRoformer
+    from attend import Attend
+except ImportError:
+    pass
+
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+def safe_attend_forward(self, q, k, v, mask=None):
+    return F.scaled_dot_product_attention(q, k, v, attn_mask=mask, dropout_p=0.0)
+
+try:
+    Attend.forward = safe_attend_forward
+except Exception:
+    pass
+
+def load_model():
+    ckpt = hf_hub_download(
+        repo_id="Tachyeon/IAM-RoFormer-Model-Weights",
+        filename="v11_consensus_epoch_30.pt"
+    )
+    model = BSRoformer(
+        dim=512, depth=12, stereo=True, num_stems=4,
+        time_transformer_depth=1, freq_transformer_depth=1,
+        flash_attn=True
+    ).to(DEVICE)
+
+    state = torch.load(ckpt, map_location=DEVICE)
+    model.load_state_dict(state["model"] if "model" in state else state)
+    model.eval()
+    return model
+
+model = load_model()
+
+def separate_audio(path):
+    if not path:
+        return [None]*4
+
+    mix, sr = librosa.load(path, sr=44100, mono=False)
+    if mix.ndim == 1:
+        mix = np.stack([mix, mix])
+
+    x = torch.tensor(mix).float().to(DEVICE)[None]
+    L = x.shape[-1]
+
+    out = torch.zeros(1,4,2,L, device=DEVICE)
+    cnt = torch.zeros_like(out)
+
+    chunk = 44100*10
+    hop = chunk - 44100
+
+    with torch.no_grad(), torch.autocast("cuda", enabled=torch.cuda.is_available()):
+        for s in range(0, L, hop):
+            e = min(s+chunk, L)
+            part = x[:,:,s:e]
+            if part.shape[-1] < chunk:
+                part = F.pad(part,(0,chunk-part.shape[-1]))
+            pred = model(part)
+            out[:,:,:,s:e] += pred[:,:,:,:e-s]
+            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 + layout (polished) =================
 css = """
-:root {
-  --bg1: #1a0c12;
-  --bg2: #3a1523;
-  --panel: rgba(255,255,255,0.04);
-  --panel-soft: rgba(255,255,255,0.03);
-  --border: rgba(255,255,255,0.10);
-  --text: #f4f4f4;
-  --muted: rgba(255,255,255,0.65);
-  --accent: #ff5c8a;
+@import url('https://fonts.googleapis.com/css2?family=Anton&family=Poppins:wght@400;500;600&display=swap');
+
+:root{
+  --bg1:#2a141d;
+  --bg2:#14080d;
+  --ink:#f3ece6;
+  --muted:#b6aeb0;
+  --accent:#ff6f9f;
+  --panel: rgba(255,255,255,0.03);
+  --panel-2: rgba(255,255,255,0.02);
+  --radius: 12px;
 }
 
-body {
-  background: linear-gradient(135deg, var(--bg1), var(--bg2));
-  font-family: Inter, system-ui, sans-serif;
+/* page reset */
+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;
+  overflow: auto;
 }
 
-.gradio-container {
-  max-width: 1200px !important;
+/* wrapper */
+.app {
+  max-width: 1160px;
   margin: 0 auto;
+  padding: 48px 40px;
+  box-sizing: border-box;
+  display: grid;
+  grid-template-rows: auto 1fr;
+  gap: 36px;
 }
 
-/* Header */
-.header {
-  margin-bottom: 26px;
+/* logo + tagline (left-anchored brand lockup) */
+.brand {
+  display:flex;
+  flex-direction:column;
+  gap:6px;
 }
 .logo {
-  font-size: 42px;
-  font-weight: 800;
-  letter-spacing: 1px;
+  font-family: Anton, sans-serif;
+  font-size:46px;
+  letter-spacing:1px;
 }
 .tagline {
-  margin-top: 6px;
-  font-size: 14px;
-  color: var(--accent);
+  font-size:14px;
+  color:var(--accent);
+  opacity:0.9;
+  text-transform: none; /* keep exact caps as user wanted */
 }
 
-/* Section */
-.section-title {
-  font-size: 20px;
-  font-weight: 600;
+/* main grid: left upload (dominant) and right stems column */
+.main {
+  display:grid;
+  grid-template-columns: 1fr 420px; /* fixed right column to keep alignment predictable */
+  gap: 48px;
+  align-items: start; /* ensure top baseline alignment */
 }
-.section-sub {
-  font-size: 14px;
-  color: var(--muted);
-  margin-bottom: 14px;
+
+/* left column content */
+.left h3 { margin: 0; font-size:18px; font-weight:600; }
+.left p { margin:6px 0 18px; color:var(--muted); font-size:13px; }
+
+/* style the Gradio audio drop area so it's visually lighter and slightly shorter */
+.left .gradio-audio {
+  background: var(--panel) !important;
+  border-radius: var(--radius) !important;
+  border: 1px solid rgba(255,255,255,0.04) !important;
+  padding: 12px !important;
+  box-sizing: border-box;
+  min-height: 260px; /* reduced height so stems line up better */
+  display:flex;
+  align-items:center;
+  justify-content:center;
+}
+
+/* tone-down the upload inner texts to make them less dominant */
+.left .gradio-audio .dropzone, .left .gradio-audio .component, .left .gradio-audio .file {
+  color: var(--muted) !important;
+}
+
+/* Make the CTA authoritative but balanced */
+.button-primary {
+  margin-top: 18px;
+  height:46px;
+  width:100%;
+  font-size:15px !important;
+  font-weight:600 !important;
+  background: linear-gradient(90deg,#ff6f9f,#ffbf7a) !important;
+  color: #14080d !important;
+  border-radius: 10px !important;
+  border: none !important;
+  box-shadow: 0 8px 30px rgba(0,0,0,0.25);
 }
 
-/* GRID */
-.main-grid {
+/* RIGHT COLUMN (stems) - keep fixed width, grid with 2 columns */
+.stems {
   display: grid;
-  grid-template-columns: 2fr 1fr 1fr;
-  grid-template-rows: auto auto;
+  grid-template-columns: 1fr 1fr;
   gap: 22px;
+  align-items: start;
 }
 
-/* Upload */
-.upload-box {
-  grid-row: span 2;
-  background: var(--panel);
-  border: 1px solid var(--border);
+/* each stem slot now reads like a subdued surface */
+.stem-surface {
+  background: var(--panel-2);
   border-radius: 14px;
-  padding: 18px;
+  padding: 12px;
+  border: 1px solid rgba(255,255,255,0.03);
+  box-sizing: border-box;
+  min-height: 140px; /* consistent stem surface height */
+  display: flex;
+  flex-direction: column;
+  gap: 10px;
+  justify-content: center;
 }
 
-/* Stems */
-.stem-box {
-  background: var(--panel-soft);
-  border: 1px solid var(--border);
-  border-radius: 14px;
-  padding: 14px;
+/* stem label: same system rhythm as body (no competing font) */
+.stem-label {
+  font-size:13px;
+  font-weight:500;
+  color: var(--accent);
+  margin-left: 4px;
 }
 
-.stem-title {
-  font-size: 14px;
-  font-weight: 500;
-  color: var(--accent);
-  margin-bottom: 6px;
+/* neutralize Gradio's internal "Audio" label visually so it doesn't fight the stem label */
+.stem-surface .gradio-audio label {
+  display:none !important;
 }
 
-/* Audio cleanup */
-.gradio-audio {
+/* smaller inner audio player to avoid odd spacing */
+.stem-surface .gradio-audio {
   background: transparent !important;
-  border: none !important;
+  border: 0 !important;
+  padding: 0 !important;
+  min-height: 88px !important;
+  display:flex;
+  align-items:center;
+  justify-content:center;
 }
 
-/* Button */
-.primary-btn {
-  height: 48px;
-  border-radius: 12px;
-  background: linear-gradient(135deg, #ff7aa2, #ffb36b);
-  color: #111;
-  font-weight: 700;
-  letter-spacing: 0.3px;
+/* ensure the internal audio element doesn't expand */
+.stem-surface audio {
+  width: 92%;
+  max-height: 36px;
+  opacity: 0.95;
 }
 
-footer { display: none !important; }
+/* remove any extra left margin coming from gradio wrapper */
+.gradio-container .wrap { margin: 0; padding: 0; }
+
+/* footer info (small) */
+.footer {
+  margin-top: 28px;
+  text-align: center;
+  color: var(--muted);
+  font-size: 12px;
+  opacity: 0.95;
+}
+
+/* small responsive fallback */
+@media (max-width: 980px) {
+  .main { grid-template-columns: 1fr; }
+  .stems { grid-template-columns: 1fr 1fr; gap:16px; }
+  .app { padding: 28px 20px; }
+}
 """
 
-# -----------------------------
-# APP
-# -----------------------------
 with gr.Blocks() as demo:
+    with gr.Column(elem_classes="app"):
+
+        # Brand lockup (left-anchored)
+        with gr.Row(elem_classes="brand"):
+            gr.HTML('<div class="logo">SWARA STUDIO</div>')
+            # exact user-caps as requested
+            gr.HTML('<div class="tagline">Separating Music Into Its elements</div>')
+
+        # Main content (left upload, right stems). Top baselines match now.
+        with gr.Row(elem_classes="main"):
+            # LEFT: upload + CTA
+            with gr.Column(elem_classes="left"):
+                gr.HTML("""
+                    <h3>Select a track</h3>
+                    <p>We’ll break it down into individual parts</p>
+                """)
+                input_audio = gr.Audio(type="filepath")
+                run_btn = gr.Button("Separate", elem_classes="button-primary")
+
+            # RIGHT: stems grid (fixed width column)
+            with gr.Column():
+                with gr.Row(elem_classes="stems"):
+                    with gr.Column(elem_classes="stem-surface"):
+                        gr.HTML('<div class="stem-label">Vocals</div>')
+                        out_vocals = gr.Audio(interactive=False)
+                    with gr.Column(elem_classes="stem-surface"):
+                        gr.HTML('<div class="stem-label">Drums</div>')
+                        out_drums = gr.Audio(interactive=False)
+                    with gr.Column(elem_classes="stem-surface"):
+                        gr.HTML('<div class="stem-label">Bass</div>')
+                        out_bass = gr.Audio(interactive=False)
+                    with gr.Column(elem_classes="stem-surface"):
+                        gr.HTML('<div class="stem-label">Other</div>')
+                        out_other = gr.Audio(interactive=False)
+
+    # wiring — unchanged logic
+    run_btn.click(separate_audio, input_audio, [out_vocals, out_drums, out_bass, out_other])
 
-    # Header
-    gr.HTML("""
-        <div class="header">
-          <div class="logo">SWARA STUDIO</div>
-          <div class="tagline">Separating Music Into Its elements</div>
-        </div>
-    """)
-
-    # Intro
-    gr.HTML("""
-        <div class="section-title">Select a track</div>
-        <div class="section-sub">We’ll break it down into individual parts</div>
-    """)
-
-    # Main grid
-    with gr.Row(elem_classes="main-grid"):
-
-        # Upload (aligned by grid, not hacks)
-        with gr.Column(elem_classes="upload-box"):
-            audio_in = gr.Audio(label="Audio", type="filepath")
-
-        # Vocals
-        with gr.Column(elem_classes="stem-box"):
-            gr.HTML('<div class="stem-title">Vocals</div>')
-            vocals = gr.Audio(label=None)
-
-        # Drums
-        with gr.Column(elem_classes="stem-box"):
-            gr.HTML('<div class="stem-title">Drums</div>')
-            drums = gr.Audio(label=None)
-
-        # Bass
-        with gr.Column(elem_classes="stem-box"):
-            gr.HTML('<div class="stem-title">Bass</div>')
-            bass = gr.Audio(label=None)
-
-        # Other
-        with gr.Column(elem_classes="stem-box"):
-            gr.HTML('<div class="stem-title">Other</div>')
-            other = gr.Audio(label=None)
-
-    # Action
-    separate = gr.Button("Separate", elem_classes="primary-btn")
-
-demo.launch(css=css)
+if __name__ == "__main__":
+    demo.launch(css=css, theme=gr.themes.Base())