Update app.py

app.py

@@ -1,291 +1,187 @@
 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
 
-# ================= 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 (FINAL, POLISHED)
+# -----------------------------
 css = """
-@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;
+:root {
+  --bg-1: #1a0c12;
+  --bg-2: #3a1523;
+  --surface: rgba(255,255,255,0.035);
+  --surface-soft: rgba(255,255,255,0.025);
+  --border: rgba(255,255,255,0.08);
+  --text: #f5f5f5;
+  --muted: rgba(255,255,255,0.65);
+  --accent: #ff5c8a;
 }
 
-/* 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;
+body {
+  background: linear-gradient(135deg, var(--bg-1), var(--bg-2));
+  font-family: Poppins, system-ui, sans-serif;
 }
 
-/* wrapper */
-.app {
-  max-width: 1160px;
+.gradio-container {
+  max-width: 1200px !important;
   margin: 0 auto;
-  padding: 48px 40px;
-  box-sizing: border-box;
-  display: grid;
-  grid-template-rows: auto 1fr;
-  gap: 36px;
 }
 
-/* logo + tagline (left-anchored brand lockup) */
-.brand {
-  display:flex;
-  flex-direction:column;
-  gap:6px;
+/* Header */
+.header {
+  margin-bottom: 28px;
 }
+
 .logo {
-  font-family: Anton, sans-serif;
-  font-size:46px;
-  letter-spacing:1px;
+  font-size: 42px;
+  font-weight: 800;
+  letter-spacing: 1px;
+  color: var(--text);
 }
+
 .tagline {
-  font-size:14px;
-  color:var(--accent);
-  opacity:0.9;
-  text-transform: none; /* keep exact caps as user wanted */
+  margin-top: 6px;
+  font-size: 14px;
+  color: var(--accent);
+  letter-spacing: 0.3px;
 }
 
-/* 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 text */
+.section-title {
+  font-size: 20px;
+  font-weight: 600;
+  margin-bottom: 4px;
 }
 
-/* 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;
+.section-sub {
+  font-size: 14px;
+  color: var(--muted);
+  margin-bottom: 14px;
 }
 
-/* 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;
+/* Main grid — THIS FIXES ALIGNMENT */
+.main-grid {
+  display: grid;
+  grid-template-columns: 2fr 1fr 1fr;
+  grid-template-rows: auto auto;
+  gap: 22px;
+  align-items: stretch;
 }
 
-/* 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);
+/* Upload area */
+.upload-box {
+  grid-row: 1 / span 2;
+  background: var(--surface);
+  border: 1px solid var(--border);
+  border-radius: 14px;
+  padding: 20px;
+  display: flex;
+  flex-direction: column;
+  justify-content: space-between;
 }
 
-/* RIGHT COLUMN (stems) - keep fixed width, grid with 2 columns */
-.stems {
-  display: grid;
-  grid-template-columns: 1fr 1fr;
-  gap: 22px;
-  align-items: start;
+.upload-box .gradio-audio {
+  background: transparent !important;
+  border: none !important;
 }
 
-/* each stem slot now reads like a subdued surface */
-.stem-surface {
-  background: var(--panel-2);
+/* Stem boxes */
+.stem {
+  background: var(--surface-soft);
+  border: 1px solid var(--border);
   border-radius: 14px;
-  padding: 12px;
-  border: 1px solid rgba(255,255,255,0.03);
-  box-sizing: border-box;
-  min-height: 140px; /* consistent stem surface height */
+  padding: 14px;
   display: flex;
   flex-direction: column;
-  gap: 10px;
-  justify-content: center;
+  justify-content: space-between;
 }
 
-/* stem label: same system rhythm as body (no competing font) */
-.stem-label {
-  font-size:13px;
-  font-weight:500;
+.stem-title {
+  font-size: 14px;
+  font-weight: 500;
   color: var(--accent);
-  margin-left: 4px;
+  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;
-}
-
-/* smaller inner audio player to avoid odd spacing */
-.stem-surface .gradio-audio {
+.stem .gradio-audio {
   background: transparent !important;
-  border: 0 !important;
-  padding: 0 !important;
-  min-height: 88px !important;
-  display:flex;
-  align-items:center;
-  justify-content:center;
+  border: none !important;
 }
 
-/* ensure the internal audio element doesn't expand */
-.stem-surface audio {
-  width: 92%;
-  max-height: 36px;
-  opacity: 0.95;
+/* Button */
+.action-row {
+  margin-top: 18px;
 }
 
-/* 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;
+.primary-btn {
+  width: 100%;
+  height: 48px;
+  border-radius: 12px;
+  background: linear-gradient(135deg, #ff7aa2, #ffb36b);
+  color: #111;
+  font-weight: 700;
+  letter-spacing: 0.3px;
 }
 
-/* small responsive fallback */
-@media (max-width: 980px) {
-  .main { grid-template-columns: 1fr; }
-  .stems { grid-template-columns: 1fr 1fr; gap:16px; }
-  .app { padding: 28px 20px; }
+/* Footer cleanup */
+footer {
+  display: none !important;
 }
 """
 
-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])
-
-if __name__ == "__main__":
-    demo.launch(css=css, theme=gr.themes.Base())
+# -----------------------------
+# APP
+# -----------------------------
+with gr.Blocks(css=css) as demo:
+
+    # Header
+    gr.HTML("""
+    <div class="header">
+      <div class="logo">SWARA STUDIO</div>
+      <div class="tagline">Separating Music Into Its elements</div>
+    </div>
+    """)
+
+    # Section 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 aligned grid
+    with gr.HTML('<div class="main-grid">'):
+
+        # Upload box (left, spans 2 rows)
+        with gr.HTML('<div class="upload-box">'):
+            audio_in = gr.Audio(label="Audio", type="filepath")
+        gr.HTML('</div>')
+
+        # Vocals
+        with gr.HTML('<div class="stem">'):
+            gr.HTML('<div class="stem-title">Vocals</div>')
+            vocals = gr.Audio(label=None)
+        gr.HTML('</div>')
+
+        # Drums
+        with gr.HTML('<div class="stem">'):
+            gr.HTML('<div class="stem-title">Drums</div>')
+            drums = gr.Audio(label=None)
+        gr.HTML('</div>')
+
+        # Bass
+        with gr.HTML('<div class="stem">'):
+            gr.HTML('<div class="stem-title">Bass</div>')
+            bass = gr.Audio(label=None)
+        gr.HTML('</div>')
+
+        # Other
+        with gr.HTML('<div class="stem">'):
+            gr.HTML('<div class="stem-title">Other</div>')
+            other = gr.Audio(label=None)
+        gr.HTML('</div>')
+
+    gr.HTML('</div>')  # end grid
+
+    # Action button
+    with gr.Row(elem_classes="action-row"):
+        separate = gr.Button("Separate", elem_classes="primary-btn")
+
+demo.launch()