Update app.py

app.py

@@ -1,28 +1,28 @@
 #!/usr/bin/env python3
 """
 GhostAI Music Generator — ZeroGPU Space
-• Pulls facebook/musicgen-medium from the Hub (needs HF_TOKEN)
-• Installs audiocraft==1.3.0 *at runtime* ( --no-deps ) to dodge spaCy/typer clash
-• Decorated with @spaces.GPU so each call gets a transient GPU
+• Each request grabs a transient GPU via Hugging Face ZeroGPU.
+• Model: facebook/musicgen-medium (downloaded with HF_TOKEN secret).
+• Runtime installs audiocraft==1.3.0 --no-deps to dodge spaCy/typer issues.
 """
 
-# ── ZERO-GPU decorator must come BEFORE torch/CUDA ────────────────
-import spaces                  # ① import first
+# ── 0. ZeroGPU must be imported *before* torch is touched ──────────
+import spaces  # <-- DO NOT move this lower
 
-# ── Standard libs ─────────────────────────────────────────────────
+# ── 1. Standard libs ───────────────────────────────────────────────
 import os, sys, gc, time, warnings, tempfile, subprocess, random
 import numpy as np, psutil
 
-# ── Torch (CPU wheels already in container; tensors move to GPU) ──
+# ── 2. Torch stack (torch 2.1.0 & torchaudio 2.1.0 from requirements) ─
 import torch, torchaudio
 
-# ── Other deps ────────────────────────────────────────────────────
+# ── 3. UI / misc deps ─────────────────────────────────────────────
 import gradio as gr
 from pydub import AudioSegment
 from huggingface_hub import login
 from torch.cuda.amp import autocast
 
