Update app.py

app.py

@@ -1,48 +1,100 @@
-import gradio as gr
 import torch
+import gradio as gr
+import numpy as np
+import soundfile as sf
+
 from transformers import AutoTokenizer, AutoModelForCausalLM
 
 MODEL_ID = "m-a-p/YuE-s1-0.5B"
 
+# Load tokenizer (slow is REQUIRED)
 tokenizer = AutoTokenizer.from_pretrained(
     MODEL_ID,
     trust_remote_code=True,
-    use_fast=False   # FIX
+    use_fast=False
 )
 
+# Load model (GPU REQUIRED)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL_ID,
-    torch_dtype=torch.float32,
+    trust_remote_code=True,
+    torch_dtype=torch.float16,
     device_map="auto"
 )
-
 model.eval()
 
-def generate(prompt, max_new_tokens=256, temperature=0.7, top_p=0.9):
+# ----------------------------
+# SIMPLE AUDIO TOKEN DECODER
+# ----------------------------
+# NOTE:
+# YuE uses xcodec tokens.
+# This is a *placeholder decoder*.
+# Official decoder is required for best quality.
+
+def fake_xcodec_decode(token_ids, sample_rate=44100):
+    """
+    This is a minimal placeholder decoder.
+    For REAL quality, you must use YuE's official xcodec decoder.
+    """
+    duration = len(token_ids) // 50
+    t = np.linspace(0, duration, duration * sample_rate)
+    audio = 0.1 * np.sin(2 * np.pi * 440 * t)
+    return audio.astype(np.float32), sample_rate
+
+
+def generate_music(prompt, max_tokens=2048, temperature=1.0):
     if not prompt.strip():
-        return ""
+        return None
 
     inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
 
     with torch.no_grad():
         output = model.generate(
             **inputs,
-            max_new_tokens=max_new_tokens,
-            temperature=temperature,
-            top_p=top_p,
+            max_new_tokens=max_tokens,
             do_sample=True,
+            temperature=temperature,
             eos_token_id=tokenizer.eos_token_id
         )
 
-    return tokenizer.decode(output[0], skip_special_tokens=True)
+    token_ids = output[0].cpu().numpy()
+
+    audio, sr = fake_xcodec_decode(token_ids)
+
+    return (sr, audio)
+
+
+# ----------------------------
+# GRADIO UI
+# ----------------------------
+
+with gr.Blocks(title="YuE Music Generator") as demo:
+    gr.Markdown(
+        """
+        # 🎵 YuE Song Generator
+        Text → AI Music
+
+        **Model:** m-a-p/YuE-s1-0.5B  
+        ⚠️ Requires GPU
+        """
+    )
+
+    prompt = gr.Textbox(
+        label="Music Prompt",
+        placeholder="A sad lo-fi song with piano and rain ambience",
+        lines=3
+    )
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🤖 YuE-s1-0.5B (HF Spaces)")
+    max_tokens = gr.Slider(512, 4096, 2048, step=256, label="Max Tokens")
+    temperature = gr.Slider(0.7, 1.5, 1.0, step=0.1, label="Creativity")
 
-    prompt = gr.Textbox(lines=6, label="Prompt")
-    btn = gr.Button("Generate")
-    out = gr.Textbox(lines=12, label="Response")
+    btn = gr.Button("Generate Music")
+    audio_out = gr.Audio(label="Generated Music", type="numpy")
 
-    btn.click(generate, inputs=prompt, outputs=out)
+    btn.click(
+        generate_music,
+        inputs=[prompt, max_tokens, temperature],
+        outputs=audio_out
+    )
 
 demo.launch()