Update app.py

app.py

@@ -1,108 +1,213 @@
 #!/usr/bin/env python3
-# GhostAI Music Generator Hugging Face Spaces GPU-Compatible
-
-import spaces  # <--- Must be imported FIRST before torch and CUDA
-
-import os
-import sys
-import gc
-import warnings
-import tempfile
-import torch
-import torchaudio
-import numpy as np
-import gradio as gr
+# GhostAI Music Generator – HF-download version
+import os, sys, gc, time, warnings, random, tempfile
+import torch, torchaudio, numpy as np, gradio as gr, psutil
 from pydub import AudioSegment
+from torch.cuda.amp import autocast
 from audiocraft.models import MusicGen
 from huggingface_hub import login
 
+# ------------------------------------------------------------------ #
+#  ❓ Torch <2.3 shim (adds 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"
+    )
+
 warnings.filterwarnings("ignore")
+os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 
-# Hugging Face token auth
+# ------------------------------------------------------------------ #
+#  🔑 Login to HF (model download)                                   #
+# ------------------------------------------------------------------ #
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
-    sys.exit("ERROR: HF_TOKEN not set.")
+    sys.exit("ERROR: environment variable HF_TOKEN not set in the Space settings.")
 login(HF_TOKEN)
 
-# Device setup
+# ------------------------------------------------------------------ #
+#  🖥  Device setup                                                  #
+# ------------------------------------------------------------------ #
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Running on {device.upper()}")
-
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
-
-# Load MusicGen model explicitly on GPU
-musicgen = MusicGen.get_pretrained("medium")
-musicgen.lm.to(device)
-musicgen.set_generation_params(duration=10)
-
-def clean_resources():
+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
+
+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)
+    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
+    gtr  = f", {gtr} guitar riffs"  if gtr!="none" else dflt_gtr
+    return f"{base}{bass}{gtr}{dr}{syn}{step} at {bpm} BPM."
+
+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 …
+
+# ------------------------------------------------------------------ #
+#  🎚  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                                                     #
+# ------------------------------------------------------------------ #
+def generate_music(prompt,cfg,k,p,temp,
+                   total_len,chunk_len,xfade,
+                   bpm,dr,syn,step,bass,gtr):
 
-@spaces.GPU  # <-- Correct GPU decorator for HF
-def generate_music(prompt, cfg, top_k, top_p, temp, total_len, chunk_len, crossfade):
     if not prompt.strip():
-        return None, "⚠️ Enter a valid prompt."
-
-    sample_rate = musicgen.sample_rate
-    segments = []
-    chunks = max(1, total_len // chunk_len)
-
-    for _ in range(chunks):
-        with torch.no_grad():
-            audio = musicgen.generate(
-                [prompt],
-                temperature=temp,
-                cfg_coef=cfg,
-                top_k=top_k,
-                top_p=top_p,
-                duration=chunk_len,
-                progress=False
-            )[0].cpu().float()
-
-        if audio.dim() == 1:
-            audio = audio.unsqueeze(0).repeat(2, 1)
-        elif audio.shape[0] == 1:
-            audio = audio.repeat(2, 1)
+        return None, "❌ Prompt is empty."
+    if not vram_ok():
+        return None, "❌ Not enough VRAM."
+
+    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   = []
+
+    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)
+        with torch.no_grad(), autocast():
+            audio = musicgen_model.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)
 
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
             torchaudio.save(tmp.name, audio, sample_rate)
-            segment = AudioSegment.from_wav(tmp.name)
+            seg = AudioSegment.from_wav(tmp.name)
         os.unlink(tmp.name)
-        segments.append(segment)
-
-        clean_resources()
+        segments.append(seg)
+        clean(); log(f"after chunk {i+1}")
 
-    final = segments[0]
+    track = segments[0]
     for seg in segments[1:]:
-        final = final.append(seg, crossfade=crossfade)
-
-    final = final[:total_len * 1000].fade_in(1000).fade_out(1000).normalize(headroom=-9.0)
+        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"
-    final.export(out_path, format="mp3", bitrate="128k", tags={"title": "GhostAI Track", "artist": "GhostAI"})
-
+    track.export(out_path, format="mp3", bitrate="128k",
+                 tags={"title":"GhostAI Track","artist":"GhostAI"})
+    log("final"); print(f"⏱  {time.time()-t0:.1f}s total")
     return out_path, "✅ Done!"
 
-demo = gr.Interface(
-    fn=generate_music,
-    inputs=[
-        gr.Textbox(label="Instrumental Prompt"),
-        gr.Slider(1.0, 10.0, value=3.0, step=0.1, label="CFG Scale"),
-        gr.Slider(10, 500, value=250, step=10, label="Top-K"),
-        gr.Slider(0.0, 1.0, value=0.9, step=0.05, label="Top-P"),
-        gr.Slider(0.1, 2.0, value=1.0, step=0.1, label="Temperature"),
-        gr.Radio([30, 60, 90, 120], value=30, label="Length (seconds)"),
-        gr.Slider(5, 15, value=10, step=1, label="Chunk Length (seconds)"),
-        gr.Slider(100, 2000, value=1000, step=100, label="Crossfade (ms)")
-    ],
-    outputs=[
-        gr.Audio(label="Generated Music", type="filepath"),
-        gr.Textbox(label="Status")
-    ],
-    title="👻 GhostAI Music Generator",
-    description="Generate instrumental music using MusicGen Medium model."
-)
+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 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)
+
+    # genre buttons (showing two; add the rest as needed)
+    with gr.Row():
+        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,
+            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 🎼")
+    clr = gr.Button("Clear 🧹")
+    audio_out = 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])
 
 demo.launch(share=False)