HeartMuLa Gradio Space deployment

.dockerignore

@@ -0,0 +1,18 @@
+.git
+.github
+.venv
+.env
+__pycache__
+*.pyc
+*.pyo
+*.egg-info
+.pytest_cache
+.coverage
+*.mp3
+*.wav
+*.ogg
+assets/*.mp3
+assets/*.wav
+.gitignore
+README.md
+LICENSE

.github/copilot-instructions.md

@@ -0,0 +1,89 @@
+# GitHub Copilot instructions for heartlib
+
+## What this repo is
+- HeartMuLa music generation stack: converts lyrics + style tags → audio via two-stage pipeline (HeartMuLa LLM → audio tokens, HeartCodec flow-matching codec → waveform).
+- Supports lyrics transcription via Whisper-based HeartTranscriptor.
+- Python package with main entry points: `heartlib.HeartMuLaGenPipeline` and `heartlib.HeartTranscriptorPipeline` (see `src/heartlib/__init__.py`).
+- Examples/reference CLIs in `examples/`; production use: install via `pip install -e .`
+
+## Core architecture & data flow
+**Music generation:** `Inputs(lyrics, tags)` → Tokenizer → HeartMuLa (frame-by-frame token generation) → HeartCodec (flow-matching detokenization) → MP3
+- **HeartMuLa** (LLaMA3.2 backbone, 3B/300M/7B flavors): generates 8 parallel audio codebook streams + 1 prompt guidance stream (9 total, `_parallel_number=9`)
+- **HeartCodec**: VQ codec that reconstructs waveforms from codebook frames in overlapping windows, fixed 48 kHz output
+- **HeartTranscriptor**: Whisper variant fine-tuned for vocal transcription; works on 30-second chunks, batch=16
+
+## Repo map
+- `src/heartlib/pipelines/music_generation.py`: orchestrates tokenization → HeartMuLa inference → HeartCodec detokenize → `torchaudio.save()`
+  - `HeartMuLaGenPipeline.from_pretrained()`: factory with device/dtype/lazy_load config
+  - `_resolve_paths()`: validates checkpoint layout early (hard error if missing)
+  - `_resolve_devices()`: handles scalar/dict device specs; forces `lazy_load=False` for multi-device
+- `src/heartlib/heartmula/modeling_heartmula.py`: backbone (llama3_2_3B/7B/300M factory functions), token generator with CFG support
+- `src/heartlib/heartcodec/modeling_heartcodec.py`: VQ codec + flow-matching decoder, detokenizes `(codebooks, time)` frames
+- `src/heartlib/pipelines/lyrics_transcription.py`: wraps transformers' Whisper; fixed chunk=30s, batch=16
+- `src/heartlib/heartmula/configuration_heartmula.py`, `src/heartlib/heartcodec/configuration_heartcodec.py`: model configs
+
+## Checkpoints & required layout
+Directory structure after downloads (see README or `hf download` commands):
+```
+./ckpt/
+  HeartMuLa-oss-3B/  (or -7B, -300M)
+    config.json
+    model-*.safetensors
+    model.safetensors.index.json
+  HeartCodec-oss/
+    config.json
+    model-*.safetensors
+    model.safetensors.index.json
+  HeartTranscriptor-oss/
+    config.json
+    pytorch_model.bin (or safetensors)
+  tokenizer.json
+  gen_config.json
+```
+- `_resolve_paths(pretrained_path, version)` validates all required files; raises FileNotFoundError if missing
+- Latest checkpoint: HeartMuLa-RL-oss-3B-20260123 (RL-tuned, recommended for style control)
+
+## Generation pipeline behaviors to know
+- **Inputs:** dict with `lyrics`, `tags` (both strings or file paths); auto-lowercased, tags wrapped with `<tag>...</tag>` if missing
+- **Tokenization:** uses `tokenizers.Tokenizer` (from `tokenizer.json`); token IDs from `HeartMuLaGenConfig` (text_bos_id=128000, text_eos_id=128001, audio_eos_id=8193)
+- **CFG (classifier-free guidance):** if `cfg_scale != 1.0`, batch duplicated for unconditional pass (bs becomes 2×); enables style control tradeoff
+- **Audio generation loop:** runs max `max_audio_length_ms // 80` frames (~12.5 Hz generation rate); stops early if any token ≥ audio_eos_id
+- **Memory optimization:**
+  - `lazy_load=True` defers model loading, unloads after generation (saves CUDA between uses)
+  - Forced `lazy_load=False` if mula_device ≠ codec_device (different device types can't swap)
+  - Uses `torch.autocast` with specified dtype to reduce memory footprint
+- **Output:** via `torchaudio.save(save_path, wav, 48000)` at fixed 48 kHz sample rate
+
+## Codec specifics
+- `HeartCodec.detokenize(frames)` expects shape `(codebooks, time)` where codebooks ≤ 8; pads/repeats to uniform length internally
+- Uses flow-matching inference in overlapping windows (reduces boundary artifacts), then scalar decoder → PCM waveform
+- Fixed 48 kHz output; non-standard rates must be resampled post-generation
+- Model config (`config.json`) defines number of codebooks and codec architecture
+
+## Transcription pipeline behaviors
+- `HeartTranscriptorPipeline.from_pretrained(model_path, device, dtype)` wraps `WhisperForConditionalGeneration` from `HeartTranscriptor-oss`
+- Fixed at 30-second chunks, batch size 16; no dynamic chunking
+- Note: trained on separated vocals; best results with source-separated inputs (use demucs or similar pre-pipeline)
+- Supports beam search and temperature kwargs via `__call__()` decoding_kwargs (see `examples/run_lyrics_transcription.py`)
+
+## Dev workflows & commands
+- **Install:** `pip install -e .` (Python ≥3.9, 3.10 recommended; CUDA deps: torch 2.4.1, torchaudio 2.4.1, torchtune 0.4.0, bitsandbytes 0.49.0)
+- **Generate:** `python examples/run_music_generation.py --model_path ./ckpt --version 3B --lyrics ./assets/lyrics.txt --tags ./assets/tags.txt`
+  - Key flags: `--mula_device cuda --codec_device cuda` (or separate); `--lazy_load true` (single-GPU VRAM relief); `--cfg_scale 1.5` (style strength)
+- **Transcribe:** `python examples/run_lyrics_transcription.py --model_path ./ckpt --music_path ./assets/output.mp3`
+- No test suite; validate changes via example scripts
+
+## Coding conventions & critical patterns
+- **Device specs:** `from_pretrained(..., device=X)` accepts `torch.device` (both models→X) or dict `{"mula": dev1, "codec": dev2}` (forces `lazy_load=False`)
+- **Dtype specs:** mirrors device—scalar dtype or dict with `"mula"`, `"codec"` keys
+- **Token/text handling:** always lowercase inputs, auto-wrap tags with `<tag>...</tag>`, append BOS/EOS via tokenizer config (callers depend on this)
+- **Unimplemented:** reference audio path exists but raises `NotImplementedError`; don't add stub without full end-to-end implementation
+- **Generation loop internals:** `tqdm` progress, `torch.autocast` scope—avoid breaking these or model cache setup (`setup_caches()`)
+- **Memory patterns:** properties `self.mula` / `self.codec` lazy-load on first access if `lazy_load=True`, then can unload via `_unload_models()`
+
+## Quick pointers for agents
+- Extend pipelines, not models, unless changing core LLM/codec logic
+- Validate paths early (mirror `_resolve_paths` style) for new entry points
+- Preserve 48 kHz sample rate and codebook count (8) in outputs
+- When modifying tokenization or BOS/EOS logic, verify examples still run end-to-end
+- Device/dtype flexibility is intentional—test multi-GPU configs if changing device dispatch logic

.gitignore

@@ -0,0 +1,18 @@
+__pycache__/
+**.pyc
+checkpoint/
+.DS_Store
+**.wav
+**.mp3
+**.png
+**.jpeg
+**.jpg
+.vscode/
+**.egg-info/
+build/
+.idea/
+ckpt/
+.venv*/
+.env
+models/
+assets/

DEPLOYMENT_GUIDE.md

@@ -0,0 +1,115 @@
+# HeartMuLa HF Space Deployment Guide
+
+## What's Been Set Up
+
+✅ **Gradio Web UI** (`app.py`)
+- Interactive music generation interface
+- Real-time parameter adjustment
+- Audio preview and download
+- Example prompts included
+
+✅ **Docker Environment** (`Dockerfile` + `requirements.txt`)
+- CUDA 12.1 with GPU support
+- All dependencies pre-configured
+- Automatic model downloading on startup
+
+✅ **Space Configuration**
+- `README_SPACE.md` - Space documentation
+- `.dockerignore` - Optimized Docker builds
+
+## Deployment Steps
+
+### 1. Push to Your HF Space Repository
+```bash
+cd f:\Projects\heartlib
+git remote add space https://huggingface.co/spaces/brandongraves08/test
+git push space main
+```
+
+### 2. Configure Space Settings (in HF UI)
+1. Go to https://huggingface.co/spaces/brandongraves08/test/settings
+2. Set **Runtime** to **Docker**
+3. Select **GPU** hardware (A100, T4, or A10 recommended)
+4. Save settings
+
+### 3. Space Will Auto-Deploy
+The Dockerfile will:
+- Install all dependencies
+- Download HeartMuLa and HeartCodec models
+- Start the Gradio app on port 7860
+
+## Features
+
+### Generation Parameters
+- **Lyrics**: Custom lyrics for the song
+- **Tags**: Style descriptors (pop, rock, ambient, etc.)
+- **Duration**: 5-60 seconds
+- **Temperature**: 0.1-2.0 (creativity level)
+- **CFG Scale**: 1.0-3.0 (style control strength)
+- **Top-K**: 10-100 (sampling parameter)
+
+### Model Information
+- **HeartMuLa-RL-oss-3B**: RL-tuned 3B model (recommended)
+- **HeartCodec-oss**: High-fidelity codec (48 kHz)
+- **Inference Speed**: ~RTF 1.0
+
+## Local Testing (Before Deploying)
+
+Test locally with CUDA GPU:
+```bash
+cd f:\Projects\heartlib
+pip install gradio
+python app.py
+```
+
+Then open http://localhost:7860 in your browser.
+
+## Troubleshooting
+
+### Models Not Downloading
+If the Space fails to download models:
+1. Check HF token is configured: `huggingface-cli login`
+2. Verify model access on HF
+3. Check Space logs for download errors
+
+### Out of Memory
+- Reduce duration slider
+- Use smaller model version (if available)
+- Enable lazy_load in pipeline (already done)
+
+### Slow Generation
+- Generation is ~RTF 1.0 (real-time speed)
+- First run may be slower due to model loading
+- CPU is ~10x slower than GPU
+
+## File Structure
+
+```
+heartlib/
+├── app.py                    # Gradio app (main entry point)
+├── Dockerfile                # Docker build config
+├── requirements.txt          # Python dependencies
+├── setup.sh                  # Model download script
+├── README_SPACE.md          # Space documentation
+├── .dockerignore            # Docker build optimization
+├── src/heartlib/            # Core library
+│   ├── pipelines/
+│   │   ├── music_generation.py
+│   │   └── lyrics_transcription.py
+│   ├── heartmula/
+│   └── heartcodec/
+└── examples/                 # Example scripts
+```
+
+## Next Steps
+
+1. **Commit & Push** to your Space repository
+2. **Monitor** the Space build logs
+3. **Test** once deployment completes
+4. **Share** the Space URL!
+
+## Support
+
+- **Paper**: https://arxiv.org/pdf/2601.10547
+- **GitHub**: https://github.com/HeartMuLa/heartlib
+- **Discord**: https://discord.gg/BKXF5FgH

Dockerfile