-# ── tiny shim for transformers on Torch <2.3 ──────────────────────
+# ── 4. Torch <2.3 shim (transformers 4.38 calls get_default_device) ──
 if not hasattr(torch, "get_default_device"):
     torch.get_default_device = lambda: torch.device(
         "cuda" if torch.cuda.is_available() else "cpu"
@@ -31,49 +31,52 @@ if not hasattr(torch, "get_default_device"):
 warnings.filterwarnings("ignore")
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 
-# ── 1. make sure audiocraft is present (no deps) ──────────────────
+# ── 5. Install audiocraft at runtime (no deps) ─────────────────────
 try:
     from audiocraft.models import MusicGen
 except ModuleNotFoundError:
     print("🔧 Installing audiocraft 1.3.0 (no-deps)…")
-    subprocess.check_call([sys.executable, "-m", "pip", "install",
-                           "audiocraft==1.3.0", "--no-deps", "--quiet"])
+    subprocess.check_call(
+        [sys.executable, "-m", "pip", "install", "audiocraft==1.3.0", "--no-deps", "-q"]
+    )
     from audiocraft.models import MusicGen
 
-# ── 2. HF authentication ─────────────────────────────────────────
+# ── 6. Hugging Face authentication ────────────────────────────────
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    sys.exit("ERROR: Please add HF_TOKEN secret in your Space.")
+    sys.exit("ERROR: Please add HF_TOKEN as a secret in your Space.")
 login(HF_TOKEN)
 
-# ── 3. load model from Hub ────────────────────────────────────────
-print("⏬  Loading facebook/musicgen-medium (first run ≈ 6 GB)…")
+# ── 7. Load MusicGen model ────────────────────────────────────────
+print("⏬  Loading facebook/musicgen-medium (first load ≈ 6 GB)…")
 musicgen = MusicGen.get_pretrained("facebook/musicgen-medium")
 musicgen.to(torch.get_default_device())
 musicgen.set_generation_params(duration=10, two_step_cfg=False)
 SR = musicgen.sample_rate
 
-# ── 4. prompt helpers (two shown; add the rest as before) ─────────
-def _p(base,bpm,dr,syn,st,bass,gtr,db,dg,flow):
+# ── 8. Prompt builders (add more in same style) ───────────────────
+def _build(base,bpm,dr,syn,st,bass,gtr,def_bass,def_gtr,flow):
     stxt = f" with {st}" if st!="none" else flow.format(bpm=bpm)
-    return (f"{base}{db if bass=='none' else ', '+bass}"
-            f"{dg if gtr=='none' else ', '+gtr+' guitar riffs'}"
+    return (f"{base}"
+            f"{def_bass if bass=='none' else ', '+bass}"
+            f"{def_gtr if gtr=='none' else ', '+gtr+' guitar riffs'}"
             f"{'' if dr=='none' else ', '+dr+' drums'}"
-            f"{'' if syn=='none' else ', '+syn+' accents'}{stxt} at {bpm} BPM.")
+            f"{'' if syn=='none' else ', '+syn+' accents'}"
+            f"{stxt} at {bpm} BPM.")
 
-def set_red_hot_chili(bpm,dr,syn,st,bass,gtr):
-    return _p("Instrumental funk rock",bpm,dr,syn,st,bass,gtr,
-              ", groovy basslines",", syncopated guitar riffs",
-              "{bpm} BPM funky flow" if bpm>120 else "groovy rhythmic flow")
+def set_rhcp(bpm,dr,syn,st,bass,gtr):
+    return _build("Instrumental funk rock",bpm,dr,syn,st,bass,gtr,
+                  ", groovy basslines",", syncopated guitar riffs",
+                  "{bpm} BPM funky flow" if bpm>120 else "groovy rhythmic flow")
 
-def set_nirvana_grunge(bpm,dr,syn,st,bass,gtr):
-    return _p("Instrumental grunge",bpm,dr,syn,st,bass,gtr,
-              ", melodic basslines",", raw distorted guitar riffs",
-              "{bpm} BPM grungy pulse" if bpm>120 else "grungy rhythmic pulse")
+def set_nirvana(bpm,dr,syn,st,bass,gtr):
+    return _build("Instrumental grunge",bpm,dr,syn,st,bass,gtr,
+                  ", melodic basslines",", raw distorted guitar riffs",
+                  "{bpm} BPM grungy pulse" if bpm>120 else "grungy rhythmic pulse")
 
-# (… add your remaining genre prompt functions unchanged …)
+# (⚠️ Add your remaining genre functions here)
 
-# ── 5. audio FX helpers ───────────────────────────────────────────
+# ── 9. Audio post-FX helpers ──────────────────────────────────────
 def apply_eq(seg):   return seg.low_pass_filter(8000).high_pass_filter(80)
 def apply_fade(seg): return seg.fade_in(1000).fade_out(1000)
 
@@ -83,75 +86,84 @@ def log(stage=""):
         a = torch.cuda.memory_allocated()/1024**3
         r = torch.cuda.memory_reserved()/1024**3
         print(f"GPU alloc {a:.2f} GB  reserved {r:.2f} GB")
-    print(f"CPU {psutil.virtual_memory().percent}%")
+    print(f"CPU {psutil.virtual_memory().percent}% used")
 
-# ── 6. main generation API (ZeroGPU wrapper) ───────────────────────
+# ── 10. Core generation (ZeroGPU decorator) ───────────────────────
 @spaces.GPU
 def generate(prompt,cfg,k,p,temp,
              total_len,chunk_len,xfade,
              bpm,dr,syn,step,bass,gtr):
+
     if not prompt.strip():
         return None, "⚠️ Empty prompt."
 
-    total_len, chunk_len = int(total_len), max(5, min(int(chunk_len), 15))
-    n_chunks  = max(1, total_len // chunk_len)
-    chunk_len = total_len / n_chunks
-    overlap   = min(1.0, xfade / 1000.0)
-    render    = chunk_len + overlap
-    parts     = []
+    total_len  = int(total_len)
+    chunk_len  = max(5, min(int(chunk_len), 15))
+    n_chunks   = max(1, total_len // chunk_len)
+    chunk_len  = total_len / n_chunks
+    overlap    = min(1.0, xfade / 1000.0)
+    render     = chunk_len + overlap
+    segments   = []
+
     torch.manual_seed(42); np.random.seed(42)
 
     for i in range(n_chunks):
-        log(f"chunk {i+1} start")
+        log(f"chunk {i+1}")
         musicgen.set_generation_params(duration=render,use_sampling=True,
                                        top_k=k,top_p=p,temperature=temp,cfg_coef=cfg)
         with torch.no_grad(), autocast():
             audio = musicgen.generate([prompt], progress=False)[0]
+
         audio = audio.cpu().float()
         if audio.dim()==1 or audio.shape[0]==1:
             audio = audio.repeat(2,1)
+
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
             torchaudio.save(tmp.name, audio, SR)
-            parts.append(AudioSegment.from_wav(tmp.name))
+            segments.append(AudioSegment.from_wav(tmp.name))
         os.unlink(tmp.name)
-        torch.cuda.empty_cache(); gc.collect(); log(f"chunk {i+1} done")
+        torch.cuda.empty_cache(); gc.collect()
 
-    track = parts[0]
-    for seg in parts[1:]:
+    track = segments[0]
+    for seg in segments[1:]:
         track = track.append(seg, crossfade=xfade)
-    track = track[:total_len*1000]
+    track = track[: total_len*1000]
     track = apply_fade(apply_eq(track).normalize(headroom=-9.0))
 
-    out_f = "output_cleaned.mp3"
-    track.export(out_f, format="mp3", bitrate="128k",
+    out_path = "output_cleaned.mp3"
+    track.export(out_path, format="mp3", bitrate="128k",
                  tags={"title":"GhostAI Track","artist":"GhostAI"})
-    return out_f, "✅ Done!"
+    return out_path, "✅ Done!"
 
 def clear_inputs():
     return ("",3.0,250,0.9,1.0,30,10,1000,
             120,"none","none","none","none","none")
 
-# ── 7. minimal Gradio 4 UI (add more buttons if needed) ───────────
+# ── 11. Gradio Blocks UI ──────────────────────────────────────────
 with gr.Blocks(css="body{background:#0A0A0A;color:#E0E0E0;font-family:'Orbitron',sans-serif}") as demo:
-    gr.HTML("<h1 style='text-align:center'>👻 GhostAI MusicGen</h1>")
-    prompt = gr.Textbox(lines=4, label="Prompt")
-    with gr.Row():
-        gr.Button("RHCP 🌶️").click(set_red_hot_chili,
-            inputs=[gr.State(120),"none","none","none","none","none"],
-            outputs=prompt)
-        gr.Button("Nirvana 🎸").click(set_nirvana_grunge,
-            inputs=[gr.State(120),"none","none","none","none","none"],
-            outputs=prompt)
+    gr.Markdown("<h1 style='text-align:center'>👻 GhostAI MusicGen</h1>")
+    prompt = gr.Textbox(lines=4, label="Instrumental Prompt")
 
+    with gr.Row():
+        gr.Button("RHCP 🌶️").click(set_rhcp,
+                                   inputs=[gr.State(120),"none","none","none","none","none"],
+                                   outputs=prompt)
+        gr.Button("Nirvana 🎸").click(set_nirvana,
+                                     inputs=[gr.State(120),"none","none","none","none","none"],
+                                     outputs=prompt)
+        # ➕ Add more genre buttons here …
+
+    # Parameter controls
     cfg   = gr.Slider(1,10,3,label="CFG")
-    k     = gr.Slider(10,500,250,step=10,label="Top-K")
-    p     = gr.Slider(0,1,0.9,step=0.05,label="Top-P")
-    temp  = gr.Slider(0.1,2,1,label="Temp")
-    length= gr.Radio([30,60,90,120],value=30,label="Length")
-    chunk = gr.Slider(5,15,10,label="Chunk")
-    xfade = gr.Slider(100,2000,1000,label="Cross-fade")
+    top_k = gr.Slider(10,500,250,step=10,label="Top-K")
+    top_p = gr.Slider(0,1,0.9,step=0.05,label="Top-P")
+    temp  = gr.Slider(0.1,2,1,label="Temperature")
+    length= gr.Radio([30,60,90,120],value=30,label="Length (s)")
+    chunk = gr.Slider(5,15,10,label="Chunk (s)")
+    xfade = gr.Slider(100,2000,1000,label="Cross-fade (ms)")
+
     bpm   = gr.Slider(60,180,120,label="BPM")
-    drum  = gr.Dropdown(["none","standard rock","funk groove","techno kick","jazz swing"],"none","Drum")
+    drum  = gr.Dropdown(["none","standard rock","funk groove","techno kick","jazz swing"],"none","Drums")
     synth = gr.Dropdown(["none","analog synth","digital pad","arpeggiated synth"],"none","Synth")
     steps = gr.Dropdown(["none","syncopated steps","steady steps","complex steps"],"none","Steps")
     bass  = gr.Dropdown(["none","slap bass","deep bass","melodic bass"],"none","Bass")
@@ -159,13 +171,15 @@ with gr.Blocks(css="body{background:#0A0A0A;color:#E0E0E0;font-family:'Orbitron'
 
     gen = gr.Button("Generate 🚀")
     clr = gr.Button("Clear 🧹")
-    out = gr.Audio(type="filepath")
-    status = gr.Textbox(interactive=False)
+    audio_out = gr.Audio(type="filepath")
+    status    = gr.Textbox(interactive=False)
 
     gen.click(generate,
-              [prompt,cfg,k,p,temp,length,chunk,xfade,bpm,drum,synth,steps,bass,gtr],
-              [out,status])
+              inputs=[prompt,cfg,top_k,top_p,temp,length,chunk,xfade,
+                      bpm,drum,synth,steps,bass,gtr],
+              outputs=[audio_out,status])
     clr.click(clear_inputs, None,
-              [prompt,cfg,k,p,temp,length,chunk,xfade,bpm,drum,synth,steps,bass,gtr])
+              [prompt,cfg,top_k,top_p,temp,length,chunk,xfade,
+               bpm,drum,synth,steps,bass,gtr])
 
 demo.launch(share=False)