Update app.py

app.py

@@ -1,55 +1,153 @@
-import gradio as gr
-import subprocess
+#!/usr/bin/env python3
+# Minimal Muse-0.6b HF Space App (EXPERIMENTAL)
+
 import os
-import uuid
-
-OUTPUT_DIR = "outputs"
-os.makedirs(OUTPUT_DIR, exist_ok=True)
-
-def generate_music(prompt):
-    uid = str(uuid.uuid4())[:8]
-    out_wav = os.path.join(OUTPUT_DIR, f"{uid}.wav")
-
-    cmd = [
-        "python",
-        "scripts/infer.py",
-        "--model", "bolshyC/Muse-0.6b",
-        "--prompt", prompt,
-        "--out", out_wav
-    ]
-
-    try:
-        subprocess.run(
-            cmd,
-            check=True,
-            timeout=600
-        )
-    except Exception as e:
-        return None, f"生成失败：{e}"
+import sys
+import tempfile
+import torch
+import gradio as gr
+import numpy as np
+import torchaudio
+from huggingface_hub import snapshot_download
+from transformers import AutoTokenizer, AutoModelForCausalLM
 
-    if not os.path.exists(out_wav):
-        return None, "生成失败：未找到音频文件"
+# -----------------------------------------------------------------------------
+# Basic config
+# -----------------------------------------------------------------------------
 
-    return out_wav, "生成完成"
+MODEL_ID = "bolshyC/Muse-0.6b"
+DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+SAMPLE_RATE = 48000
 
-with gr.Blocks() as demo:
-    gr.Markdown("# 🎵 Muse Music Generator (HF Space)")
-    gr.Markdown("使用 Muse-0.6b 在 HF Space 上生成音乐（CLI 接入）")
+# Force HF cache to writable dir
+os.environ["HF_HOME"] = "/tmp/hf"
+os.environ["TRANSFORMERS_CACHE"] = "/tmp/hf/transformers"
+
+# -----------------------------------------------------------------------------
+# Load Muse language model
+# -----------------------------------------------------------------------------
+
+print("Loading Muse-0.6b...")
+tokenizer = AutoTokenizer.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True
+)
+
+model = AutoModelForCausalLM.from_pretrained(
+    MODEL_ID,
+    trust_remote_code=True,
+    torch_dtype=torch.float16 if DEVICE == "cuda" else torch.float32,
+    device_map="auto" if DEVICE == "cuda" else None,
+)
+model.eval()
+print("Muse loaded.")
+
+# -----------------------------------------------------------------------------
+# Load MuCodec + AudioLDM (VERY HEAVY)
+# -----------------------------------------------------------------------------
+
+print("Loading MuCodec / AudioLDM...")
+
+sys.path.insert(0, "./MuCodec")
+from MuCodec.model import PromptCondAudioDiffusion
+from MuCodec.tools.get_melvaehifigan48k import build_pretrained_models
+
+# Download AudioLDM to /tmp
+audioldm_dir = snapshot_download(
+    "haoheliu/audioldm_48k",
+    local_dir="/tmp/audioldm",
+    local_dir_use_symlinks=False
+)
+audioldm_path = os.path.join(audioldm_dir, "audioldm_48k.pth")
+
+vae, stft = build_pretrained_models(audioldm_path)
+vae = vae.to(DEVICE).eval()
+
+mucodec = PromptCondAudioDiffusion(
+    num_channels=32,
+    unet_model_name=None,
+    unet_model_config_path="./MuCodec/configs/models/transformer2D.json",
+    snr_gamma=None,
+)
+
+ckpt = torch.load("./MuCodec/ckpt/mucodec.pt", map_location="cpu")
+mucodec.load_state_dict(ckpt, strict=False)
+mucodec = mucodec.to(DEVICE).eval()
+
+print("MuCodec loaded.")
 
-    prompt = gr.Textbox(
-        label="Prompt",
-        placeholder="A calm piano melody, slow tempo"
+# -----------------------------------------------------------------------------
+# Helpers
+# -----------------------------------------------------------------------------
+
+def extract_audio_tokens(text: str):
+    if "<|audio_0|>" not in text:
+        return None
+    start = text.find("<|audio_0|>") + len("<|audio_0|>")
+    end = text.find("<|audio_1|>")
+    tokens = [int(x) for x in text[start:end].split() if x.isdigit()]
+    if not tokens:
+        return None
+    return torch.tensor(tokens).unsqueeze(0).unsqueeze(0)
+
+# -----------------------------------------------------------------------------
+# Generation
+# -----------------------------------------------------------------------------
+
+def generate(prompt):
+    if not prompt.strip():
+        return None, "Empty prompt"
+
+    inputs = tokenizer(prompt, return_tensors="pt").to(DEVICE)
+
+    with torch.no_grad():
+        out = model.generate(
+            **inputs,
+            max_new_tokens=1024,
+            do_sample=False
+        )
+
+    text = tokenizer.decode(out[0], skip_special_tokens=False)
+    codes = extract_audio_tokens(text)
+    if codes is None:
+        return None, "Failed to parse audio tokens"
+
+    codes = codes.to(DEVICE)
+
+    # Extremely reduced diffusion steps
+    latents = mucodec.inference_codes(
+        [codes[:, :, :1024]],
+        torch.zeros([1, 32, 1, 32], device=DEVICE),
+        torch.randn(1, 32, 512, 32, device=DEVICE),
+        latent_length=512,
+        first_latent_length=0,
+        additional_feats=[],
+        guidance_scale=1.0,
+        num_steps=10,
+        disable_progress=True,
+        scenario="other_seg"
     )
 
+    mel = vae.decode_first_stage(latents.float())
+    wav = vae.decode_to_waveform(mel)
+
+    with tempfile.NamedTemporaryFile(delete=False, suffix=".wav") as f:
+        out_path = f.name
+
+    torchaudio.save(out_path, torch.from_numpy(wav), SAMPLE_RATE)
+    return out_path, "Done"
+
+# -----------------------------------------------------------------------------
+# UI
+# -----------------------------------------------------------------------------
+
+with gr.Blocks() as demo:
+    gr.Markdown("# Muse-0.6b (Experimental HF Space)")
+    prompt = gr.Textbox(label="Prompt")
     btn = gr.Button("Generate")
-    audio = gr.Audio(label="Output", type="filepath")
-    status = gr.Textbox(label="Status")
+    audio = gr.Audio(type="filepath")
+    status = gr.Textbox()
 
-    btn.click(
-        fn=generate_music,
-        inputs=prompt,
-        outputs=[audio, status]
-    )
+    btn.click(generate, prompt, [audio, status])
 
-if __name__ == "__main__":
-    demo.launch()
+demo.launch()