@@ -0,0 +1,30 @@
+FROM nvidia/cuda:12.1.0-runtime-ubuntu22.04
+
+WORKDIR /app
+
+# Install system dependencies
+RUN apt-get update && apt-get install -y \
+    python3.10 \
+    python3-pip \
+    git \
+    && rm -rf /var/lib/apt/lists/*
+
+# Set Python 3.10 as default
+RUN update-alternatives --install /usr/bin/python python /usr/bin/python3.10 1 && \
+    update-alternatives --install /usr/bin/pip pip /usr/bin/pip3 1
+
+# Copy requirements and install Python dependencies
+COPY requirements.txt .
+RUN pip install --no-cache-dir -r requirements.txt
+
+# Copy application code
+COPY . .
+
+# Install heartlib package
+RUN pip install -e .
+
+# Create models directory
+RUN mkdir -p ./models
+
+# Run setup (downloads models and starts app)
+CMD ["bash", "-c", "bash setup.sh && python app.py"]

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

MANUAL_DEPLOYMENT.md

@@ -0,0 +1,55 @@
+# 🚀 Manual Deployment Steps for HF Space
+
+Since HF requires token authentication, here's how to deploy manually:
+
+## Option 1: Use HF CLI (Recommended)
+1. Get your HF access token: https://huggingface.co/settings/tokens
+2. Create a write-enabled token
+3. Run this command with your token:
+
+```bash
+huggingface-cli login --token hf_YOUR_TOKEN_HERE --add_to_git_credential
+```
+
+Then push:
+```bash
+cd F:\Projects\heartlib
+git remote add space https://huggingface.co/spaces/brandongraves08/test
+git push space main --force
+```
+
+## Option 2: Manual Upload via Web UI
+1. Go to https://huggingface.co/spaces/brandongraves08/test
+2. Click **Files** → **Upload files**
+3. Upload these files:
+   - app.py
+   - Dockerfile
+   - requirements.txt
+   - setup.sh
+   - .dockerignore
+   - README_SPACE.md
+
+## Option 3: Use Git with SSH Key
+1. Set up SSH key on HF: https://huggingface.co/settings/keys
+2. Update git remote:
+```bash
+git remote set-url space git@huggingface.co:spaces/brandongraves08/test.git
+git push space main
+```
+
+## Space Configuration (After Upload)
+Once files are on the Space:
+1. Go to **Settings** (gear icon)
+2. Set **Runtime** → **Docker**
+3. Select **GPU Hardware** (A100/T4/A10G recommended)
+4. Click **Save**
+
+The Space will auto-build and start the Gradio app!
+
+## Verify Deployment
+Check status at: https://brandongraves08-test.hf.space
+
+## Troubleshooting
+- **Models not downloading**: Check HF token has read access
+- **Build fails**: Check Space logs for error details
+- **Out of memory**: Reduce model size or use better GPU

README.md

@@ -0,0 +1,255 @@
+<p align="center">
+    <picture>
+        <source srcset="./assets/logo.png" media="(prefers-color-scheme: dark)">
+        <img src="./assets/logo.png" width="30%">
+    </picture>
+    
+</p>
+
+<p align="center">
+    <a href="https://heartmula.github.io/">Demo 🎶</a> &nbsp;|&nbsp; 📑 <a href="https://arxiv.org/pdf/2601.10547">Paper</a>
+    <br>
+    <a href="https://huggingface.co/HeartMuLa/HeartMuLa-oss-3B">HeartMuLa-oss-3B 🤗</a> &nbsp;|&nbsp; <a href="https://modelscope.cn/models/HeartMuLa/HeartMuLa-oss-3B">HeartMuLa-oss-3B <picture>
+        <source srcset="./assets/badge.svg" media="(prefers-color-scheme: dark)">
+        <img src="./assets/badge.svg" width="20px">
+    </picture></a>
+    <br>
+    <a href="https://huggingface.co/HeartMuLa/HeartMuLa-RL-oss-3B-20260123"> HeartMuLa-RL-oss-3B-20260123 🤗</a> &nbsp;|&nbsp; <a href="https://modelscope.cn/models/HeartMuLa/HeartMuLa-RL-oss-3B-20260123">HeartMuLa-RL-oss-3B-20260123 <picture>
+        <source srcset="./assets/badge.svg" media="(prefers-color-scheme: dark)">
+        <img src="./assets/badge.svg" width="20px">
+    </picture></a>
+    
+</p>
+
+---
+# HeartMuLa: A Family of Open Sourced Music Foundation Models
+
+HeartMuLa is a family of open sourced music foundation models including: 
+1. HeartMuLa: a music language model that generates music conditioned on lyrics and tags with multilingual support including but not limited to English, Chinese, Japanese, Korean and Spanish.
+2. HeartCodec: a 12.5 hz music codec with high reconstruction fidelity;
+3. HeartTranscriptor: a whisper-based model specifically tuned for lyrics transcription; Check [this page](./examples/README.md) for its usage.
+4. HeartCLAP: an audio–text alignment model that establishes a unified embedding space for music descriptions and cross-modal retrieval.
+---
+
+
+Below shows the experiment result of our oss-3B version compared with other baselines.
+<p align="center">
+    <picture>
+        <source srcset="./assets/exp.png" media="(prefers-color-scheme: dark)">
+        <img src="./assets/exp.png" width="90%">
+    </picture>
+    
+</p>
+
+---
+
+## 🔥 Highlight
+
+Our latest internal version of HeartMuLa-7B achieves **comparable performance with Suno** in terms of musicality, fidelity and controllability. If you are interested, welcome to reach us out via heartmula.ai@gmail.com
+
+## 📰 News
+Join on Discord! [<img alt="join discord" src="https://img.shields.io/discord/842440537755353128?color=%237289da&logo=discord"/>](https://discord.gg/BKXF5FgH)
+
+- 🚀 **23 Jan. 2026**
+
+    By leveraging Reinforcement Learning, we have continuously refined our model and are proud to officially release **HeartMuLa-RL-oss-3B-20260123**. This version is designed to achieve more precise control over styles and tags. Simultaneously, we are launching **HeartCodec-oss-20260123**, which optimizes audio decoding quality.
+
+- 🫶 **20 Jan. 2026** 
+    
+    [Benji](https://github.com/benjiyaya) has created a wonderful [ComfyUI custom node](https://github.com/benjiyaya/HeartMuLa_ComfyUI) for HeartMuLa. Thanks Benji!
+- ⚖️ **20 Jan. 2026** 
+
+    License update: We update the license of this repo and all related model weights to **Apache 2.0**.
+- 🚀 **14 Jan. 2026**  
+    The official release of **HeartTranscriptor-oss** and the first **HeartMuLa-oss-3B** version along with our **HeartCodec-oss**.
+
+---
+## 🧭 TODOs
+
+- ⏳ Release scripts for inference acceleration and streaming inference. The current inference speed is around RTF $\approx 1.0$.
+- ⏳ Support **reference audio conditioning**, **fine-grained controllable music generation**, **hot song generation**.
+- ⏳ Release the **HeartMuLa-oss-7B** version.
+- ✅ Release inference code and pretrained checkpoints of  
+  **HeartCodec-oss, HeartMuLa-oss-3B, and HeartTranscriptor-oss**.
+
+---
+
+## 🛠️ Local Deployment
+
+### ⚙️ Environment Setup
+
+We recommend using `python=3.10` for local deployment.
+
+Clone this repo and install locally.
+
+```
+git clone https://github.com/HeartMuLa/heartlib.git
+cd heartlib
+pip install -e .
+```
+
+Download our pretrained checkpoints from huggingface or modelscope using the following command:
+
+```
+# if you are using huggingface
+hf download --local-dir './ckpt' 'HeartMuLa/HeartMuLaGen'
+
+## To use version released on 20260123 (recommended)
+hf download --local-dir './ckpt/HeartMuLa-oss-3B' 'HeartMuLa/HeartMuLa-RL-oss-3B-20260123'
+hf download --local-dir './ckpt/HeartCodec-oss' HeartMuLa/HeartCodec-oss-20260123
+
+## To use oss-3B version
+hf download --local-dir './ckpt/HeartMuLa-oss-3B' 'HeartMuLa/HeartMuLa-oss-3B'
+hf download --local-dir './ckpt/HeartCodec-oss' 'HeartMuLa/HeartCodec-oss'
+
+# if you are using modelscope
+modelscope download --model 'HeartMuLa/HeartMuLaGen' --local_dir './ckpt'
+
+## To use version released on 20260123 (recommended)
+modelscope download --model 'HeartMuLa/HeartMuLa-RL-oss-3B-20260123' --local_dir './ckpt/HeartMuLa-oss-3B'
+modelscope download --model 'HeartMuLa/HeartCodec-oss-20260123' --local_dir './ckpt/HeartCodec-oss'
+
+## To use oss-3B version
+modelscope download --model 'HeartMuLa/HeartMuLa-oss-3B' --local_dir './ckpt/HeartMuLa-oss-3B'
+modelscope download --model 'HeartMuLa/HeartCodec-oss' --local_dir './ckpt/HeartCodec-oss'
+```
+
+After downloading, the `./ckpt` subfolder should structure like this:
+```
+./ckpt/
+├── HeartCodec-oss/
+├── HeartMuLa-oss-3B/
+├── gen_config.json
+└── tokenizer.json
+```
+
+
+### ▶️ Example Usage
+
+To generate music, run:
+
+```
+python ./examples/run_music_generation.py --model_path=./ckpt --version="3B"
+```
+
+By default this command will generate a piece of music conditioned on lyrics and tags provided in `./assets` folder. The output music will be saved at `./assets/output.mp3`.
+
+#### FAQs
+
+1. How to specify lyrics and tags?
+
+    The model will load lyrics from the txt file `--lyrics` link to (by default `./assets/lyrics.txt`). If you would like to use your own lyrics, just modify the content in `./assets/lyrics.txt`. If you would like to save your lyrics to another path, e.g. `my_awesome_lyrics.txt`, remember to input arguments `--lyrics my_awesome_lyrics.txt`.
+
+    For tags it's basically the same.
+
+2. CUDA out of memory?
+
+    If you have multi-GPUs (e.g. 2 4090s), we recommend placing the params of HeartMuLa and HeartCodec separately on different devices. You can do it by typing `--mula_device cuda:0 --codec_device cuda:1`
+
+    If you are running on a single GPU, use `--lazy_load true` so that modules will be loaded on demand and deleted once inference completed to save GPU memory.
+
+All parameters:
+
+- `--model_path` (required): Path to the pretrained model checkpoint
+- `--lyrics`: Path to lyrics file (default: `./assets/lyrics.txt`)
+- `--tags`: Path to tags file (default: `./assets/tags.txt`)
+- `--save_path`: Output audio file path (default: `./assets/output.mp3`)
+- `--max_audio_length_ms`: Maximum audio length in milliseconds (default: 240000)
+- `--topk`: Top-k sampling parameter for generation (default: 50)
+- `--temperature`: Sampling temperature for generation (default: 1.0)
+- `--cfg_scale`: Classifier-free guidance scale (default: 1.5)
+- `--version`: The version of HeartMuLa, choose between [`3B`, `7B`]. (default: `3B`) # `7B` version not released yet.
+- `--mula_device/--codec_device`: The device where params will be placed. Both are set to `cuda` by default. You can use `--mula_device cuda:0 --codec_device cuda:1` to explicitly place different modules to different devices.
+- `--mula_dtype/--codec_dtype`: Inference dtype. By default is `bf16` for HeartMuLa and `fp32` for HeartCodec. Setting `bf16` for HeartCodec may result in the degradation of audio quality.
+- `--lazy_load`: Whether or not to use lazy loading (default: false). If turned on, modules will be loaded on demand to save GPU usage. 
+Recommended format of lyrics and tags:
+```txt
+[Intro]
+
+[Verse]
+The sun creeps in across the floor
+I hear the traffic outside the door
+The coffee pot begins to hiss
+It is another morning just like this
+
+[Prechorus]
+The world keeps spinning round and round
+Feet are planted on the ground
+I find my rhythm in the sound
+
+[Chorus]
+Every day the light returns
+Every day the fire burns
+We keep on walking down this street
+Moving to the same steady beat
+It is the ordinary magic that we meet
+
+[Verse]
+The hours tick deeply into noon
+Chasing shadows,chasing the moon
+Work is done and the lights go low
+Watching the city start to glow
+
+[Bridge]
+It is not always easy,not always bright
+Sometimes we wrestle with the night
+But we make it to the morning light
+
+[Chorus]
+Every day the light returns
+Every day the fire burns
+We keep on walking down this street
+Moving to the same steady beat
+
+[Outro]
+Just another day
+Every single day
+```
+
+Regarding tags, check this [issue](https://github.com/HeartMuLa/heartlib/issues/17) for reference.
+Our different tags are comma-separated without spaces as illustrated below:
+```txt
+piano,happy,wedding,synthesizer,romantic
+```
+
+---
+
+
+## ⚖️ License
+
+This repository is licensed under the Apache 2.0 License.
+
+---
+
+## 📚 Citation
+
+```
+@misc{yang2026heartmulafamilyopensourced,
+      title={HeartMuLa: A Family of Open Sourced Music Foundation Models}, 
+      author={Dongchao Yang and Yuxin Xie and Yuguo Yin and Zheyu Wang and Xiaoyu Yi and Gongxi Zhu and Xiaolong Weng and Zihan Xiong and Yingzhe Ma and Dading Cong and Jingliang Liu and Zihang Huang and Jinghan Ru and Rongjie Huang and Haoran Wan and Peixu Wang and Kuoxi Yu and Helin Wang and Liming Liang and Xianwei Zhuang and Yuanyuan Wang and Haohan Guo and Junjie Cao and Zeqian Ju and Songxiang Liu and Yuewen Cao and Heming Weng and Yuexian Zou},
+      year={2026},
+      eprint={2601.10547},
+      archivePrefix={arXiv},
+      primaryClass={cs.SD},
+      url={https://arxiv.org/abs/2601.10547}, 
+}
+```
+
+## 📬 Contact
+If you are interested in HeartMuLa, feel free to reach us at heartmula.ai@gmail.com
+
+Welcome to join us through [discord](https://discord.gg/BKXF5FgH) or Wechat group.
+
+Scan the QR code on the left to join our Wechat group. If it expires, feel free to raise an issue to remind us of updating. 
+
+If the number of group members exceeds 200, joining the group via directly scanning the QR code is restricted by WeChat. In this case, scan our team member's QR code on the right and send a request writing **HeartMuLa Group Invite**. We will invite you into the group manually.
+<p align="center">
+    <picture>
+        <source srcset="./assets/group_wx.jpeg" media="(prefers-color-scheme: dark)">
+        <img src="./assets/group_wx.jpeg" width="40%">
+    </picture>
+    <picture>
+        <source srcset="./assets/lead_wx.jpeg" media="(prefers-color-scheme: dark)">
+        <img src="./assets/lead_wx.jpeg" width="40%">
+    </picture>
+</p>

README_SPACE.md

@@ -0,0 +1,54 @@
+---
+title: HeartMuLa Music Generation
+emoji: 🎵
+colorFrom: purple
+colorTo: pink
+sdk: docker
+sdk_version: latest
+app_file: app.py
+pinned: true
+duplicated_from: HeartMuLa/HeartMuLa-oss
+---
+
+# HeartMuLa Music Generation Space
+
+Generate music from lyrics and style tags using the HeartMuLa family of open-source music foundation models.
+
+## Features
+
+- **Music Generation**: Convert lyrics + style tags → audio via two-stage pipeline
+- **HeartMuLa LLM**: Frame-by-frame audio token generation with style control
+- **HeartCodec**: High-fidelity flow-matching codec (48 kHz output)
+- **Multiple Model Sizes**: 3B, 7B, and 300M versions available
+
+## Setup
+
+The Space will automatically download and set up the required models on first run.
+
+## Usage
+
+1. Enter your **lyrics** in the text field
+2. Add **style tags** (e.g., "pop, upbeat, energetic")
+3. Adjust generation parameters:
+   - **Duration**: Length of generated music (5-60 seconds)
+   - **Temperature**: Creativity level (0.1-2.0)
+   - **CFG Scale**: Style control strength (1.0-3.0)
+   - **Top-K**: Sampling parameter (10-100)
+4. Click **Generate Music** to create your track
+
+## Model Information
+
+- **HeartMuLa-RL-oss-3B-20260123**: RL-tuned version with improved style control (recommended)
+- **HeartCodec-oss-20260123**: Optimized audio decoding quality
+
+## Performance
+
+- RTF ≈ 1.0 (real-time inference speed)
+- 48 kHz sample rate output
+- Supports multiple languages
+
+## References
+
+- [Paper](https://arxiv.org/pdf/2601.10547)
+- [GitHub](https://github.com/HeartMuLa/heartlib)
+- [Discord](https://discord.gg/BKXF5FgH)

app.py

@@ -0,0 +1,228 @@
+#!/usr/bin/env python3
+"""
+HeartMuLa Music Generation Gradio App for Hugging Face Spaces
+"""
+
+import gradio as gr
+import torch
+import os
+from pathlib import Path
+from heartlib import HeartMuLaGenPipeline
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+# Configuration
+MODEL_PATH = "./models"
+DEFAULT_VERSION = "3B"
+
+# Global pipeline instance
+pipeline = None
+
+
+def load_pipeline():
+    """Load the HeartMuLa pipeline"""
+    global pipeline
+    
+    if pipeline is not None:
+        return pipeline
+    
+    logger.info("Loading HeartMuLa pipeline...")
+    
+    # Determine device
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    dtype = torch.float16 if device == "cuda" else torch.float32
+    
+    logger.info(f"Using device: {device}, dtype: {dtype}")
+    
+    try:
+        pipeline = HeartMuLaGenPipeline.from_pretrained(
+            MODEL_PATH,
+            device={
+                "mula": torch.device(device),
+                "codec": torch.device(device),
+            },
+            dtype={
+                "mula": dtype,
+                "codec": dtype,
+            },
+            version=DEFAULT_VERSION,
+            lazy_load=True,
+        )
+        logger.info("Pipeline loaded successfully!")
+        return pipeline
+    except Exception as e:
+        logger.error(f"Failed to load pipeline: {e}")
+        raise gr.Error(f"Failed to load model: {e}")
+
+
+def generate_music(
+    lyrics: str,
+    tags: str,
+    max_duration: int = 30,
+    temperature: float = 1.0,
+    top_k: int = 50,
+    cfg_scale: float = 1.5,
+):
+    """Generate music from lyrics and tags"""
+    
+    if not lyrics.strip():
+        raise gr.Error("Please enter lyrics")
+    if not tags.strip():
+        raise gr.Error("Please enter tags")
+    
+    try:
+        logger.info(f"Generating music with lyrics: {lyrics[:50]}... and tags: {tags}")
+        
+        # Load pipeline
+        pipe = load_pipeline()
+        
+        # Convert duration to milliseconds
+        max_audio_length_ms = max_duration * 1000
+        
+        # Generate music
+        output_path = "/tmp/generated_music.mp3"
+        os.makedirs("/tmp", exist_ok=True)
+        
+        with torch.no_grad():
+            pipe(
+                {
+                    "lyrics": lyrics,
+                    "tags": tags,
+                },
+                max_audio_length_ms=max_audio_length_ms,
+                save_path=output_path,
+                topk=top_k,
+                temperature=temperature,
+                cfg_scale=cfg_scale,
+            )
+        
+        logger.info(f"Music generated successfully: {output_path}")
+        return output_path
+        
+    except Exception as e:
+        logger.error(f"Error during generation: {e}")
+        raise gr.Error(f"Generation failed: {e}")
+
+
+def main():
+    """Create Gradio interface"""
+    
+    # Check if models exist
+    if not Path(MODEL_PATH).exists():
+        logger.warning(f"Models directory not found at {MODEL_PATH}")
+        logger.info("You need to download the models first:")
+        logger.info("hf download --local-dir './models/HeartMuLa-oss-3B' HeartMuLa/HeartMuLa-RL-oss-3B-20260123")
+        logger.info("hf download --local-dir './models/HeartCodec-oss' HeartMuLa/HeartCodec-oss-20260123")
+    
+    with gr.Blocks(title="HeartMuLa Music Generation") as demo:
+        gr.Markdown("""
+        # 🎵 HeartMuLa Music Generation
+        
+        Generate music from lyrics and style tags using HeartMuLa, a family of open-source music foundation models.
+        """)
+        
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Input")
+                
+                lyrics_input = gr.Textbox(
+                    label="Lyrics",
+                    placeholder="Enter your lyrics here...",
+                    lines=4,
+                )
+                
+                tags_input = gr.Textbox(
+                    label="Style Tags",
+                    placeholder="e.g., pop, upbeat, energetic",
+                    lines=2,
+                )
+                
+                with gr.Row():
+                    duration_slider = gr.Slider(
+                        minimum=5,
+                        maximum=60,
+                        value=30,
+                        step=5,
+                        label="Duration (seconds)",
+                    )
+                
+                with gr.Row():
+                    temp_slider = gr.Slider(
+                        minimum=0.1,
+                        maximum=2.0,
+                        value=1.0,
+                        step=0.1,
+                        label="Temperature",
+                    )
+                    
+                    cfg_slider = gr.Slider(
+                        minimum=1.0,
+                        maximum=3.0,
+                        value=1.5,
+                        step=0.1,
+                        label="CFG Scale (style strength)",
+                    )
+                
+                with gr.Row():
+                    topk_slider = gr.Slider(
+                        minimum=10,
+                        maximum=100,
+                        value=50,
+                        step=5,
+                        label="Top-K",
+                    )
+                
+                generate_btn = gr.Button("🎵 Generate Music", variant="primary", size="lg")
+            
+            with gr.Column():
+                gr.Markdown("### Output")
+                audio_output = gr.Audio(label="Generated Music", type="filepath")
+                
+                gr.Markdown("""
+                ### 📝 Tips
+                - **Lyrics**: Describe the vocals and melody
+                - **Tags**: Use style descriptors like "pop", "rock", "ambient", "upbeat", etc.
+                - **CFG Scale**: Higher values = stronger style control (1.5 is recommended)
+                - **Temperature**: Higher = more creative, lower = more consistent
+                """)
+        
+        # Connect button to generation function
+        generate_btn.click(
+            fn=generate_music,
+            inputs=[
+                lyrics_input,
+                tags_input,
+                duration_slider,
+                temp_slider,
+                topk_slider,
+                cfg_slider,
+            ],
+            outputs=audio_output,
+        )
+        
+        # Example inputs
+        gr.Examples(
+            examples=[
+                [
+                    "Love is in the air, feel the magic",
+                    "pop, upbeat, romantic",
+                ],
+                [
+                    "Dark skies falling down, lonely tonight",
+                    "rock, emotional, melancholic",
+                ],
+            ],
+            inputs=[lyrics_input, tags_input],
+            outputs=audio_output,
+            fn=generate_music,
+            cache_examples=False,
+        )
+    
+    return demo
+
+
+if __name__ == "__main__":
+    demo = main()
+    demo.launch(server_name="0.0.0.0", server_port=7860, share=False)

examples/README.md

@@ -0,0 +1,15 @@
+# 🎤 Lyrics Transcription
+
+Download checkpoint using any of the following command:
+```
+hf download --local_dir './ckpt/HeartTranscriptor-oss' 'HeartMuLa/HeartTranscriptor-oss' 
+modelscope download --model 'HeartMuLa/HeartTranscriptor-oss' --local_dir './ckpt/HeartTranscriptor-oss'
+```
+
+```
+python ./examples/run_lyrics_transcription.py --model_path=./ckpt
+```
+
+By default this command will load the generated music file at `./assets/output.mp3` and print the transcribed lyrics. Use `--music_path` to specify the path to the music file.
+
+Note that our HeartTranscriptor is trained on separated vocal tracks. In this example usage part, we directly demonstrate on unseparated music tracks, which is purely for simplicity of illustration. We recommend using source separation tools like demucs to separate the tracks before transcribing lyrics to achieve better results.

examples/run_lyrics_transcription.py

@@ -0,0 +1,35 @@
+from heartlib import HeartTranscriptorPipeline
+import argparse
+import torch
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--music_path", type=str, default="./assets/output.mp3")
+
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    pipe = HeartTranscriptorPipeline.from_pretrained(
+        args.model_path,
+        device=torch.device("cuda"),
+        dtype=torch.float16,
+    )
+    with torch.no_grad():
+        result = pipe(
+            args.music_path,
+            **{
+                "max_new_tokens": 256,
+                "num_beams": 2,
+                "task": "transcribe",
+                "condition_on_prev_tokens": False,
+                "compression_ratio_threshold": 1.8,
+                "temperature": (0.0, 0.1, 0.2, 0.4),
+                "logprob_threshold": -1.0,
+                "no_speech_threshold": 0.4,
+            },
+        )
+    print(result)

examples/run_music_generation.py

@@ -0,0 +1,81 @@
+from heartlib import HeartMuLaGenPipeline
+import argparse
+import torch
+
+
+def str2bool(value):
+    if isinstance(value, bool):
+        return value
+    if value.lower() in ("yes", "y", "true", "t", "1"):
+        return True
+    elif value.lower() in ("no", "n", "false", "f", "0"):
+        return False
+    else:
+        raise argparse.ArgumentTypeError(f"Boolean value expected. Got: {value}")
+
+
+def str2dtype(value):
+    value = value.lower()
+    if value == "float32" or value == "fp32":
+        return torch.float32
+    elif value == "float16" or value == "fp16":
+        return torch.float16
+    elif value == "bfloat16" or value == "bf16":
+        return torch.bfloat16
+    else:
+        raise argparse.ArgumentTypeError(f"Dtype not recognized: {value}")
+
+
+def str2device(value):
+    value = value.lower()
+    return torch.device(value)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model_path", type=str, required=True)
+    parser.add_argument("--version", type=str, default="3B")
+    parser.add_argument("--lyrics", type=str, default="./assets/lyrics.txt")
+    parser.add_argument("--tags", type=str, default="./assets/tags.txt")
+    parser.add_argument("--save_path", type=str, default="./assets/output.mp3")
+
+    parser.add_argument("--max_audio_length_ms", type=int, default=240_000)
+    parser.add_argument("--topk", type=int, default=50)
+    parser.add_argument("--temperature", type=float, default=1.0)
+    parser.add_argument("--cfg_scale", type=float, default=1.5)
+    parser.add_argument("--mula_device", type=str2device, default="cuda")
+    parser.add_argument("--codec_device", type=str2device, default="cuda")
+    parser.add_argument("--mula_dtype", type=str2dtype, default="bfloat16")
+    parser.add_argument("--codec_dtype", type=str2dtype, default="float32")
+    parser.add_argument("--lazy_load", type=str2bool, default=False)
+    return parser.parse_args()
+
+
+if __name__ == "__main__":
+    args = parse_args()
+    pipe = HeartMuLaGenPipeline.from_pretrained(
+        args.model_path,
+        device={
+            "mula": torch.device(args.mula_device),
+            "codec": torch.device(args.codec_device),
+        },
+        dtype={
+            "mula": args.mula_dtype,
+            "codec": args.codec_dtype,
+        },
+        version=args.version,
+        lazy_load=args.lazy_load,
+    )
+    with torch.no_grad():
+        pipe(
+            {
+                "lyrics": args.lyrics,
+                "tags": args.tags,
+            },
+            max_audio_length_ms=args.max_audio_length_ms,
+            save_path=args.save_path,
+            topk=args.topk,
+            temperature=args.temperature,
+            cfg_scale=args.cfg_scale,
+        )
+    print(f"Generated music saved to {args.save_path}")

pyproject.toml

@@ -0,0 +1,46 @@
+[build-system]
+requires = ["setuptools>=61", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "heartlib"
+version = "0.1.0"
+description = "A Python Library."
+readme = "README.md"
+requires-python = ">=3.9"
+license = {text = "Apache-2.0"}
+authors = [
+    {name = "HeartMuLa Team", email = "heartmula.ai@gmail.com"}
+]
+dependencies = [
+    "numpy==2.0.2",
+    "torch==2.4.1",
+    "torchaudio==2.4.1",
+    "torchtune==0.4.0",
+    "torchao==0.9.0",
+    "torchvision==0.19.1",
+    "tqdm==4.67.1",
+    "traitlets==5.7.1",
+    "traittypes==0.2.3",
+    "transformers==4.57.0",
+    "tokenizers==0.22.1",
+    "ipykernel==6.17.1",
+    "einops==0.8.1",
+    "accelerate==1.12.0",
+    "bitsandbytes==0.49.0",
+    "vector-quantize-pytorch==1.27.15",
+    "modelscope==1.33.0",
+    "soundfile"
+]
+urls = { "homepage" = "https://heartmula.github.io/" }
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "Operating System :: OS Independent"
+]
+
+[tool.setuptools]
+package-dir = {"" = "src"}
+
+[tool.setuptools.packages.find]
+where = ["src"]
+

requirements.txt

@@ -0,0 +1,11 @@
+torch==2.4.1
+torchaudio==2.4.1
+transformers==4.40.0
+safetensors==0.4.1
+bitsandbytes==0.49.0
+torchtune==0.4.0
+tokenizers==0.15.0
+tqdm==4.66.1
+gradio==4.36.1
+pydantic==2.5.0
+numpy==1.24.3

setup.sh

@@ -0,0 +1,28 @@
+#!/bin/bash
+# Setup script for HF Space deployment
+
+echo "Installing HeartMuLa package..."
+pip install -e .
+
+echo "Downloading HeartMuLa models..."
+mkdir -p ./models
+
+# Download HeartMuLa 3B model (RL-tuned version - recommended)
+echo "Downloading HeartMuLa-RL-oss-3B-20260123..."
+huggingface-cli download \
+  --local-dir "./models/HeartMuLa-oss-3B" \
+  HeartMuLa/HeartMuLa-RL-oss-3B-20260123
+
+# Download HeartCodec
+echo "Downloading HeartCodec-oss-20260123..."
+huggingface-cli download \
+  --local-dir "./models/HeartCodec-oss" \
+  HeartMuLa/HeartCodec-oss-20260123
+
+# Copy tokenizer and config
+echo "Copying tokenizer and config..."
+huggingface-cli download \
+  --local-dir "./models" \
+  HeartMuLa/HeartMuLaGen
+
+echo "Setup complete! Starting Gradio app..."

src/heartlib/__init__.py

@@ -0,0 +1,7 @@
+from .pipelines.music_generation import HeartMuLaGenPipeline
+from .pipelines.lyrics_transcription import HeartTranscriptorPipeline
+
+__all__ = [
+    "HeartMuLaGenPipeline",
+    "HeartTranscriptorPipeline"
+]

src/heartlib/heartcodec/configuration_heartcodec.py

@@ -0,0 +1,73 @@
+from transformers.configuration_utils import PretrainedConfig
+from typing import List
+
+
+class HeartCodecConfig(PretrainedConfig):
+    model_type = "heartcodec"
+
+    def __init__(
+        self,
+        # config for rvq
+        dim: int = 512,
+        codebook_size: int = 8192,
+        decay: float = 0.9,
+        commitment_weight: float = 1.0,
+        threshold_ema_dead_code: int = 2,
+        use_cosine_sim: bool = False,
+        codebook_dim: int = 32,
+        num_quantizers: int = 8,
+        # config for diffusion transformer
+        attention_head_dim: int = 64,
+        in_channels: int = 1024,
+        norm_type: str = "ada_norm_single",
+        num_attention_heads: int = 24,
+        num_layers: int = 24,
+        num_layers_2: int = 6,
+        out_channels: int = 256,
+        # config for sq codec
+        num_bands: int = 1,
+        sample_rate: int = 48000,
+        causal: bool = True,
+        num_samples: int = 2,
+        downsample_factors: List[int] = [3, 4, 4, 4, 5],
+        downsample_kernel_sizes: List[int] = [6, 8, 8, 8, 10],
+        upsample_factors: List[int] = [5, 4, 4, 4, 3],
+        upsample_kernel_sizes: List[int] = [10, 8, 8, 8, 6],
+        latent_hidden_dim: int = 128,
+        default_kernel_size: int = 7,
+        delay_kernel_size: int = 5,
+        init_channel: int = 64,
+        res_kernel_size: int = 7,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.dim = dim
+        self.codebook_size = codebook_size
+        self.decay = decay
+        self.commitment_weight = commitment_weight
+        self.threshold_ema_dead_code = threshold_ema_dead_code
+        self.use_cosine_sim = use_cosine_sim
+        self.codebook_dim = codebook_dim
+        self.num_quantizers = num_quantizers
+
+        self.attention_head_dim = attention_head_dim
+        self.in_channels = in_channels
+        self.norm_type = norm_type
+        self.num_attention_heads = num_attention_heads
+        self.num_layers = num_layers
+        self.num_layers_2 = num_layers_2
+        self.out_channels = out_channels
+
+        self.num_bands = num_bands
+        self.sample_rate = sample_rate
+        self.causal = causal
+        self.num_samples = num_samples
+        self.downsample_factors = downsample_factors
+        self.downsample_kernel_sizes = downsample_kernel_sizes
+        self.upsample_factors = upsample_factors
+        self.upsample_kernel_sizes = upsample_kernel_sizes
+        self.latent_hidden_dim = latent_hidden_dim
+        self.default_kernel_size = default_kernel_size
+        self.delay_kernel_size = delay_kernel_size
+        self.init_channel = init_channel
+        self.res_kernel_size = res_kernel_size

src/heartlib/heartcodec/modeling_heartcodec.py

@@ -0,0 +1,183 @@
+import torch
+from .models.flow_matching import FlowMatching
+from .models.sq_codec import ScalarModel
+from .configuration_heartcodec import HeartCodecConfig
+from transformers.modeling_utils import PreTrainedModel
+import math
+import numpy as np
+
+
+class HeartCodec(PreTrainedModel):
+    config_class = HeartCodecConfig
+
+    def __init__(
+        self,
+        config: HeartCodecConfig,
+    ):
+        super(HeartCodec, self).__init__(config)
+
+        self.config = config
+
+        self.flow_matching = FlowMatching(
+            dim=config.dim,
+            codebook_size=config.codebook_size,
+            decay=config.decay,
+            commitment_weight=config.commitment_weight,
+            threshold_ema_dead_code=config.threshold_ema_dead_code,
+            use_cosine_sim=config.use_cosine_sim,
+            codebook_dim=config.codebook_dim,
+            num_quantizers=config.num_quantizers,
+            attention_head_dim=config.attention_head_dim,
+            in_channels=config.in_channels,
+            norm_type=config.norm_type,
+            num_attention_heads=config.num_attention_heads,
+            num_layers=config.num_layers,
+            num_layers_2=config.num_layers_2,
+            out_channels=config.out_channels,
+        )
+        self.scalar_model = ScalarModel(
+            num_bands=config.num_bands,
+            sample_rate=config.sample_rate,
+            causal=config.causal,
+            num_samples=config.num_samples,
+            downsample_factors=config.downsample_factors,
+            downsample_kernel_sizes=config.downsample_kernel_sizes,
+            upsample_factors=config.upsample_factors,
+            upsample_kernel_sizes=config.upsample_kernel_sizes,
+            latent_hidden_dim=config.latent_hidden_dim,
+            default_kernel_size=config.default_kernel_size,
+            delay_kernel_size=config.delay_kernel_size,
+            init_channel=config.init_channel,
+            res_kernel_size=config.res_kernel_size,
+        )
+        self.post_init()
+
+        self.sample_rate = config.sample_rate
+
+    @torch.inference_mode()
+    def detokenize(
+        self,
+        codes,
+        duration=29.76,
+        num_steps=10,
+        disable_progress=False,
+        guidance_scale=1.25,
+    ):
+        codes = codes.unsqueeze(0).to(self.device)
+        first_latent = torch.randn(
+            codes.shape[0], int(duration * 25), 256, dtype=self.dtype
+        ).to(
+            self.device
+        )  # B, T, 64
+        first_latent_length = 0
+        first_latent_codes_length = 0
+        min_samples = int(duration * 12.5)
+        hop_samples = min_samples // 93 * 80
+        ovlp_samples = min_samples - hop_samples
+        ovlp_frames = ovlp_samples * 2
+        codes_len = codes.shape[-1]  #
+        target_len = int(
+            (codes_len - first_latent_codes_length) / 12.5 * self.sample_rate
+        )
+
+        # code repeat
+        if codes_len < min_samples:
+            while codes.shape[-1] < min_samples:
+                codes = torch.cat([codes, codes], -1)
+            codes = codes[:, :, 0:min_samples]
+        codes_len = codes.shape[-1]
+        if (codes_len - ovlp_frames) % hop_samples > 0:
+            len_codes = (
+                math.ceil((codes_len - ovlp_samples) / float(hop_samples)) * hop_samples
+                + ovlp_samples
+            )
+            while codes.shape[-1] < len_codes:
+                codes = torch.cat([codes, codes], -1)
+            codes = codes[:, :, 0:len_codes]
+        latent_length = int(duration * 25)
+        latent_list = []
+
+        for sinx in range(0, codes.shape[-1] - hop_samples + 1, hop_samples):
+            codes_input = []
+            codes_input.append(codes[:, :, sinx : sinx + min_samples])
+            if sinx == 0 or ovlp_frames == 0:
+                incontext_length = first_latent_length
+                latents = self.flow_matching.inference_codes(
+                    codes_input,
+                    first_latent,
+                    latent_length,
+                    incontext_length,
+                    guidance_scale=guidance_scale,
+                    num_steps=num_steps,
+                    disable_progress=disable_progress,
+                    scenario="other_seg",
+                )
+                latent_list.append(latents)
+            else:
+                true_latent = latent_list[-1][:, -ovlp_frames:, :]
+                len_add_to_latent = latent_length - true_latent.shape[1]  #
+                incontext_length = true_latent.shape[1]
+                true_latent = torch.cat(
+                    [
+                        true_latent,
+                        torch.randn(
+                            true_latent.shape[0],
+                            len_add_to_latent,
+                            true_latent.shape[-1],
+                            dtype=self.dtype,
+                        ).to(self.device),
+                    ],
+                    1,
+                )
+                latents = self.flow_matching.inference_codes(
+                    codes_input,
+                    true_latent,
+                    latent_length,
+                    incontext_length,
+                    guidance_scale=guidance_scale,
+                    num_steps=num_steps,
+                    disable_progress=disable_progress,
+                    scenario="other_seg",
+                )
+                latent_list.append(latents)
+
+        # latent_list = [l.float() for l in latent_list]
+        latent_list[0] = latent_list[0][:, first_latent_length:, :]
+        min_samples = int(duration * self.sample_rate)
+        hop_samples = min_samples // 93 * 80
+        ovlp_samples = min_samples - hop_samples
+
+        output = None
+        for i in range(len(latent_list)):
+            latent = latent_list[i]
+            bsz, t, f = latent.shape
+
+            latent = latent.reshape(
+                latent.shape[0], latent.shape[1], 2, latent.shape[2] // 2
+            ).permute(0, 2, 1, 3)
+            latent = latent.reshape(
+                latent.shape[0] * 2, latent.shape[2], latent.shape[3]
+            )
+            cur_output = (
+                self.scalar_model.decode(latent.transpose(1, 2)).squeeze(0).squeeze(1)
+            )  # 1 512 256
+
+            cur_output = cur_output[:, 0:min_samples].detach().cpu()  # B, T
+            if cur_output.dim() == 3:
+                cur_output = cur_output[0]
+
+            if output is None:
+                output = cur_output
+            else:
+                if ovlp_samples == 0:
+                    output = torch.cat([output, cur_output], -1)
+                else:
+                    ov_win = torch.from_numpy(np.linspace(0, 1, ovlp_samples)[None, :])
+                    ov_win = torch.cat([ov_win, 1 - ov_win], -1)
+                    output[:, -ovlp_samples:] = (
+                        output[:, -ovlp_samples:] * ov_win[:, -ovlp_samples:]
+                        + cur_output[:, 0:ovlp_samples] * ov_win[:, 0:ovlp_samples]
+                    )
+                    output = torch.cat([output, cur_output[:, ovlp_samples:]], -1)
+        output = output[:, 0:target_len]
+        return output

src/heartlib/heartcodec/models/flow_matching.py

@@ -0,0 +1,176 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from tqdm import tqdm
+from vector_quantize_pytorch import ResidualVQ
+from .transformer import LlamaTransformer
+
+
+class FlowMatching(nn.Module):
+    def __init__(
+        self,
+        # rvq stuff
+        dim: int = 512,
+        codebook_size: int = 8192,
+        decay: float = 0.9,
+        commitment_weight: float = 1.0,
+        threshold_ema_dead_code: int = 2,
+        use_cosine_sim: bool = False,
+        codebook_dim: int = 32,
+        num_quantizers: int = 8,
+        # dit backbone stuff
+        attention_head_dim: int = 64,
+        in_channels: int = 1024,
+        norm_type: str = "ada_norm_single",
+        num_attention_heads: int = 24,
+        num_layers: int = 24,
+        num_layers_2: int = 6,
+        out_channels: int = 256,
+    ):
+        super().__init__()
+
+        self.vq_embed = ResidualVQ(
+            dim=dim,
+            codebook_size=codebook_size,
+            decay=decay,
+            commitment_weight=commitment_weight,
+            threshold_ema_dead_code=threshold_ema_dead_code,
+            use_cosine_sim=use_cosine_sim,
+            codebook_dim=codebook_dim,
+            num_quantizers=num_quantizers,
+        )
+        self.cond_feature_emb = nn.Linear(dim, dim)
+        self.zero_cond_embedding1 = nn.Parameter(torch.randn(dim))
+        self.estimator = LlamaTransformer(
+            attention_head_dim=attention_head_dim,
+            in_channels=in_channels,
+            norm_type=norm_type,
+            num_attention_heads=num_attention_heads,
+            num_layers=num_layers,
+            num_layers_2=num_layers_2,
+            out_channels=out_channels,
+        )
+
+        self.latent_dim = out_channels
+
+    @torch.no_grad()
+    def inference_codes(
+        self,
+        codes,
+        true_latents,
+        latent_length,
+        incontext_length,
+        guidance_scale=2.0,
+        num_steps=20,
+        disable_progress=True,
+        scenario="start_seg",
+    ):
+        device = true_latents.device
+        dtype = true_latents.dtype
+        # codes_bestrq_middle, codes_bestrq_last = codes
+        codes_bestrq_emb = codes[0]
+
+        batch_size = codes_bestrq_emb.shape[0]
+        self.vq_embed.eval()
+        quantized_feature_emb = self.vq_embed.get_output_from_indices(
+            codes_bestrq_emb.transpose(1, 2)
+        )
+        quantized_feature_emb = self.cond_feature_emb(quantized_feature_emb)  # b t 512
+        # assert 1==2
+        quantized_feature_emb = F.interpolate(
+            quantized_feature_emb.permute(0, 2, 1), scale_factor=2, mode="nearest"
+        ).permute(0, 2, 1)
+
+        num_frames = quantized_feature_emb.shape[1]  #
+        latents = torch.randn(
+            (batch_size, num_frames, self.latent_dim), device=device, dtype=dtype
+        )
+        latent_masks = torch.zeros(
+            latents.shape[0], latents.shape[1], dtype=torch.int64, device=latents.device
+        )
+        latent_masks[:, 0:latent_length] = 2
+        if scenario == "other_seg":
+            latent_masks[:, 0:incontext_length] = 1
+
+        quantized_feature_emb = (latent_masks > 0.5).unsqueeze(
+            -1
+        ) * quantized_feature_emb + (latent_masks < 0.5).unsqueeze(
+            -1
+        ) * self.zero_cond_embedding1.unsqueeze(
+            0
+        )
+
+        incontext_latents = (
+            true_latents
+            * ((latent_masks > 0.5) * (latent_masks < 1.5)).unsqueeze(-1).float()
+        )
+        incontext_length = ((latent_masks > 0.5) * (latent_masks < 1.5)).sum(-1)[0]
+
+        additional_model_input = torch.cat([quantized_feature_emb], 1)
+        temperature = 1.0
+        t_span = torch.linspace(
+            0, 1, num_steps + 1, device=quantized_feature_emb.device
+        )
+        latents = self.solve_euler(
+            latents * temperature,
+            incontext_latents.to(dtype),
+            incontext_length,
+            t_span,
+            additional_model_input,
+            guidance_scale,
+        )
+
+        latents[:, 0:incontext_length, :] = incontext_latents[
+            :, 0:incontext_length, :
+        ]  # B, T, dim
+        return latents
+
+    def solve_euler(self, x, incontext_x, incontext_length, t_span, mu, guidance_scale):
+        """
+        Fixed euler solver for ODEs.
+        Args:
+            x (torch.Tensor): random noise
+            t_span (torch.Tensor): n_timesteps interpolated
+                shape: (n_timesteps + 1,)
+            mu (torch.Tensor): output of encoder
+                shape: (batch_size, n_feats, mel_timesteps)
+        """
+        t, _, dt = t_span[0], t_span[-1], t_span[1] - t_span[0]
+        noise = x.clone()
+        # I am storing this because I can later plot it by putting a debugger here and saving it to a file
+        # Or in future might add like a return_all_steps flag
+        sol = []
+        for step in tqdm(range(1, len(t_span))):
+            x[:, 0:incontext_length, :] = (1 - (1 - 1e-6) * t) * noise[
+                :, 0:incontext_length, :
+            ] + t * incontext_x[:, 0:incontext_length, :]
+            if guidance_scale > 1.0:
+                dphi_dt = self.estimator(
+                    torch.cat(
+                        [
+                            torch.cat([x, x], 0),
+                            torch.cat([incontext_x, incontext_x], 0),
+                            torch.cat([torch.zeros_like(mu), mu], 0),
+                        ],
+                        2,
+                    ),
+                    timestep=t.unsqueeze(-1).repeat(2),
+                )
+                dphi_dt_uncond, dhpi_dt_cond = dphi_dt.chunk(2, 0)
+                dphi_dt = dphi_dt_uncond + guidance_scale * (
+                    dhpi_dt_cond - dphi_dt_uncond
+                )
+            else:
+                dphi_dt = self.estimator(
+                    torch.cat([x, incontext_x, mu], 2), timestep=t.unsqueeze(-1)
+                )
+
+            x = x + dt * dphi_dt
+            t = t + dt
+            sol.append(x)
+            if step < len(t_span) - 1:
+                dt = t_span[step + 1] - t
+
+        result = sol[-1]
+
+        return result

src/heartlib/heartcodec/models/sq_codec.py

@@ -0,0 +1,539 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from torch.nn.utils.parametrizations import weight_norm
+from torch.nn.utils import remove_weight_norm
+from torch.autograd.function import InplaceFunction
+
+
+def get_padding(kernel_size, dilation=1):
+    return int((kernel_size * dilation - dilation) / 2)
+
+
+# Scripting this brings model speed up 1.4x
+@torch.jit.script
+def snake(x, alpha):
+    shape = x.shape
+    x = x.reshape(shape[0], shape[1], -1)
+    x = x + (alpha + 1e-9).reciprocal() * torch.sin(alpha * x).pow(2)
+    x = x.reshape(shape)
+    return x
+
+
+class Snake1d(nn.Module):
+    def __init__(self, channels):
+        super().__init__()
+        self.alpha = nn.Parameter(torch.ones(1, channels, 1))
+
+    def forward(self, x):
+        return snake(x, self.alpha)
+
+
+class Conv1d(nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        padding_mode: str = "zeros",
+        bias: bool = True,
+        padding=None,
+        causal: bool = False,
+        w_init_gain=None,
+    ):
+        self.causal = causal
+        if padding is None:
+            if causal:
+                padding = 0
+                self.left_padding = dilation * (kernel_size - 1)
+            else:
+                padding = get_padding(kernel_size, dilation)
+        super(Conv1d, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            dilation=dilation,
+            groups=groups,
+            padding_mode=padding_mode,
+            bias=bias,
+        )
+        if w_init_gain is not None:
+            torch.nn.init.xavier_uniform_(
+                self.weight, gain=torch.nn.init.calculate_gain(w_init_gain)
+            )
+
+    def forward(self, x):
+        if self.causal:
+            x = F.pad(x.unsqueeze(2), (self.left_padding, 0, 0, 0)).squeeze(2)
+
+        return super(Conv1d, self).forward(x)
+
+
+class ConvTranspose1d(nn.ConvTranspose1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        output_padding: int = 0,
+        groups: int = 1,
+        bias: bool = True,
+        dilation: int = 1,
+        padding=None,
+        padding_mode: str = "zeros",
+        causal: bool = False,
+    ):
+        if padding is None:
+            padding = 0 if causal else (kernel_size - stride) // 2
+        if causal:
+            assert padding == 0, "padding is not allowed in causal ConvTranspose1d."
+            assert (
+                kernel_size == 2 * stride
+            ), "kernel_size must be equal to 2*stride is not allowed in causal ConvTranspose1d."
+        super(ConvTranspose1d, self).__init__(
+            in_channels,
+            out_channels,
+            kernel_size,
+            stride=stride,
+            padding=padding,
+            output_padding=output_padding,
+            groups=groups,
+            bias=bias,
+            dilation=dilation,
+            padding_mode=padding_mode,
+        )
+        self.causal = causal
+        self.stride = stride
+
+    def forward(self, x):
+        x = super(ConvTranspose1d, self).forward(x)
+        if self.causal:
+            x = x[:, :, : -self.stride]
+        return x
+
+
+class PreProcessor(nn.Module):
+    def __init__(self, n_in, n_out, num_samples, kernel_size=7, causal=False):
+        super(PreProcessor, self).__init__()
+        self.pooling = torch.nn.AvgPool1d(kernel_size=num_samples)
+        self.conv = Conv1d(n_in, n_out, kernel_size=kernel_size, causal=causal)
+        self.activation = nn.PReLU()
+
+    def forward(self, x):
+        output = self.activation(self.conv(x))
+        output = self.pooling(output)
+        return output
+
+
+class PostProcessor(nn.Module):
+    def __init__(self, n_in, n_out, num_samples, kernel_size=7, causal=False):
+        super(PostProcessor, self).__init__()
+        self.num_samples = num_samples
+        self.conv = Conv1d(n_in, n_out, kernel_size=kernel_size, causal=causal)
+        self.activation = nn.PReLU()
+
+    def forward(self, x):
+        x = torch.transpose(x, 1, 2)
+        B, T, C = x.size()
+        x = x.repeat(1, 1, self.num_samples).view(B, -1, C)
+        x = torch.transpose(x, 1, 2)
+        output = self.activation(self.conv(x))
+        return output
+
+
+class ResidualUnit(nn.Module):
+    def __init__(self, n_in, n_out, dilation, res_kernel_size=7, causal=False):
+        super(ResidualUnit, self).__init__()
+        self.conv1 = weight_norm(
+            Conv1d(
+                n_in,
+                n_out,
+                kernel_size=res_kernel_size,
+                dilation=dilation,
+                causal=causal,
+            )
+        )
+        self.conv2 = weight_norm(Conv1d(n_in, n_out, kernel_size=1, causal=causal))
+        self.activation1 = nn.PReLU()
+        self.activation2 = nn.PReLU()
+
+    def forward(self, x):
+        output = self.activation1(self.conv1(x))
+        output = self.activation2(self.conv2(output))
+        return output + x
+
+
+class ResEncoderBlock(nn.Module):
+    def __init__(
+        self, n_in, n_out, stride, down_kernel_size, res_kernel_size=7, causal=False
+    ):
+        super(ResEncoderBlock, self).__init__()
+        self.convs = nn.ModuleList(
+            [
+                ResidualUnit(
+                    n_in,
+                    n_out // 2,
+                    dilation=1,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out // 2,
+                    n_out // 2,
+                    dilation=3,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out // 2,
+                    n_out // 2,
+                    dilation=5,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out // 2,
+                    n_out // 2,
+                    dilation=7,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out // 2,
+                    n_out // 2,
+                    dilation=9,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+            ]
+        )
+
+        self.down_conv = DownsampleLayer(
+            n_in, n_out, down_kernel_size, stride=stride, causal=causal
+        )
+
+    def forward(self, x):
+        for conv in self.convs:
+            x = conv(x)
+        x = self.down_conv(x)
+        return x
+
+
+class ResDecoderBlock(nn.Module):
+    def __init__(
+        self, n_in, n_out, stride, up_kernel_size, res_kernel_size=7, causal=False
+    ):
+        super(ResDecoderBlock, self).__init__()
+        self.up_conv = UpsampleLayer(
+            n_in,
+            n_out,
+            kernel_size=up_kernel_size,
+            stride=stride,
+            causal=causal,
+            activation=None,
+        )
+
+        self.convs = nn.ModuleList(
+            [
+                ResidualUnit(
+                    n_out,
+                    n_out,
+                    dilation=1,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out,
+                    n_out,
+                    dilation=3,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out,
+                    n_out,
+                    dilation=5,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out,
+                    n_out,
+                    dilation=7,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+                ResidualUnit(
+                    n_out,
+                    n_out,
+                    dilation=9,
+                    res_kernel_size=res_kernel_size,
+                    causal=causal,
+                ),
+            ]
+        )
+
+    def forward(self, x):
+        x = self.up_conv(x)
+        for conv in self.convs:
+            x = conv(x)
+        return x
+
+
+class DownsampleLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        causal: bool = False,
+        activation=nn.PReLU(),
+        use_weight_norm: bool = True,
+        pooling: bool = False,
+    ):
+        super(DownsampleLayer, self).__init__()
+        self.pooling = pooling
+        self.stride = stride
+        self.activation = nn.PReLU()
+        self.use_weight_norm = use_weight_norm
+        if pooling:
+            self.layer = Conv1d(in_channels, out_channels, kernel_size, causal=causal)
+            self.pooling = nn.AvgPool1d(kernel_size=stride)
+        else:
+            self.layer = Conv1d(
+                in_channels, out_channels, kernel_size, stride=stride, causal=causal
+            )
+        if use_weight_norm:
+            self.layer = weight_norm(self.layer)
+
+    def forward(self, x):
+        x = self.layer(x)
+        x = self.activation(x) if self.activation is not None else x
+        if self.pooling:
+            x = self.pooling(x)
+        return x
+
+    def remove_weight_norm(self):
+        if self.use_weight_norm:
+            remove_weight_norm(self.layer)
+
+
+class UpsampleLayer(nn.Module):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        causal: bool = False,
+        activation=nn.PReLU(),
+        use_weight_norm: bool = True,
+        repeat: bool = False,
+    ):
+        super(UpsampleLayer, self).__init__()
+        self.repeat = repeat
+        self.stride = stride
+        self.activation = activation
+        self.use_weight_norm = use_weight_norm
+        if repeat:
+            self.layer = Conv1d(in_channels, out_channels, kernel_size, causal=causal)
+        else:
+            self.layer = ConvTranspose1d(
+                in_channels, out_channels, kernel_size, stride=stride, causal=causal
+            )
+        if use_weight_norm:
+            self.layer = weight_norm(self.layer)
+
+    def forward(self, x):
+        x = self.layer(x)
+        x = self.activation(x) if self.activation is not None else x
+        if self.repeat:
+            x = torch.transpose(x, 1, 2)
+            B, T, C = x.size()
+            x = x.repeat(1, 1, self.stride).view(B, -1, C)
+            x = torch.transpose(x, 1, 2)
+        return x
+
+    def remove_weight_norm(self):
+        if self.use_weight_norm:
+            remove_weight_norm(self.layer)
+
+
+class round_func9(InplaceFunction):
+    @staticmethod
+    def forward(ctx, input):
+        ctx.input = input
+        return torch.round(9 * input) / 9
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_input = grad_output.clone()
+        return grad_input
+
+
+class ScalarModel(nn.Module):
+    def __init__(
+        self,
+        num_bands,
+        sample_rate,
+        causal,
+        num_samples,
+        downsample_factors,
+        downsample_kernel_sizes,
+        upsample_factors,
+        upsample_kernel_sizes,
+        latent_hidden_dim,
+        default_kernel_size,
+        delay_kernel_size,
+        init_channel,
+        res_kernel_size,
+        mode="pre_proj",
+    ):
+        super(ScalarModel, self).__init__()
+        # self.args = args
+        self.encoder = []
+        self.decoder = []
+        self.vq = round_func9()  # using 9
+        self.mode = mode
+        # Encoder parts
+        self.encoder.append(
+            weight_norm(
+                Conv1d(
+                    num_bands,
+                    init_channel,
+                    kernel_size=default_kernel_size,
+                    causal=causal,
+                )
+            )
+        )
+        if num_samples > 1:
+            # Downsampling
+            self.encoder.append(
+                PreProcessor(
+                    init_channel,
+                    init_channel,
+                    num_samples,
+                    kernel_size=default_kernel_size,
+                    causal=causal,
+                )
+            )
+        for i, down_factor in enumerate(downsample_factors):
+            self.encoder.append(
+                ResEncoderBlock(
+                    init_channel * np.power(2, i),
+                    init_channel * np.power(2, i + 1),
+                    down_factor,
+                    downsample_kernel_sizes[i],
+                    res_kernel_size,
+                    causal=causal,
+                )
+            )
+        self.encoder.append(
+            weight_norm(
+                Conv1d(
+                    init_channel * np.power(2, len(downsample_factors)),
+                    latent_hidden_dim,
+                    kernel_size=default_kernel_size,
+                    causal=causal,
+                )
+            )
+        )
+        # Decoder
+        # look ahead
+        self.decoder.append(
+            weight_norm(
+                Conv1d(
+                    latent_hidden_dim,
+                    init_channel * np.power(2, len(upsample_factors)),
+                    kernel_size=delay_kernel_size,
+                )
+            )
+        )
+        for i, upsample_factor in enumerate(upsample_factors):
+            self.decoder.append(
+                ResDecoderBlock(
+                    init_channel * np.power(2, len(upsample_factors) - i),
+                    init_channel * np.power(2, len(upsample_factors) - i - 1),
+                    upsample_factor,
+                    upsample_kernel_sizes[i],
+                    res_kernel_size,
+                    causal=causal,
+                )
+            )
+        if num_samples > 1:
+            self.decoder.append(
+                PostProcessor(
+                    init_channel,
+                    init_channel,
+                    num_samples,
+                    kernel_size=default_kernel_size,
+                    causal=causal,
+                )
+            )
+        self.decoder.append(
+            weight_norm(
+                Conv1d(
+                    init_channel,
+                    num_bands,
+                    kernel_size=default_kernel_size,
+                    causal=causal,
+                )
+            )
+        )
+        self.encoder = nn.ModuleList(self.encoder)
+        self.decoder = nn.ModuleList(self.decoder)
+
+    def forward(self, x):
+        for i, layer in enumerate(self.encoder):
+            if i != len(self.encoder) - 1:
+                x = layer(x)
+            else:
+                x = F.tanh(layer(x))
+        # import pdb; pdb.set_trace()
+        x = self.vq.apply(x)  # vq
+        for i, layer in enumerate(self.decoder):
+            x = layer(x)
+        return x
+
+    def inference(self, x):
+        for i, layer in enumerate(self.encoder):
+            if i != len(self.encoder) - 1:
+                x = layer(x)
+            else:
+                x = F.tanh(layer(x))  # reverse to tanh
+
+        emb = x
+        # import pdb; pdb.set_trace()
+        emb_quant = self.vq.apply(emb)  # vq
+        x = emb_quant
+        for i, layer in enumerate(self.decoder):
+            x = layer(x)
+        return emb, emb_quant, x
+
+    def encode(self, x):
+        for i, layer in enumerate(self.encoder):
+            if i != len(self.encoder) - 1:
+                x = layer(x)
+            else:
+                x = F.tanh(layer(x))  # reverse to tanh
+
+        emb = x
+        # import pdb; pdb.set_trace()
+        emb_quant = self.vq.apply(emb)  # vq
+        return emb
+
+    def decode(self, x):
+        x = self.vq.apply(
+            x
+        )  # make sure the prediction follow the similar disctribution
+        for i, layer in enumerate(self.decoder):
+            x = layer(x)
+        return x

src/heartlib/heartcodec/models/transformer.py

@@ -0,0 +1,501 @@
+import math
+from typing import Optional, Tuple
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        var = x.pow(2).mean(dim=-1, keepdim=True)
+        x = x * torch.rsqrt(var + self.eps)
+        return self.weight * x
+
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, base: int = 10000):
+        super().__init__()
+        self.dim = dim
+        self.base = base
+        self._cache = {}
+
+    def get_sin_cos(self, seq_len: int, device, dtype):
+        key = (seq_len, device, dtype)
+        cached = self._cache.get(key, None)
+        if cached is not None and cached[0].device == device:
+            return cached
+        inv_freq = 1.0 / (
+            self.base
+            ** (torch.arange(0, self.dim, 2, device=device, dtype=dtype) / self.dim)
+        )
+        t = torch.arange(seq_len, device=device, dtype=dtype)
+        freqs = torch.einsum("i,j->ij", t, inv_freq)
+        sin = freqs.sin()
+        cos = freqs.cos()
+        self._cache[key] = (sin, cos)
+        return sin, cos
+
+    def apply_rotary(
+        self, x: torch.Tensor, sin: torch.Tensor, cos: torch.Tensor
+    ) -> torch.Tensor:
+        x1, x2 = x[..., : self.dim // 2], x[..., self.dim // 2 : self.dim]
+        # Interleave sin/cos across pairs
+        x_rot = torch.stack((-x2, x1), dim=-1).reshape_as(x[..., : self.dim])
+        return (x[..., : self.dim] * cos.unsqueeze(-1)).reshape_as(
+            x[..., : self.dim]
+        ) + (x_rot * sin.unsqueeze(-1)).reshape_as(x[..., : self.dim])
+
+
+class LlamaAttention(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        head_dim: int,
+        bias: bool = False,
+        dropout: float = 0.0,
+        rope_dim: Optional[int] = None,
+        cross_attention_dim: Optional[int] = None,
+        use_sdpa: bool = True,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.inner_dim = n_heads * head_dim
+        self.cross_attention_dim = cross_attention_dim
+        self.q_proj = nn.Linear(dim, self.inner_dim, bias=bias)
+        k_in = dim if cross_attention_dim is None else cross_attention_dim
+        self.k_proj = nn.Linear(k_in, self.inner_dim, bias=bias)
+        self.v_proj = nn.Linear(k_in, self.inner_dim, bias=bias)
+        self.o_proj = nn.Linear(self.inner_dim, dim, bias=bias)
+        self.dropout = dropout
+        self.rope_dim = rope_dim if rope_dim is not None else head_dim
+        self.rope = RotaryEmbedding(self.rope_dim)
+        self.use_sdpa = use_sdpa
+        self._has_sdpa = hasattr(F, "scaled_dot_product_attention")
+
+    def _shape(self, x: torch.Tensor, b: int, t: int) -> torch.Tensor:
+        return x.view(b, t, self.n_heads, self.head_dim).transpose(1, 2)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        b, t, c = x.shape
+        q = self._shape(self.q_proj(x), b, t)
+        if encoder_hidden_states is None:
+            k = self._shape(self.k_proj(x), b, t)
+            v = self._shape(self.v_proj(x), b, t)
+        else:
+            bt, tk, ck = encoder_hidden_states.shape
+            k = self._shape(self.k_proj(encoder_hidden_states), b, tk)
+            v = self._shape(self.v_proj(encoder_hidden_states), b, tk)
+
+        # RoPE on first rope_dim of head_dim
+        rope_dim = min(self.rope_dim, self.head_dim)
+        seq_len_for_rope = k.shape[-2]
+        sin, cos = self.rope.get_sin_cos(
+            seq_len_for_rope, device=x.device, dtype=x.dtype
+        )
+
+        def apply_rope_vec(tensor):
+            head = tensor[..., :rope_dim]
+            tail = tensor[..., rope_dim:]
+            b, h, tt, _ = head.shape
+            head = head.view(b, h, tt, rope_dim // 2, 2)
+            sin_ = sin.view(1, 1, tt, rope_dim // 2, 1)
+            cos_ = cos.view(1, 1, tt, rope_dim // 2, 1)
+            x1 = head[..., 0:1]
+            x2 = head[..., 1:2]
+            rot = torch.cat(
+                [x1 * cos_ - x2 * sin_, x1 * sin_ + x2 * cos_], dim=-1
+            ).view(b, h, tt, rope_dim)
+            return torch.cat([rot, tail], dim=-1)
+
+        q = apply_rope_vec(q)
+        k = apply_rope_vec(k)
+
+        # Prefer PyTorch SDPA (can enable FlashAttention kernel on supported GPUs)
+        if self.use_sdpa and self._has_sdpa:
+            s = k.shape[-2]
+            attn_mask_sdpa = None
+            if attention_mask is not None:
+                m = attention_mask
+
+                if m.dim() == 2 and m.shape == (b, s):  # [b, s]
+                    m = m[:, None, None, :]  # [b,1,1,s]
+                elif m.dim() == 3 and m.shape[-2] == 1:  # [b,1,s]
+                    m = m[:, None, :, :]  # [b,1,1,s]
+                elif m.dim() == 3 and m.shape[-2] == t:  # [b,t,s]
+                    m = m[:, None, :, :]  # [b,1,t,s]
+                elif m.dim() == 4 and m.shape[1] == 1:  # [b,1,t,s] or [b,1,1,s]
+                    pass
+                attn_mask_sdpa = m
+
+            out = F.scaled_dot_product_attention(
+                q,
+                k,
+                v,
+                attn_mask=attn_mask_sdpa,
+                dropout_p=self.dropout if self.training else 0.0,
+                is_causal=False,
+            )
+            out = out.transpose(1, 2).contiguous().view(b, t, self.inner_dim)
+            return self.o_proj(out)
+        else:
+            attn_scores = torch.matmul(q, k.transpose(-2, -1)) / math.sqrt(
+                self.head_dim
+            )
+            if attention_mask is not None:
+                attn_scores = attn_scores + attention_mask
+            attn = attn_scores.softmax(dim=-1)
+            attn = F.dropout(attn, p=self.dropout, training=self.training)
+            out = torch.matmul(attn, v)
+            out = out.transpose(1, 2).contiguous().view(b, t, self.inner_dim)
+            return self.o_proj(out)
+
+
+class LlamaMLP(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        hidden_dim: Optional[int] = None,
+        multiple_of: int = 256,
+        dropout: float = 0.0,
+    ):
+        super().__init__()
+        hidden_dim = hidden_dim or 4 * dim
+        # align to multiple_of like Llama
+        hidden_dim = int(2 * hidden_dim / 3)
+        hidden_dim = multiple_of * ((hidden_dim + multiple_of - 1) // multiple_of)
+        self.gate = nn.Linear(dim, hidden_dim, bias=False)
+        self.up = nn.Linear(dim, hidden_dim, bias=False)
+        self.down = nn.Linear(hidden_dim, dim, bias=False)
+        self.dropout = dropout
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = F.silu(self.gate(x)) * self.up(x)
+        x = F.dropout(x, p=self.dropout, training=self.training)
+        return self.down(x)
+
+
+class LlamaTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        n_heads: int,
+        head_dim: int,
+        mlp_multiple_of: int = 256,
+        dropout: float = 0.0,
+        attention_bias: bool = False,
+        cross_attention_dim: Optional[int] = None,
+        use_ada_layer_norm_single: bool = False,
+    ):
+        super().__init__()
+        self.attn_norm = RMSNorm(dim, 1e-6)
+        self.attn = LlamaAttention(
+            dim,
+            n_heads,
+            head_dim,
+            bias=attention_bias,
+            dropout=dropout,
+            rope_dim=head_dim,
+            cross_attention_dim=None,
+        )
+        self.cross_attn = None
+        if cross_attention_dim is not None:
+            self.cross_attn_norm = RMSNorm(dim, 1e-6)
+            self.cross_attn = LlamaAttention(
+                dim,
+                n_heads,
+                head_dim,
+                bias=attention_bias,
+                dropout=dropout,
+                rope_dim=head_dim,
+                cross_attention_dim=cross_attention_dim,
+            )
+        self.mlp_norm = RMSNorm(dim, 1e-6)
+        self.mlp = LlamaMLP(dim, multiple_of=mlp_multiple_of, dropout=dropout)
+        self.use_ada_layer_norm_single = use_ada_layer_norm_single
+        if self.use_ada_layer_norm_single:
+            self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5)
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        timestep: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.use_ada_layer_norm_single:
+            batch_size = x.shape[0]
+            # timestep: [B, 6*D]
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+                self.scale_shift_table[None] + timestep.reshape(batch_size, 6, -1)
+            ).chunk(6, dim=1)
+
+            # Self-Attention with modulation and gating
+            norm_hidden_states = self.attn_norm(x)
+            norm_hidden_states = norm_hidden_states * (1 + scale_msa) + shift_msa
+            h = self.attn(norm_hidden_states, attention_mask=attention_mask)
+            h = gate_msa * h
+            x = x + h
+
+            # MLP with modulation and gating
+            norm_hidden_states = self.mlp_norm(x)
+            norm_hidden_states = norm_hidden_states * (1 + scale_mlp) + shift_mlp
+            h = self.mlp(norm_hidden_states)
+            h = gate_mlp * h
+            x = x + h
+            return x
+        else:
+            h = self.attn(self.attn_norm(x), attention_mask=attention_mask)
+            x = x + h
+            h = self.mlp(self.mlp_norm(x))
+            x = x + h
+            return x
+
+
+class ProjectLayer(nn.Module):
+    def __init__(self, hidden_size, filter_size, kernel_size=1, dropout=0.0):
+        super().__init__()
+        self.kernel_size = kernel_size
+        self.dropout = dropout
+        self.ffn_1 = nn.Conv1d(
+            hidden_size, filter_size, kernel_size, padding=kernel_size // 2
+        )
+        self.ffn_2 = nn.Linear(filter_size, filter_size)
+
+    def forward(self, x):
+        x = self.ffn_1(x.transpose(1, 2)).transpose(1, 2)
+        x = x * self.kernel_size**-0.5
+        x = self.ffn_2(x)
+        return x
+
+
+class LlamaTransformer(nn.Module):
+    def __init__(
+        self,
+        num_attention_heads: int,
+        attention_head_dim: int,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 12,
+        num_layers_2: int = 2,
+        dropout: float = 0.0,
+        cross_attention_dim: Optional[int] = None,
+        norm_type: str = "layer_norm",
+    ):
+        super().__init__()
+        inner_dim = num_attention_heads * attention_head_dim
+        inner_dim_2 = inner_dim * 2
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.inner_dim = inner_dim
+        self.inner_dim_2 = inner_dim_2
+        self.dropout = dropout
+
+        self.proj_in = ProjectLayer(in_channels, inner_dim, kernel_size=3)
+
+        use_ada_single = norm_type == "ada_norm_single"
+        self.transformer_blocks = nn.ModuleList(
+            [
+                LlamaTransformerBlock(
+                    dim=inner_dim,
+                    n_heads=num_attention_heads,
+                    head_dim=attention_head_dim,
+                    dropout=dropout,
+                    attention_bias=False,
+                    cross_attention_dim=cross_attention_dim,
+                    use_ada_layer_norm_single=use_ada_single,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.transformer_blocks_2 = nn.ModuleList(
+            [
+                LlamaTransformerBlock(
+                    dim=inner_dim_2,
+                    n_heads=num_attention_heads,
+                    head_dim=attention_head_dim * 2,
+                    dropout=dropout,
+                    attention_bias=False,
+                    cross_attention_dim=cross_attention_dim,
+                    use_ada_layer_norm_single=use_ada_single,
+                )
+                for _ in range(num_layers_2)
+            ]
+        )
+
+        self.connection_proj = ProjectLayer(
+            in_channels + inner_dim, inner_dim_2, kernel_size=3
+        )
+        self.norm_out = nn.LayerNorm(inner_dim, elementwise_affine=False, eps=1e-6)
+        self.norm_out_2 = nn.LayerNorm(inner_dim_2, elementwise_affine=False, eps=1e-6)
+        self.scale_shift_table = nn.Parameter(
+            torch.randn(2, inner_dim) / inner_dim**0.5
+        )
+        self.scale_shift_table_2 = nn.Parameter(
+            torch.randn(2, inner_dim_2) / inner_dim_2**0.5
+        )
+        self.proj_out = ProjectLayer(inner_dim_2, out_channels, kernel_size=3)
+        self.adaln_single = AdaLayerNormSingleFlow(inner_dim)
+        self.adaln_single_2 = AdaLayerNormSingleFlow(inner_dim_2)
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        timestep: Optional[torch.LongTensor] = None,
+    ):
+        s = self.proj_in(hidden_states)
+
+        embedded_timestep = None
+        timestep_mod = None
+        if self.adaln_single is not None and timestep is not None:
+            batch_size = s.shape[0]
+            timestep_mod, embedded_timestep = self.adaln_single(
+                timestep, hidden_dtype=s.dtype
+            )
+        for blk in self.transformer_blocks:
+            s = blk(s, timestep=timestep_mod)
+
+        if embedded_timestep is None:
+            embedded_timestep = torch.zeros(
+                s.size(0), s.size(-1), device=s.device, dtype=s.dtype
+            )
+
+        shift, scale = (
+            self.scale_shift_table[None] + embedded_timestep[:, None]
+        ).chunk(2, dim=1)
+        s = self.norm_out(s)
+        s = s * (1 + scale) + shift
+
+        x = torch.cat([hidden_states, s], dim=-1)
+        x = self.connection_proj(x)
+
+        embedded_timestep_2 = None
+        timestep_mod_2 = None
+        if self.adaln_single_2 is not None and timestep is not None:
+            batch_size = x.shape[0]
+            timestep_mod_2, embedded_timestep_2 = self.adaln_single_2(
+                timestep, hidden_dtype=x.dtype
+            )
+        for blk in self.transformer_blocks_2:
+            x = blk(x, timestep=timestep_mod_2)
+
+        if embedded_timestep_2 is None:
+            embedded_timestep_2 = torch.zeros(
+                x.size(0), x.size(-1), device=x.device, dtype=x.dtype
+            )
+
+        shift_2, scale_2 = (
+            self.scale_shift_table_2[None] + embedded_timestep_2[:, None]
+        ).chunk(2, dim=1)
+        x = self.norm_out_2(x)
+        x = x * (1 + scale_2) + shift_2
+
+        out = self.proj_out(x)
+
+        return out
+
+
+class PixArtAlphaCombinedFlowEmbeddings(nn.Module):
+    def __init__(self, embedding_dim: int, size_emb_dim: int):
+        super().__init__()
+        self.flow_t_size = 512
+        self.outdim = size_emb_dim
+        self.timestep_embedder = TimestepEmbedding(
+            in_channels=self.flow_t_size, time_embed_dim=embedding_dim
+        )
+
+    def timestep_embedding(self, timesteps, max_period=10000, scale=1000):
+        half = self.flow_t_size // 2
+        freqs = torch.exp(
+            -math.log(max_period)
+            * torch.arange(start=0, end=half, device=timesteps.device)
+            / half
+        ).type(timesteps.type())
+        args = timesteps[:, None] * freqs[None] * scale
+        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+        if self.flow_t_size % 2:
+            embedding = torch.cat(
+                [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+            )
+        return embedding
+
+    def forward(self, timestep, hidden_dtype):
+        timesteps_proj = self.timestep_embedding(timestep)
+        timesteps_emb = self.timestep_embedder(timesteps_proj.to(dtype=hidden_dtype))
+        conditioning = timesteps_emb
+        return conditioning
+
+
+class AdaLayerNormSingleFlow(nn.Module):
+    def __init__(self, embedding_dim: int):
+        super().__init__()
+        self.emb = PixArtAlphaCombinedFlowEmbeddings(
+            embedding_dim, size_emb_dim=embedding_dim // 3
+        )
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
+
+    def forward(
+        self,
+        timestep: torch.Tensor,
+        hidden_dtype: Optional[torch.dtype] = None,
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
+
+        embedded_timestep = self.emb(timestep, hidden_dtype=hidden_dtype)
+        return self.linear(self.silu(embedded_timestep)), embedded_timestep
+
+
+class TimestepEmbedding(nn.Module):
+    def __init__(self, in_channels: int, time_embed_dim: int):
+        super().__init__()
+        self.linear_1 = nn.Linear(in_channels, time_embed_dim)
+        self.act = nn.SiLU()
+        self.linear_2 = nn.Linear(time_embed_dim, time_embed_dim)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x = self.linear_1(x)
+        x = self.act(x)
+        x = self.linear_2(x)
+        return x
+
+
+class Timesteps(nn.Module):
+    def __init__(
+        self,
+        num_channels: int,
+        flip_sin_to_cos: bool = True,
+        downscale_freq_shift: float = 0,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.flip_sin_to_cos = flip_sin_to_cos
+        self.downscale_freq_shift = downscale_freq_shift
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        half_dim = self.num_channels // 2
+        exponent = (
+            -math.log(10000)
+            * torch.arange(0, half_dim, device=timesteps.device)
+            / (half_dim - self.downscale_freq_shift)
+        )
+        emb = torch.exp(exponent)[None, :] * timesteps[:, None]
+        if self.flip_sin_to_cos:
+            emb = torch.cat([torch.cos(emb), torch.sin(emb)], dim=-1)
+        else:
+            emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        if self.num_channels % 2 == 1:
+            emb = torch.nn.functional.pad(emb, (0, 1))
+        return emb

src/heartlib/heartmula/configuration_heartmula.py

@@ -0,0 +1,23 @@
+from transformers.configuration_utils import PretrainedConfig
+
+
+class HeartMuLaConfig(PretrainedConfig):
+    model_type = "heartmula"
+
+    def __init__(
+        self,
+        backbone_flavor: str = "llama-3B",
+        decoder_flavor: str = "llama-300M",
+        text_vocab_size: int = 128256,
+        audio_vocab_size: int = 8197,
+        audio_num_codebooks: int = 8,
+        muq_dim: int = 512,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.backbone_flavor = backbone_flavor
+        self.decoder_flavor = decoder_flavor
+        self.text_vocab_size = text_vocab_size
+        self.audio_vocab_size = audio_vocab_size
+        self.audio_num_codebooks = audio_num_codebooks
+        self.muq_dim = muq_dim

src/heartlib/heartmula/modeling_heartmula.py

@@ -0,0 +1,316 @@
+import torch
+import torch.nn as nn
+from .configuration_heartmula import HeartMuLaConfig
+from transformers.modeling_utils import PreTrainedModel
+import torch
+import torch.nn as nn
+import torchtune
+from torchtune.models import llama3_2
+
+
+def llama3_2_3B() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=28,
+        num_heads=24,
+        num_kv_heads=8,
+        embed_dim=3072,
+        max_seq_len=8192,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+def llama3_2_300M() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=3,
+        num_heads=8,
+        num_kv_heads=4,
+        embed_dim=3072,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+def llama3_2_7B() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=32,
+        num_heads=32,
+        num_kv_heads=8,
+        embed_dim=4096,
+        max_seq_len=8192,
+        intermediate_dim=14336,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )
+
+
+def llama3_2_400M() -> torchtune.modules.transformer.TransformerDecoder:
+    return llama3_2.llama3_2(
+        vocab_size=128_256,
+        num_layers=4,
+        num_heads=8,
+        num_kv_heads=4,
+        embed_dim=3072,
+        max_seq_len=2048,
+        intermediate_dim=8192,
+        attn_dropout=0.0,
+        norm_eps=1e-5,
+        rope_base=500_000,
+        scale_factor=32,
+    )  # 减少了num_heads和num_kv_heads之间的倍速，提升了精确度，但降低了效率
+
+
+FLAVORS = {
+    "llama-3B": llama3_2_3B,
+    "llama-300M": llama3_2_300M,
+    "llama-7B": llama3_2_7B,
+    "llama-400M": llama3_2_400M,
+}
+
+
+def _prepare_transformer(model):
+    embed_dim = model.tok_embeddings.embedding_dim
+    model.tok_embeddings = nn.Identity()
+    model.output = nn.Identity()
+    return model, embed_dim
+
+
+def _create_causal_mask(seq_len: int, device: torch.device):
+    return torch.tril(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device))
+
+
+def _index_causal_mask(mask: torch.Tensor, input_pos: torch.Tensor):
+    r = mask[input_pos, :]
+    return r
+
+
+def _multinomial_sample_one_no_sync(
+    probs,
+):  # Does multinomial sampling without a cuda synchronization
+    q = torch.empty_like(probs).exponential_(1)
+    return torch.argmax(probs / q, dim=-1, keepdim=True).to(dtype=torch.int)
+
+
+def sample_topk(logits: torch.Tensor, topk: int, temperature: float):
+    logits = logits / temperature
+
+    filter_value: float = -float("Inf")
+    indices_to_remove = logits < torch.topk(logits, topk)[0][..., -1, None]
+    scores_processed = logits.masked_fill(indices_to_remove, filter_value)
+    scores_processed = torch.nn.functional.log_softmax(scores_processed, dim=-1)
+    probs = torch.nn.functional.softmax(scores_processed, dim=-1)
+
+    sample_token = _multinomial_sample_one_no_sync(probs)
+    return sample_token
+
+
+class HeartMuLa(PreTrainedModel):
+    config_class = HeartMuLaConfig
+
+    def __init__(
+        self,
+        config: HeartMuLaConfig,
+    ):
+        super(HeartMuLa, self).__init__(config)
+
+        self.config = config
+
+        self.backbone, backbone_dim = _prepare_transformer(
+            FLAVORS[config.backbone_flavor]()
+        )
+        self.decoder, decoder_dim = _prepare_transformer(
+            FLAVORS[config.decoder_flavor]()
+        )
+
+        self.text_embeddings = nn.Embedding(config.text_vocab_size, backbone_dim)
+        self.audio_embeddings = nn.Embedding(
+            config.audio_vocab_size * config.audio_num_codebooks, backbone_dim
+        )
+        self.unconditional_text_embedding = nn.Embedding(1, backbone_dim)
+
+        self.projection = nn.Linear(backbone_dim, decoder_dim, bias=False)
+        self.codebook0_head = nn.Linear(
+            backbone_dim, config.audio_vocab_size, bias=False
+        )
+        self.audio_head = nn.Parameter(
+            torch.empty(
+                config.audio_num_codebooks - 1, decoder_dim, config.audio_vocab_size
+            )
+        )
+        self.muq_linear = nn.Linear(config.muq_dim, backbone_dim)
+        self.post_init()
+
+    def setup_caches(self, max_batch_size: int):
+        dtype = next(self.parameters()).dtype
+        device = next(self.parameters()).device
+
+        try:
+            self.reset_caches()
+        except RuntimeError:
+            pass
+
+        with device:
+            self.backbone.setup_caches(max_batch_size, dtype)
+            self.decoder.setup_caches(
+                max_batch_size,
+                dtype,
+                decoder_max_seq_len=self.config.audio_num_codebooks,
+            )
+
+        self.register_buffer(
+            "backbone_causal_mask",
+            _create_causal_mask(self.backbone.max_seq_len, device),
+        )
+        self.register_buffer(
+            "decoder_causal_mask",
+            _create_causal_mask(self.config.audio_num_codebooks, device),
+        )
+
+    def generate_frame(
+        self,
+        tokens: torch.Tensor,
+        tokens_mask: torch.Tensor,
+        input_pos: torch.Tensor,
+        temperature: float,
+        topk: int,
+        cfg_scale: float,
+        continuous_segments: torch.Tensor = None,
+        starts=None,
+    ) -> torch.Tensor:
+        b, s, _ = tokens.size()
+
+        assert self.backbone.caches_are_enabled(), "backbone caches are not enabled"
+        curr_backbone_mask = _index_causal_mask(self.backbone_causal_mask, input_pos)
+
+        uncond_mask = None
+        if cfg_scale > 1.0 and b > 1:
+            actual_B = b // 2
+            uncond_mask = torch.cat(
+                [
+                    torch.zeros(actual_B, dtype=torch.bool, device=tokens.device),
+                    torch.ones(actual_B, dtype=torch.bool, device=tokens.device),
+                ]
+            )
+
+        embeds = self._embed_tokens(tokens, uncond_mask=uncond_mask)
+        masked_embeds = embeds * tokens_mask.unsqueeze(-1)
+        h = masked_embeds.sum(dim=2, dtype=embeds.dtype)  # merge
+        if continuous_segments is not None:
+            continuous_segments = self.muq_linear(continuous_segments)
+            if uncond_mask is not None:
+                uncond_embed = self.unconditional_text_embedding(
+                    torch.zeros(1, device=tokens.device, dtype=torch.long)
+                )
+                mask_expanded = uncond_mask.view(b, 1).expand_as(continuous_segments)
+                continuous_segments = torch.where(
+                    mask_expanded, uncond_embed, continuous_segments
+                )
+            batch_indices = torch.arange(h.shape[0], device=h.device)
+            h[batch_indices, starts] = continuous_segments
+        h = self.backbone(h, input_pos=input_pos, mask=curr_backbone_mask)
+        last_h = h[:, -1, :]  # the last frame
+        c0_logits = self.codebook0_head(last_h)  # only predict the audio part
+
+        if cfg_scale > 1.0 and b > 1 and (b % 2 == 0):
+            actual_B = b // 2
+            cond_logits = c0_logits[:actual_B, :]
+            uncond_logits = c0_logits[actual_B:, :]
+            guided_logits = uncond_logits + (cond_logits - uncond_logits) * cfg_scale
+            c0_sample = sample_topk(guided_logits, topk, temperature)
+            c0_sample = c0_sample.repeat(
+                2, 1
+            )  # repeat to both branches to keep alignment
+        else:
+            c0_sample = sample_topk(c0_logits, topk, temperature)
+
+        c0_embed = self._embed_audio(0, c0_sample)
+
+        self.decoder.reset_caches()
+        curr_h = torch.cat([last_h.unsqueeze(1), c0_embed], dim=1)
+        curr_sample = c0_sample.clone()
+        curr_pos = (
+            torch.arange(0, curr_h.size(1), device=curr_h.device)
+            .unsqueeze(0)
+            .repeat(curr_h.size(0), 1)
+        )
+        curr_h = curr_h.to(embeds.dtype)
+        for i in range(1, self.config.audio_num_codebooks):
+            curr_decoder_mask = _index_causal_mask(self.decoder_causal_mask, curr_pos)
+            decoder_h = self.decoder(
+                self.projection(curr_h), input_pos=curr_pos, mask=curr_decoder_mask
+            )
+            ci_logits = torch.mm(decoder_h[:, -1, :], self.audio_head[i - 1])
+            if cfg_scale > 1.0 and b > 1 and (b % 2 == 0):
+                actual_B = b // 2
+                cond_ci = ci_logits[:actual_B, :]
+                uncond_ci = ci_logits[actual_B:, :]
+                guided_ci = uncond_ci + (cond_ci - uncond_ci) * cfg_scale
+
+                ci_sample = sample_topk(guided_ci, topk, temperature)
+                ci_sample = ci_sample.repeat(2, 1)
+            else:
+                ci_sample = sample_topk(ci_logits, topk, temperature)
+            ci_embed = self._embed_audio(i, ci_sample)
+            curr_h = ci_embed
+            curr_sample = torch.cat([curr_sample, ci_sample], dim=1)
+            curr_pos = curr_pos[:, -1:] + 1
+
+        return curr_sample
+
+    def reset_caches(self):
+        self.backbone.reset_caches()
+        self.decoder.reset_caches()
+
+    def _embed_local_audio(self, tokens):
+        """the token from 0-30"""
+        audio_tokens = tokens + (
+            self.config.audio_vocab_size
+            * torch.arange(self.config.audio_num_codebooks - 1, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.config.audio_num_codebooks - 1, -1
+        )
+        return audio_embeds
+
+    def _embed_audio(self, codebook: int, tokens: torch.Tensor) -> torch.Tensor:
+        return self.audio_embeddings(tokens + codebook * self.config.audio_vocab_size)
+
+    def _embed_tokens(
+        self, tokens: torch.Tensor, uncond_mask: torch.Tensor | None
+    ) -> torch.Tensor:
+        B, S, _ = tokens.size()
+        text_embeds = self.text_embeddings(tokens[:, :, -1])
+
+        if uncond_mask is not None:
+            uncond_text_embed = self.unconditional_text_embedding(
+                torch.zeros(1, device=tokens.device, dtype=torch.long)
+            )
+            mask_expanded = uncond_mask.view(B, 1, 1).expand_as(text_embeds)
+            text_embeds = torch.where(
+                mask_expanded,
+                uncond_text_embed,
+                text_embeds,
+            )
+
+        text_embeds = text_embeds.unsqueeze(-2)
+
+        audio_tokens = tokens[:, :, :-1] + (
+            self.config.audio_vocab_size
+            * torch.arange(self.config.audio_num_codebooks, device=tokens.device)
+        )
+        audio_embeds = self.audio_embeddings(audio_tokens.view(-1)).reshape(
+            tokens.size(0), tokens.size(1), self.config.audio_num_codebooks, -1
+        )
+        return torch.cat([audio_embeds, text_embeds], dim=-2)

src/heartlib/pipelines/lyrics_transcription.py

@@ -0,0 +1,40 @@
+from transformers.pipelines.automatic_speech_recognition import (
+    AutomaticSpeechRecognitionPipeline,
+)
+from transformers.models.whisper.modeling_whisper import WhisperForConditionalGeneration
+from transformers.models.whisper.processing_whisper import WhisperProcessor
+import torch
+import os
+
+
+class HeartTranscriptorPipeline(AutomaticSpeechRecognitionPipeline):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+
+    @classmethod
+    def from_pretrained(
+        cls, pretrained_path: str, device: torch.device, dtype: torch.dtype
+    ):
+        if os.path.exists(
+            hearttranscriptor_path := os.path.join(
+                pretrained_path, "HeartTranscriptor-oss"
+            )
+        ):
+            model = WhisperForConditionalGeneration.from_pretrained(
+                hearttranscriptor_path, torch_dtype=dtype, low_cpu_mem_usage=True
+            )
+            processor = WhisperProcessor.from_pretrained(hearttranscriptor_path)
+        else:
+            raise FileNotFoundError(
+                f"Expected to find checkpoint for HeartTranscriptor at {hearttranscriptor_path} but not found. Please check your folder {pretrained_path}."
+            )
+
+        return cls(
+            model=model,
+            tokenizer=processor.tokenizer,
+            feature_extractor=processor.feature_extractor,
+            device=device,
+            dtype=dtype,
+            chunk_length_s=30,
+            batch_size=16,
+        )

src/heartlib/pipelines/music_generation.py

@@ -0,0 +1,383 @@
+from tokenizers import Tokenizer
+from ..heartmula.modeling_heartmula import HeartMuLa
+from ..heartcodec.modeling_heartcodec import HeartCodec
+import torch
+from typing import Dict, Any, Optional, Union
+import os
+from dataclasses import dataclass
+from tqdm import tqdm
+import torchaudio
+import json
+from contextlib import contextmanager
+import gc
+
+
+def _resolve_paths(pretrained_path: str, version: str):
+
+    heartmula_path = os.path.join(pretrained_path, f"HeartMuLa-oss-{version}")
+    heartcodec_path = os.path.join(pretrained_path, "HeartCodec-oss")
+    tokenizer_path = os.path.join(pretrained_path, "tokenizer.json")
+    gen_config_path = os.path.join(pretrained_path, "gen_config.json")
+
+    if not os.path.exists(heartmula_path):
+        raise FileNotFoundError(
+            f"Expected to find checkpoint for HeartMuLa at {heartmula_path} but not found. Please check your folder {pretrained_path}."
+        )
+    if not os.path.exists(heartcodec_path):
+        raise FileNotFoundError(
+            f"Expected to find checkpoint for HeartCodec at {heartcodec_path} but not found. Please check your folder {pretrained_path}."
+        )
+    if not os.path.isfile(tokenizer_path):
+        raise FileNotFoundError(
+            f"Expected to find tokenizer.json for HeartMuLa at {tokenizer_path} but not found. Please check your folder {pretrained_path}."
+        )
+    if not os.path.isfile(gen_config_path):
+        raise FileNotFoundError(
+            f"Expected to find gen_config.json for HeartMuLa at {gen_config_path} but not found. Please check your folder {pretrained_path}."
+        )
+
+    return heartmula_path, heartcodec_path, tokenizer_path, gen_config_path
+
+
+def _resolve_devices(
+    device: Union[torch.device, Dict[str, torch.device]], lazy_load: bool
+):
+    if isinstance(device, torch.device):
+        print(f"All model components will be loaded to device: {device}.")
+        mula_device = device
+        codec_device = device
+    elif isinstance(device, dict):
+        print("Model components will be loaded to devices as specified:")
+        for k, v in device.items():
+            print(f"  {k}: {v}")
+        mula_device = device["mula"]
+        codec_device = device["codec"]
+    else:
+        raise ValueError(
+            "device must be either torch.device or Dict[str, torch.device]"
+        )
+
+    single_device = mula_device == codec_device
+    if not single_device:
+        print(
+            f"HeartMuLa and HeartCodec will be loaded to different devices. In this case, lazy_load is turned off."
+        )
+        lazy_load = False
+
+    return mula_device, codec_device, lazy_load
+
+
+@dataclass
+class HeartMuLaGenConfig:
+    text_bos_id: int = 128000
+    text_eos_id: int = 128001
+    audio_eos_id: int = 8193
+    empty_id: int = 0
+
+    @classmethod
+    def from_file(cls, path: str):
+        with open(path, encoding="utf-8") as fp:
+            data = json.load(fp)
+        return cls(**data)
+
+
+class HeartMuLaGenPipeline:
+    def __init__(
+        self,
+        heartmula_path: str,
+        heartcodec_path: str,
+        heartmula_device: torch.device,
+        heartcodec_device: torch.device,
+        heartmula_dtype: torch.dtype,
+        heartcodec_dtype: torch.dtype,
+        lazy_load: bool,
+        muq_mulan: Optional[Any],
+        text_tokenizer: Tokenizer,
+        config: HeartMuLaGenConfig,
+    ):
+
+        self.muq_mulan = muq_mulan
+        self.text_tokenizer = text_tokenizer
+        self.config = config
+
+        # Remain fixed here for simplicity.
+        self._parallel_number = 8 + 1
+        self._muq_dim = 512
+
+        self.mula_dtype = heartmula_dtype
+        self.mula_path = heartmula_path
+        self.mula_device = heartmula_device
+        self.codec_dtype = heartcodec_dtype
+        self.codec_path = heartcodec_path
+        self.codec_device = heartcodec_device
+
+        self._mula: Optional[HeartMuLa] = None
+        self._codec: Optional[HeartCodec] = None
+        if not lazy_load:
+            print(
+                f"You have set lazy_load = False. Loading HeartMuLa and HeartCodec onto device..."
+            )
+            self._mula = HeartMuLa.from_pretrained(
+                self.mula_path,
+                device_map=self.mula_device,
+                torch_dtype=self.mula_dtype,
+            )
+            self._codec = HeartCodec.from_pretrained(
+                self.codec_path,
+                device_map=self.codec_device,
+                torch_dtype=self.codec_dtype,
+            )
+        self.lazy_load = lazy_load
+
+    @property
+    def mula(self) -> HeartMuLa:
+        if isinstance(self._mula, HeartMuLa):
+            return self._mula
+        self._mula = HeartMuLa.from_pretrained(
+            self.mula_path,
+            device_map=self.mula_device,
+            torch_dtype=self.mula_dtype,
+        )
+        return self._mula
+
+    @property
+    def codec(self) -> HeartCodec:
+        if isinstance(self._codec, HeartCodec):
+            return self._codec
+        self._codec = HeartCodec.from_pretrained(
+            self.codec_path,
+            device_map=self.codec_device,
+            torch_dtype=self.codec_dtype,
+        )
+        return self._codec
+
+    def _unload(self):
+        if not self.lazy_load:
+            return
+        if isinstance(self._mula, HeartMuLa):
+            print(f"You have set lazy_load=True. Unloading HeartMuLa from device.")
+            print(
+                f"CUDA memory before unloading: {torch.cuda.memory_allocated(self.mula_device) / 1024**3:.2f} GB"
+            )
+            del self._mula
+            gc.collect()
+            torch.cuda.empty_cache()
+            print(
+                f"CUDA memory after unloading: {torch.cuda.memory_allocated(self.mula_device) / 1024**3:.2f} GB"
+            )
+            self._mula = None
+        if isinstance(self._codec, HeartCodec):
+            print(f"You have set lazy_load=True. Unloading HeartCodec from device.")
+            print(
+                f"CUDA memory before unloading: {torch.cuda.memory_allocated(self.codec_device) / 1024**3:.2f} GB"
+            )
+            del self._codec
+            gc.collect()
+            torch.cuda.empty_cache()
+            print(
+                f"CUDA memory after unloading: {torch.cuda.memory_allocated(self.codec_device) / 1024**3:.2f} GB"
+            )
+            self._codec = None
+        return
+
+    def _sanitize_parameters(self, **kwargs):
+        preprocess_kwargs = {"cfg_scale": kwargs.get("cfg_scale", 1.5)}
+        forward_kwargs = {
+            "max_audio_length_ms": kwargs.get("max_audio_length_ms", 120_000),
+            "temperature": kwargs.get("temperature", 1.0),
+            "topk": kwargs.get("topk", 50),
+            "cfg_scale": kwargs.get("cfg_scale", 1.5),
+        }
+        postprocess_kwargs = {
+            "save_path": kwargs.get("save_path", "output.mp3"),
+        }
+        return preprocess_kwargs, forward_kwargs, postprocess_kwargs
+
+    def preprocess(self, inputs: Dict[str, Any], cfg_scale: float):
+
+        # process tags
+        tags = inputs["tags"]
+        if os.path.isfile(tags):
+            with open(tags, encoding="utf-8") as fp:
+                tags = fp.read()
+        assert isinstance(tags, str), f"tags must be a string, but got {type(tags)}"
+
+        tags = tags.lower()
+        # encapsulate with special <tag> and </tag> tokens
+        if not tags.startswith("<tag>"):
+            tags = f"<tag>{tags}"
+        if not tags.endswith("</tag>"):
+            tags = f"{tags}</tag>"
+
+        tags_ids = self.text_tokenizer.encode(tags).ids
+        if tags_ids[0] != self.config.text_bos_id:
+            tags_ids = [self.config.text_bos_id] + tags_ids
+        if tags_ids[-1] != self.config.text_eos_id:
+            tags_ids = tags_ids + [self.config.text_eos_id]
+
+        # process reference audio
+        ref_audio = inputs.get("ref_audio", None)
+        if ref_audio is not None:
+            raise NotImplementedError("ref_audio is not supported yet.")
+        muq_embed = torch.zeros([self._muq_dim], dtype=self.mula_dtype)
+        muq_idx = len(tags_ids)
+
+        # process lyrics
+        lyrics = inputs["lyrics"]
+        if os.path.isfile(lyrics):
+            with open(lyrics, encoding="utf-8") as fp:
+                lyrics = fp.read()
+        assert isinstance(
+            lyrics, str
+        ), f"lyrics must be a string, but got {type(lyrics)}"
+        lyrics = lyrics.lower()
+
+        lyrics_ids = self.text_tokenizer.encode(lyrics).ids
+        if lyrics_ids[0] != self.config.text_bos_id:
+            lyrics_ids = [self.config.text_bos_id] + lyrics_ids
+        if lyrics_ids[-1] != self.config.text_eos_id:
+            lyrics_ids = lyrics_ids + [self.config.text_eos_id]
+
+        # cat them together. tags, ref_audio, lyrics
+        prompt_len = len(tags_ids) + 1 + len(lyrics_ids)
+
+        tokens = torch.zeros([prompt_len, self._parallel_number], dtype=torch.long)
+        tokens[: len(tags_ids), -1] = torch.tensor(tags_ids)
+        tokens[len(tags_ids) + 1 :, -1] = torch.tensor(lyrics_ids)
+
+        tokens_mask = torch.zeros_like(tokens, dtype=torch.bool)
+        tokens_mask[:, -1] = True
+
+        bs_size = 2 if cfg_scale != 1.0 else 1
+
+        def _cfg_cat(tensor: torch.Tensor, cfg_scale: float):
+            tensor = tensor.unsqueeze(0)
+            if cfg_scale != 1.0:
+                tensor = torch.cat([tensor, tensor], dim=0)
+            return tensor
+
+        return {
+            "tokens": _cfg_cat(tokens, cfg_scale),
+            "tokens_mask": _cfg_cat(tokens_mask, cfg_scale),
+            "muq_embed": _cfg_cat(muq_embed, cfg_scale),
+            "muq_idx": [muq_idx] * bs_size,
+            "pos": _cfg_cat(torch.arange(prompt_len, dtype=torch.long), cfg_scale),
+        }
+
+    def _forward(
+        self,
+        model_inputs: Dict[str, Any],
+        max_audio_length_ms: int,
+        temperature: float,
+        topk: int,
+        cfg_scale: float,
+    ):
+        prompt_tokens = model_inputs["tokens"].to(self.mula_device)
+        prompt_tokens_mask = model_inputs["tokens_mask"].to(self.mula_device)
+        continuous_segment = model_inputs["muq_embed"].to(self.mula_device)
+        starts = model_inputs["muq_idx"]
+        prompt_pos = model_inputs["pos"].to(self.mula_device)
+        frames = []
+
+        bs_size = 2 if cfg_scale != 1.0 else 1
+        self.mula.setup_caches(bs_size)
+        with torch.autocast(device_type=self.mula_device.type, dtype=self.mula_dtype):
+            curr_token = self.mula.generate_frame(
+                tokens=prompt_tokens,
+                tokens_mask=prompt_tokens_mask,
+                input_pos=prompt_pos,
+                temperature=temperature,
+                topk=topk,
+                cfg_scale=cfg_scale,
+                continuous_segments=continuous_segment,
+                starts=starts,
+            )
+        frames.append(curr_token[0:1,])
+
+        def _pad_audio_token(token: torch.Tensor):
+            padded_token = (
+                torch.ones(
+                    (token.shape[0], self._parallel_number),
+                    device=token.device,
+                    dtype=torch.long,
+                )
+                * self.config.empty_id
+            )
+            padded_token[:, :-1] = token
+            padded_token = padded_token.unsqueeze(1)
+            padded_token_mask = torch.ones_like(
+                padded_token, device=token.device, dtype=torch.bool
+            )
+            padded_token_mask[..., -1] = False
+            return padded_token, padded_token_mask
+
+        max_audio_frames = max_audio_length_ms // 80
+
+        for i in tqdm(range(max_audio_frames)):
+            curr_token, curr_token_mask = _pad_audio_token(curr_token)
+            with torch.autocast(
+                device_type=self.mula_device.type, dtype=self.mula_dtype
+            ):
+                curr_token = self.mula.generate_frame(
+                    tokens=curr_token,
+                    tokens_mask=curr_token_mask,
+                    input_pos=prompt_pos[..., -1:] + i + 1,
+                    temperature=temperature,
+                    topk=topk,
+                    cfg_scale=cfg_scale,
+                    continuous_segments=None,
+                    starts=None,
+                )
+            if torch.any(curr_token[0:1, :] >= self.config.audio_eos_id):
+                break
+            frames.append(curr_token[0:1,])
+        frames = torch.stack(frames).permute(1, 2, 0).squeeze(0)
+        self._unload()
+        return {"frames": frames}
+
+    def postprocess(self, model_outputs: Dict[str, Any], save_path: str):
+        frames = model_outputs["frames"].to(self.codec_device)
+        wav = self.codec.detokenize(frames)
+        self._unload()
+        torchaudio.save(save_path, wav.to(torch.float32).cpu(), 48000)
+
+    def __call__(self, inputs: Dict[str, Any], **kwargs):
+        preprocess_kwargs, forward_kwargs, postprocess_kwargs = (
+            self._sanitize_parameters(**kwargs)
+        )
+        model_inputs = self.preprocess(inputs, **preprocess_kwargs)
+        model_outputs = self._forward(model_inputs, **forward_kwargs)
+        self.postprocess(model_outputs, **postprocess_kwargs)
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        pretrained_path: str,
+        device: Union[torch.device, Dict[str, torch.device]],
+        dtype: Union[torch.dtype, Dict[str, torch.dtype]],
+        version: str,
+        lazy_load: bool = False,
+    ):
+
+        mula_path, codec_path, tokenizer_path, gen_config_path = _resolve_paths(
+            pretrained_path, version
+        )
+        mula_device, codec_device, lazy_load = _resolve_devices(device, lazy_load)
+        tokenizer = Tokenizer.from_file(tokenizer_path)
+        gen_config = HeartMuLaGenConfig.from_file(gen_config_path)
+
+        mula_dtype = dtype["mula"] if isinstance(dtype, dict) else dtype
+        codec_dtype = dtype["codec"] if isinstance(dtype, dict) else dtype
+
+        return cls(
+            heartmula_path=mula_path,
+            heartcodec_path=codec_path,
+            heartmula_device=mula_device,
+            heartcodec_device=codec_device,
+            lazy_load=lazy_load,
+            muq_mulan=None,
+            text_tokenizer=tokenizer,
+            config=gen_config,
+            heartmula_dtype=mula_dtype,
+            heartcodec_dtype=codec_dtype,
+        )