Update app.py

app.py

@@ -1,15 +1,28 @@
 #!/usr/bin/env python3
-# GhostAI Music Generator – HF-download version
+"""
+GhostAI Music Generator — ZeroGPU Space
+Streams facebook/musicgen-medium with dynamic GPU bursts.
+"""
+
+# 0️⃣  Import spaces *first* so CUDA isn’t touched beforehand
+import spaces                          # HF ZeroGPU decorator
+
+# 1️⃣  Standard libs
 import os, sys, gc, time, warnings, random, tempfile
-import torch, torchaudio, numpy as np, gradio as gr, psutil
+import numpy as np, psutil
+
+# 2️⃣  Torch (CPU wheels; ZeroGPU migrates tensors when needed)
+import torch, torchaudio
+
+# 3️⃣  Other deps
+import gradio as gr
 from pydub import AudioSegment
-from torch.cuda.amp import autocast
 from audiocraft.models import MusicGen
 from huggingface_hub import login
+from torch.cuda.amp import autocast
 
-# ------------------------------------------------------------------ #
-#  ❓ Torch <2.3 shim (adds get_default_device)                       #
-# ------------------------------------------------------------------ #
+# ────────────────────────────────────────────────────────────────────
+# Torch <2.3 shim (transformers may call torch.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"
@@ -18,61 +31,27 @@ if not hasattr(torch, "get_default_device"):
 warnings.filterwarnings("ignore")
 os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 
-# ------------------------------------------------------------------ #
-#  🔑 Login to HF (model download)                                   #
-# ------------------------------------------------------------------ #
+# 🔑  Authenticate so we can pull the model
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    sys.exit("ERROR: environment variable HF_TOKEN not set in the Space settings.")
+    sys.exit("ERROR: Add HF_TOKEN as a secret in your Space.")
 login(HF_TOKEN)
 
-# ------------------------------------------------------------------ #
-#  🖥  Device setup                                                  #
-# ------------------------------------------------------------------ #
 device = "cuda" if torch.cuda.is_available() else "cpu"
-print(f"Running on {device.upper()}")
-if device == "cuda":
-    print(f"GPU  : {torch.cuda.get_device_name(0)}")
-def clean():
-    if device == "cuda":
-        torch.cuda.empty_cache()
-    gc.collect()
-clean()
-
-# ------------------------------------------------------------------ #
-#  📥 Load MusicGen from HF Hub                                      #
-# ------------------------------------------------------------------ #
-print("Downloading & loading facebook/musicgen-medium …")
-musicgen_model = MusicGen.get_pretrained("facebook/musicgen-medium", device=device)
-musicgen_model.set_generation_params(duration=10, two_step_cfg=False)
-
-sample_rate = musicgen_model.sample_rate
-
-# ------------------------------------------------------------------ #
-#  📊 Helpers                                                       #
-# ------------------------------------------------------------------ #
-def vram_ok(req=3.5):
-    if device != "cuda":
-        return True
-    total = torch.cuda.get_device_properties(0).total_memory / 1024**3
-    free  = total - torch.cuda.memory_allocated() / 1024**3
-    if free < req:
-        print(f"⚠️  Only {free:.2f} GB free (< {req} GB).")
-    return free >= req
+print(f"ZeroGPU detected → initial device is {device.upper()}")
 
-def log(stage=""):
-    if stage: print(f"── {stage} ──")
-    if device == "cuda":
-        a = torch.cuda.memory_allocated() / 1024**3
-        r = torch.cuda.memory_reserved()  / 1024**3
-        print(f"GPU mem  alloc {a:.2f} GB  reserved {r:.2f} GB")
-    print(f"CPU mem  {psutil.virtual_memory().percent}% used")
-
-# ------------------------------------------------------------------ #
-#  🎛 Prompt builders (unchanged)                                    #
-# ------------------------------------------------------------------ #
-def _p(base,bpm,dr,syn,st,bass,gtr,dflt_bass,dflt_gtr,vibe):
-    step = f" with {st}" if st!="none" else vibe.format(bpm=bpm)
+# 📥  Download model from Hub
+print("Loading facebook/musicgen-medium … (first run may take ~6 GB download)")
+musicgen = MusicGen.get_pretrained("facebook/musicgen-medium")
+musicgen.set_generation_params(duration=10, two_step_cfg=False)
+SAMPLE_RATE = musicgen.sample_rate
+
+
+# ╭──────────────────────────────────────────────────────────────╮
+# │ Prompt helpers (kept exactly from your original script)      │
+# ╰──────────────────────────────────────────────────────────────╯
+def _p(base,bpm,dr,syn,st,bass,gtr,dflt_bass,dflt_gtr,flow):
+    step = f" with {st}" if st!="none" else flow.format(bpm=bpm)
     dr   = f", {dr} drums"          if dr!="none" else ""
     syn  = f", {syn} accents"       if syn!="none" else ""
     bass = f", {bass}"              if bass!="none" else dflt_bass
