import os import tempfile import gradio as gr import spaces import torch # ============================================================================= # Both models are loaded and run LOCALLY inside the ZeroGPU allocation. # Neither is served through HF's Inference Providers, so InferenceClient # can't reach either one — that's why this uses heartlib / diffusers directly. # # requirements.txt needs: # gradio # spaces # torch # huggingface_hub # soundfile # heartlib @ git+https://github.com/HeartMuLa/heartlib.git # diffusers @ git+https://github.com/huggingface/diffusers.git # AceStepPipeline isn't in a release yet # transformers # accelerate # # packages.txt needs: # ffmpeg # ============================================================================= from huggingface_hub import snapshot_download MODEL_HEARTMULA = "HeartMuLa" MODEL_ACESTEP = "ACE-Step 1.5" HEARTMULA_CKPT_DIR = os.environ.get("HEARTMULA_CKPT_DIR", "./ckpt-heartmula") # NOTE: the diffusers-loadable checkpoint lives under a different repo than # the raw "ACE-Step/Ace-Step1.5" release. This is the "turbo" (guidance-distilled, # 8-step) variant — fastest for research iteration. "base"/"sft" variants exist # too (ACE-Step/acestep-v15-base, ACE-Step/acestep-v15-sft) and use real CFG. ACESTEP_REPO = "ACE-Step/acestep-v15-xl-turbo-diffusers" _heartmula_pipe = None _acestep_pipe = None # --- HeartMuLa ----------------------------------------------------------- def _get_heartmula_pipeline(): global _heartmula_pipe if _heartmula_pipe is None: from heartlib import HeartMuLaGenPipeline if not os.path.exists(os.path.join(HEARTMULA_CKPT_DIR, "gen_config.json")): os.makedirs(HEARTMULA_CKPT_DIR, exist_ok=True) snapshot_download("HeartMuLa/HeartMuLaGen", local_dir=HEARTMULA_CKPT_DIR) snapshot_download( "HeartMuLa/HeartMuLa-oss-3B-happy-new-year", local_dir=os.path.join(HEARTMULA_CKPT_DIR, "HeartMuLa-oss-3B"), ) snapshot_download( "HeartMuLa/HeartCodec-oss-20260123", local_dir=os.path.join(HEARTMULA_CKPT_DIR, "HeartCodec-oss"), ) _heartmula_pipe = HeartMuLaGenPipeline.from_pretrained( HEARTMULA_CKPT_DIR, device="cuda", # heartlib takes per-component dtype (matches its CLI's separate # --mula_dtype / --codec_dtype flags), not a single torch_dtype dtype={"mula": torch.bfloat16, "codec": torch.float32}, version="3B", lazy_load=True, ) return _heartmula_pipe def _run_heartmula(prompt: str, lyrics: str, tags: str, duration: int): import soundfile as sf pipe = _get_heartmula_pipeline() # NOTE: verify current kwarg names against examples/run_music_generation.py # in https://github.com/HeartMuLa/heartlib — these mirror its CLI flags. audio, sample_rate = pipe.generate( lyrics=lyrics, tags=tags, max_audio_length_ms=duration * 1000, topk=50, temperature=1.0, cfg_scale=1.5, ) tmp_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name sf.write(tmp_path, audio, sample_rate) return tmp_path # --- ACE-Step 1.5 --------------------------------------------------------- def _get_acestep_pipeline(): global _acestep_pipe if _acestep_pipe is None: from diffusers import AceStepPipeline _acestep_pipe = AceStepPipeline.from_pretrained( ACESTEP_REPO, torch_dtype=torch.bfloat16, ).to("cuda") return _acestep_pipe def _run_acestep(prompt: str, lyrics: str, tags: str, duration: int): import soundfile as sf pipe = _get_acestep_pipeline() output = pipe( prompt=prompt, lyrics=lyrics or "", audio_duration=float(duration), # turbo checkpoint: 8 steps is the intended default, guidance_scale # is ignored automatically for turbo weights (per diffusers docs) num_inference_steps=8, guidance_scale=7.0, shift=3.0, generator=torch.Generator(device="cuda").manual_seed(0), ) # confirmed shape: audios[0] is (channels, samples); transpose for soundfile audio = output.audios[0].T.cpu().float().numpy() sample_rate = pipe.sample_rate tmp_path = tempfile.NamedTemporaryFile(suffix=".wav", delete=False).name sf.write(tmp_path, audio, sample_rate) return tmp_path # --- Shared entry point, still the only thing wrapped in @spaces.GPU ----- @spaces.GPU(duration=180) def generate_music(model_choice: str, prompt: str, lyrics: str, tags: str, duration: int): try: if model_choice == MODEL_HEARTMULA: path = _run_heartmula(prompt, lyrics, tags, duration) msg = "Generated with HeartMuLa-oss-3B." else: path = _run_acestep(prompt, lyrics, tags, duration) msg = "Generated with ACE-Step 1.5." return path, msg except Exception as e: return None, f"Generation error ({model_choice}): {e}" def build_ui(): with gr.Blocks() as demo: gr.Markdown( "# Music Generation Research (ZeroGPU)\n" "Compare HeartMuLa and ACE-Step 1.5, both running locally on a " "ZeroGPU-allocated GPU (no Inference API involved)." ) with gr.Row(): with gr.Column(scale=2): model_choice = gr.Radio( choices=[MODEL_HEARTMULA, MODEL_ACESTEP], value=MODEL_HEARTMULA, label="Model", ) prompt = gr.Textbox( label="Prompt / style description (used by ACE-Step 1.5)", placeholder="An upbeat synthwave track with driving drums", lines=2, ) lyrics = gr.Textbox( label="Lyrics (used by both models)", placeholder="[Verse]\n...\n[Chorus]\n...", lines=8, ) tags = gr.Textbox( label="Tags — HeartMuLa only, comma-separated, no spaces", placeholder="piano,happy,synthesizer,romantic", visible=True, ) duration = gr.Slider(5, 60, step=5, value=15, label="Duration (seconds)") generate_button = gr.Button("Generate audio") with gr.Column(scale=1): audio_output = gr.Audio(label="Generated audio", type="filepath") status = gr.Textbox(label="Status", interactive=False) # Only HeartMuLa uses the tags field — hide it for ACE-Step to avoid confusion def _toggle_tags(choice): return gr.update(visible=(choice == MODEL_HEARTMULA)) model_choice.change(fn=_toggle_tags, inputs=model_choice, outputs=tags) generate_button.click( fn=generate_music, inputs=[model_choice, prompt, lyrics, tags, duration], outputs=[audio_output, status], ) return demo demo = build_ui() if __name__ == "__main__": demo.launch()