Update app.py

app.py

@@ -1,4 +1,5 @@
-
+#!/usr/bin/env python3
+# GhostAI Music Generator for Hugging Face Spaces
 import os, sys, gc, time, warnings, tempfile
 import torch, torchaudio, numpy as np, gradio as gr
 from pydub import AudioSegment
@@ -6,66 +7,90 @@ from audiocraft.models import MusicGen
 from huggingface_hub import login
 
 warnings.filterwarnings("ignore")
-os.environ["PYTORCH_CUDA_ALLOC_CONF"] = "max_split_size_mb:128"
 
-# Auth setup
+# Hugging Face token authentication
 HF_TOKEN = os.getenv("HF_TOKEN")
 if not HF_TOKEN:
     sys.exit("ERROR: HF_TOKEN not set.")
 login(HF_TOKEN)
 
-# Device setup
+# Simple GPU check suitable for Hugging Face
 device = "cuda" if torch.cuda.is_available() else "cpu"
 print(f"Running on {device.upper()}")
-torch.cuda.empty_cache(); gc.collect()
 
-# Load model correctly
-musicgen = MusicGen.get_pretrained("medium", device=device)
-musicgen.set_generation_params(duration=10, two_step_cfg=False)
+def clean_resources():
+    if device == "cuda":
+        torch.cuda.empty_cache()
+    gc.collect()
+
+clean_resources()
 
+# Load MusicGen model explicitly on correct device
+print("Loading MusicGen 'medium' model...")
+musicgen = MusicGen.get_pretrained("medium")
+musicgen.lm.to(device)
+musicgen.set_generation_params(duration=10)
+
+# Core generation logic
 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 = []
-
-    for _ in range(total_len // chunk_len):
+    segments, sr = [], musicgen.sample_rate
+    chunks = max(1, total_len // chunk_len)
+    
+    for _ in range(chunks):
         with torch.no_grad():
-            audio = musicgen.generate([prompt])[0].cpu().float()
+            audio = musicgen.generate(
+                [prompt],
+                progress=False,
+                temperature=temp,
+                cfg_coef=cfg,
+                top_k=top_k,
+                top_p=top_p,
+                duration=chunk_len
+            )[0].cpu().float()
+
+        if audio.dim() == 1:
+            audio = audio.unsqueeze(0).repeat(2, 1)
+        elif audio.shape[0] == 1:
+            audio = audio.repeat(2, 1)
 
-        audio = audio if audio.dim() == 2 else audio.repeat(2, 1)
         with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as tmp:
-            torchaudio.save(tmp.name, audio, sample_rate)
-            seg = AudioSegment.from_wav(tmp.name)
+            torchaudio.save(tmp.name, audio, sr)
+            segment = AudioSegment.from_wav(tmp.name)
         os.unlink(tmp.name)
-        segments.append(seg)
-        torch.cuda.empty_cache(); gc.collect()
+        segments.append(segment)
+
+        clean_resources()
 
+    # Concatenate audio segments
     final = 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(-9.0)
+
+    final = final[:total_len * 1000].fade_in(1000).fade_out(1000).normalize(headroom=-9.0)
 
     out_path = "output_cleaned.mp3"
     final.export(out_path, format="mp3", bitrate="128k", tags={"title": "GhostAI Track", "artist": "GhostAI"})
-    return out_path, "✅ Done!"
 
-# Gradio Interface
+    return out_path, "✅ Music Generation Complete!"
+
+# Simple Gradio Interface
 demo = gr.Interface(
     fn=generate_music,
     inputs=[
         gr.Textbox(label="Instrumental Prompt"),
-        gr.Slider(1.0, 10.0, value=3.0, label="CFG Scale"),
-        gr.Slider(10, 500, value=250, label="Top-K"),
-        gr.Slider(0.0, 1.0, value=0.9, label="Top-P"),
-        gr.Slider(0.1, 2.0, value=1.0, label="Temperature"),
-        gr.Radio([30, 60, 90, 120], value=30, label="Length (s)"),
-        gr.Slider(5, 15, value=10, label="Chunk Length (s)"),
-        gr.Slider(100, 2000, value=1000, label="Crossfade (ms)")
+        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"),
+        gr.Audio(label="Generated Music", type="filepath"),
         gr.Textbox(label="Status")
     ],
     title="👻 GhostAI Music Generator",