@@ -83,131 +62,129 @@ def set_red_hot_chili_peppers_prompt(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")
-# … keep your other 17 prompt functions exactly as before …
+def set_nirvana_grunge_prompt(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")
+# … include your other genre functions unchanged …
+
+# Audio FX
+def apply_eq(s):   return s.low_pass_filter(8000).high_pass_filter(80)
+def apply_fade(s): return s.fade_in(1000).fade_out(1000)
+
+def log(stage=""):
+    if stage: print(f"── {stage} ──")
+    if torch.cuda.is_available():
+        alloc = torch.cuda.memory_allocated()/1024**3
+        res   = torch.cuda.memory_reserved()/1024**3
+        print(f"GPU mem alloc {alloc:.2f} GB  reserved {res:.2f} GB")
+    print(f"CPU mem {psutil.virtual_memory().percent}% used")
 
-# ------------------------------------------------------------------ #
-#  🎚  Audio post-processing                                         #
-# ------------------------------------------------------------------ #
-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)
 
-# ------------------------------------------------------------------ #
-#  🚀 Generator                                                     #
-# ------------------------------------------------------------------ #
+# ╭──────────────────────────────────────────────────────────────╮
+# │  Core generator — wrapped with @spaces.GPU                   │
+# ╰──────────────────────────────────────────────────────────────╯
+@spaces.GPU
 def generate_music(prompt,cfg,k,p,temp,
                    total_len,chunk_len,xfade,
                    bpm,dr,syn,step,bass,gtr):
 
     if not prompt.strip():
-        return None, "❌ Prompt is empty."
-    if not vram_ok():
-        return None, "❌ Not enough VRAM."
+        return None, "⚠️ Prompt is empty."
 
-    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_len = chunk_len + overlap
-    segments   = []
+    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
+    pieces    = []
 
-    torch.manual_seed(42)
-    np.random.seed(42)
+    torch.manual_seed(42); np.random.seed(42)
 
     t0 = time.time()
     for i in range(n_chunks):
         log(f"before chunk {i+1}")
-        musicgen_model.set_generation_params(
-            duration=render_len,use_sampling=True,
-            top_k=k,top_p=p,temperature=temp,cfg_coef=cfg)
+        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_model.generate([prompt], progress=False)[0]
+            audio = musicgen.generate([prompt], progress=False)[0]
 
         audio = audio.cpu().float()
-        if audio.dim()==1:            audio = audio.repeat(2,1)
-        elif audio.shape[0]==1:       audio = audio.repeat(2,1)
-        elif audio.shape[0]!=2:       audio = audio[:1].repeat(2,1)
+        if audio.dim()==1:      audio = audio.repeat(2,1)
+        elif audio.shape[0]==1: audio = audio.repeat(2,1)
+        elif audio.shape[0]!=2: audio = audio[:1].repeat(2,1)
 
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            torchaudio.save(tmp.name, audio, sample_rate)
+            torchaudio.save(tmp.name, audio, SAMPLE_RATE)
             seg = AudioSegment.from_wav(tmp.name)
         os.unlink(tmp.name)
-        segments.append(seg)
-        clean(); log(f"after chunk {i+1}")
+        pieces.append(seg)
+        torch.cuda.empty_cache(); gc.collect()
+        log(f"after chunk {i+1}")
 
-    track = segments[0]
-    for seg in segments[1:]:
+    track = pieces[0]
+    for seg in pieces[1:]:
         track = track.append(seg, crossfade=xfade)
     track = track[: total_len*1000]
     track = apply_fade(apply_eq(track).normalize(headroom=-9.0))
 
-    out_path = "output_cleaned.mp3"
-    track.export(out_path, format="mp3", bitrate="128k",
+    out_file = "output_cleaned.mp3"
+    track.export(out_file, format="mp3", bitrate="128k",
                  tags={"title":"GhostAI Track","artist":"GhostAI"})
-    log("final"); print(f"⏱  {time.time()-t0:.1f}s total")
-    return out_path, "✅ Done!"
+    log("final"); print(f"Total {time.time()-t0:.1f}s")
+    return out_file, "✅ Done!"
 
 def clear_inputs():
     return ("",3.0,250,0.9,1.0,30,10,1000,
             120,"none","none","none","none","none")
 
-# ------------------------------------------------------------------ #
-#  🎨 Custom CSS (unchanged)                                         #
-# ------------------------------------------------------------------ #
-css = """
-body{background:linear-gradient(135deg,#0A0A0A 0%,#1C2526 100%);color:#E0E0E0;font-family:'Orbitron',sans-serif}
-.header{padding:10px;text-align:center;background:rgba(0,0,0,.9);border-bottom:1px solid #00FF9F}
-"""
+# ╭──────────────────────────────────────────────────────────────╮
+# │  Gradio Blocks UI with your CSS & controls                   │
+# ╰──────────────────────────────────────────────────────────────╯
+css = "body{background:#0A0A0A;color:#E0E0E0;font-family:'Orbitron',sans-serif}"
 
-# ------------------------------------------------------------------ #
-#  🖼  Gradio Blocks UI                                              #
-# ------------------------------------------------------------------ #
 with gr.Blocks(css=css) as demo:
-    gr.HTML('<div class="header"><h1>👻 GhostAI Music Generator</h1></div>')
-    prompt_box = gr.Textbox(label="Instrumental Prompt", lines=4)
+    gr.HTML("<h1 style='text-align:center'>👻 GhostAI Music Generator</h1>")
+    prompt = gr.Textbox(label="Prompt", lines=4)
 
-    # genre buttons (showing two; add the rest as needed)
     with gr.Row():
-        gr.Button("RHCP 🌶️").click(
-            set_red_hot_chili_peppers_prompt,
+        gr.Button("RHCP 🌶️").click(set_red_hot_chili_peppers_prompt,
             inputs=[gr.State(120),"none","none","none","none","none"],
-            outputs=prompt_box
-        )
-        gr.Button("Nirvana 🎸").click(
-            set_nirvana_grunge_prompt,
+            outputs=prompt)
+        gr.Button("Nirvana 🎸").click(set_nirvana_grunge_prompt,
             inputs=[gr.State(120),"none","none","none","none","none"],
-            outputs=prompt_box
-        )
-
-    # parameter sliders
-    cfg_scale  = gr.Slider(1,10,3,label="CFG Scale")
-    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,step=0.1,label="Temperature")
-    total_len  = gr.Radio([30,60,90,120],value=30,label="Length (s)")
-    chunk_len  = gr.Slider(5,15,10,step=1,label="Chunk (s)")
-    crossfade  = gr.Slider(100,2000,1000,step=100,label="Cross-fade (ms)")
-
-    bpm        = gr.Slider(60,180,120,label="BPM")
-    drum       = gr.Dropdown(["none","standard rock","funk groove","techno kick","jazz swing"],value="none",label="Drum")
-    synth      = gr.Dropdown(["none","analog synth","digital pad","arpeggiated synth"],value="none",label="Synth")
-    step       = gr.Dropdown(["none","syncopated steps","steady steps","complex steps"],value="none",label="Steps")
-    bass       = gr.Dropdown(["none","slap bass","deep bass","melodic bass"],value="none",label="Bass")
-    gtr        = gr.Dropdown(["none","distorted","clean","jangle"],value="none",label="Guitar")
-
-    gen = gr.Button("Generate 🎼")
+            outputs=prompt)
+        # add more genre buttons here …
+
+    cfg   = gr.Slider(1,10,3,label="CFG")
+    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,step=0.1,label="Temperature")
+    length= gr.Radio([30,60,90,120],value=30,label="Length (s)")
+    chunk = gr.Slider(5,15,10,step=1,label="Chunk (s)")
+    xfade = gr.Slider(100,2000,1000,step=100,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")
+    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")
+    gtr   = gr.Dropdown(["none","distorted","clean","jangle"],"none","Guitar")
+
+    gen = gr.Button("Generate 🎶")
     clr = gr.Button("Clear 🧹")
-    audio_out = gr.Audio(type="filepath")
-    status    = gr.Textbox(interactive=False)
+    audio = gr.Audio(type="filepath")
+    status= gr.Textbox(interactive=False)
 
     gen.click(generate_music,
-        inputs=[prompt_box,cfg_scale,top_k,top_p,temp,
-                total_len,chunk_len,crossfade,
-                bpm,drum,synth,step,bass,gtr],
-        outputs=[audio_out,status])
-    clr.click(clear_inputs, None,
-        [prompt_box,cfg_scale,top_k,top_p,temp,
-         total_len,chunk_len,crossfade,
-         bpm,drum,synth,step,bass,gtr])
+        inputs=[prompt,cfg,top_k,top_p,temp,length,chunk,xfade,
+                bpm,drum,synth,steps,bass,gtr],
+        outputs=[audio,status])
+    clr.click(clear_inputs,None,
+        [prompt,cfg,top_k,top_p,temp,length,chunk,xfade,
+         bpm,drum,synth,steps,bass,gtr])
 
 demo.launch(share=False)