Upload folder using huggingface_hub

README.md

@@ -1,3 +1,200 @@
 ---
 license: apache-2.0
+tags:
+- music-editing
+- audio
+- multimodal
+- conversational-ai
+- lora
+- musicgen
+- qwen2.5-omni
+language:
+- en
+pipeline_tag: audio-to-audio
+datasets:
+- OpenRB-Lab/AURA-Chat-Edit
 ---
+
+# AURA: Unified Multimodal Framework for Conversational Music Editing
+
+[![arXiv](https://img.shields.io/badge/arXiv-TODO-b31b1b.svg)](https://arxiv.org/abs/TODO)
+[![GitHub](https://img.shields.io/badge/GitHub-AURA-blue.svg)](https://github.com/OpenRB-Lab/AURA)
+[![Dataset](https://img.shields.io/badge/🤗_Dataset-AURA--Chat--Edit-yellow.svg)](https://huggingface.co/datasets/OpenRB-Lab/AURA-Chat-Edit)
+
+
+## Overview
+
+**AURA** is a conversational music-editing agent that listens to a song and a natural-language instruction, replies conversationally, and renders the edited audio. The system consists of three components:
+
+- **Thinker** — A [Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B) model fine-tuned with LoRA (r=16, alpha=32). It processes audio and text, generates conversational replies, and emits typed edit-token blocks `[EDIT_<KIND>][EDIT_0..7]` (7 kinds: ADD / REMOVE / REPLACE / EXTRACT / REBALANCE / EFFECT / MOOD).
+- **Bridge** — A dual-stream fusion MusicGen decoder based on [facebook/musicgen-medium](https://huggingface.co/facebook/musicgen-medium). The 9 hidden states at the edit tokens condition the bridge via **BiFAM** (Bi-FiLM Attention Module: shared-query dual attention + FiLM modulation) and cross-attention K/V with LoRA (r=64, alpha=128). Includes learned projectors (258 MB) mapping from the thinker's hidden dimension to MusicGen's space.
+- **Classifier** — An `EditSemanticClassifier` (two-head: edit kind + instrument) used by the programmatic planner for stem routing.
+
+Localized edits are code-anchored outside the requested segment and seam-crossfaded via a stem-hybrid executor (HTDemucs-6s separation).
+
+## Model Checkpoints
+
+| File | Description | Size |
+|------|-------------|------|
+| `config.yaml` | Training configuration (paths, hyperparameters) | 1 KB |
+| `thinker/adapter_config.json` | Thinker LoRA configuration | 1 KB |
+| `thinker/adapter_model.safetensors` | Thinker LoRA weights (Qwen2.5-Omni-7B, r=16) | 2.0 GB |
+| `bridge/projectors.pt` | Learned projectors (d_llm=3584 → d_musicgen=2048) + FiLM MLPs/alphas/gates | 258 MB |
+| `bridge/lora/adapter_config.json` | Bridge LoRA configuration | 1 KB |
+| `bridge/lora/adapter_model.safetensors` | Bridge LoRA weights (MusicGen encoder_attn k/v, r=64) | 37 MB |
+| `classifier/classifier.pt` | EditSemanticClassifier (kind + instrument heads) | 14 MB |
+
+**Total checkpoint size: ~2.3 GB** (adapters only — base models downloaded separately)
+
+### Thinker Details
+
+- **Base model**: [Qwen/Qwen2.5-Omni-7B](https://huggingface.co/Qwen/Qwen2.5-Omni-7B)
+- **LoRA config**: r=16, alpha=32, dropout=0.05
+- **Target modules**: `q_proj`, `k_proj`, `v_proj`, `o_proj`, `gate_proj`, `up_proj`, `down_proj`
+- **Modules to save**: `embed_tokens`, `lm_head` (for custom edit tokens)
+- **Training**: 2-epoch SFT on dialogue data, then joint training with bridge (4000 steps, λ=0.5 convex loss)
+
+### Bridge Details
+
+- **Base model**: [facebook/musicgen-medium](https://huggingface.co/facebook/musicgen-medium) (1.5B params, frozen decoder)
+- **Fusion mechanism**: BiFAM — shared-query dual cross-attention over edit-token hidden states + gated FiLM modulation at each decoder layer
+- **LoRA config**: r=64, alpha=128, dropout=0.05, targeting `encoder_attn.{k_proj, v_proj}`
+- **Projectors**: Linear projections from thinker hidden dim (3584) to MusicGen dim (2048), plus FiLM MLP layers
+- **Cross-attention layers**: [0, 2, 4, 6, 8, 10, 12, 14]
+- **Training**: 40k steps bridge-only, then 4000 steps joint with thinker
+
+### Classifier Details
+
+- **Architecture**: Two-head classifier (edit kind: 7 classes, instrument: multi-label)
+- **Input**: 9 edit-token hidden states (pooled)
+- **Used by**: Programmatic planner for stem routing decisions
+
+## Quick Start
+
+### 1. Download base models
+
+The base models are downloaded automatically on first use, or you can pre-download them:
+
+```python
+from huggingface_hub import snapshot_download
+
+# Thinker base model (~15 GB)
+snapshot_download("Qwen/Qwen2.5-Omni-7B", cache_dir="weights")
+
+# Bridge base model (~3.3 GB)
+snapshot_download("facebook/musicgen-medium", cache_dir="weights")
+```
+
+### 2. Download AURA checkpoints
+
+```python
+from huggingface_hub import snapshot_download
+
+# Download all AURA adapters (~2.3 GB)
+repo_dir = snapshot_download("OpenRB-Lab/AURA")
+```
+
+Or download individual components:
+
+```python
+from huggingface_hub import hf_hub_download
+
+# Thinker LoRA adapter
+thinker_config = hf_hub_download("OpenRB-Lab/AURA", "thinker/adapter_config.json")
+thinker_weights = hf_hub_download("OpenRB-Lab/AURA", "thinker/adapter_model.safetensors")
+
+# Bridge projectors + LoRA
+projectors = hf_hub_download("OpenRB-Lab/AURA", "bridge/projectors.pt")
+bridge_config = hf_hub_download("OpenRB-Lab/AURA", "bridge/lora/adapter_config.json")
+bridge_weights = hf_hub_download("OpenRB-Lab/AURA", "bridge/lora/adapter_model.safetensors")
+
+# Classifier
+classifier = hf_hub_download("OpenRB-Lab/AURA", "classifier/classifier.pt")
+```
+
+### 3. Usage
+
+```python
+# Point environment to your checkpoint directory
+import os
+os.environ["AURA_QWEN"] = "path/to/aura1/thinker"
+os.environ["AURA_MG"] = "path/to/aura1/bridge"
+os.environ["AURA_CLASSIFIER"] = "path/to/aura1/classifier/classifier.pt"
+
+# Load the engine
+from serving.engine import AuraEngine
+
+engine = AuraEngine(device="cuda")
+result = engine.edit(
+    audio_path="path/to/song.wav",
+    instruction="Add a jazzy saxophone melody to the chorus",
+    guidance=2.0,
+    seed=1234,
+    max_seconds=5.0
+)
+
+# result["reply"]  -> conversational text response
+# result["wav"]    -> edited audio (float32 numpy, 32 kHz)
+# result["sr"]     -> 32000
+```
+
+### 4. Serving
+
+```bash
+# HTTP API
+API_GPU=0 API_PORT=9004 bash src/scripts/serve_musicgen_api.sh
+
+# Gradio web UI
+WORKER_URL=http://127.0.0.1:9004 \
+SFT_ADAPTER=path/to/aura1/thinker \
+WEBAPP_PORT=7862 CUDA_VISIBLE_DEVICES=1 \
+python src/edit_agent/webapp.py
+```
+
+## Training Configuration
+
+Training uses a 3-stage pipeline:
+
+1. **Stage 1 — Thinker SFT**: LoRA fine-tuning on conversational music-edit dialogues (2 epochs, lr=1e-4)
+2. **Stage 2 — Bridge**: Fusion adapter training on cached thinker hidden states (40k steps, lr=5e-5)
+3. **Stage 3 — Joint**: End-to-end training with live thinker + bridge (4k steps, loss = λ·CE_musicgen + (1−λ)·CE_LM)
+
+See `config.yaml` for the full training configuration.
+
+## Results
+
+Production checkpoints (`joint_fusion_r64/final`):
+
+| Benchmark | FAD ↓ | CLAP ↑ | SSIM ↑ |
+|-----------|-------|--------|--------|
+| IMPG Add | 1.49 | — | — |
+| IMPG Remove | 1.36 | — | — |
+| IMPG Extract | 6.13 | — | — |
+| Mixed 60-clip | 2.15 | 0.661 (MuLan cos) | — |
+
+### Fusion Ablation (IMPG Benchmark)
+
+| Method | FAD ↓ | CLAP ↑ | SSIM ↑ |
+|--------|-------|--------|--------|
+| **BiFAM (Ours)** | **0.41** | 0.218 | **0.776** |
+| Cross-Attention Only | 2.48 | **0.263** | 0.099 |
+| Concatenation Only | 13.64 | 0.181 | 0.115 |
+
+## Citation
+
+```bibtex
+@inproceedings{trinh2026aura,
+  title     = {AURA: Unified Multimodal Framework for Conversational Music Editing},
+  author    = {Trinh, Quoc-Huy and Nguyen, Minh-Van and Jha, Debesh},
+  year      = {2026},
+  note      = {arXiv preprint, arXiv:TODO (to appear)}
+}
+```
+
+## License
+
+This project is licensed under the Apache License 2.0. See [LICENSE](LICENSE) for details.
+
+**Note**: The base models have their own licenses:
+- Qwen2.5-Omni-7B: [Apache 2.0](https://huggingface.co/Qwen/Qwen2.5-Omni-7B)
+- MusicGen-medium: [CC-BY-NC 4.0](https://huggingface.co/facebook/musicgen-medium)

bridge/lora/adapter_config.json

@@ -0,0 +1,41 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": {
+    "base_model_class": "MusicgenForCausalLM",
+    "parent_library": "transformers.models.musicgen.modeling_musicgen"
+  },
+  "base_model_name_or_path": "",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 128,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": null,
+  "peft_type": "LORA",
+  "peft_version": "0.18.1",
+  "qalora_group_size": 16,
+  "r": 64,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": ".*encoder_attn\\.(k_proj|v_proj)",
+  "target_parameters": null,
+  "task_type": null,
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

bridge/lora/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f8d1ccba7e5a286137abce488cf1927f5a0a7b3f1a943fa3127804c2c171941
+size 37776848

bridge/projectors.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a1a935af50219d69f74ff8262838e1f5020edb20c4862b5860d45a9111c6a908
+size 258696730

classifier/classifier.pt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0a3881430ad35d94ed48ab466ca386fa1dfbb99cfc123ad13cea041b62e8d673
+size 14778039

config.yaml

@@ -0,0 +1,34 @@
+# Edit-agent configuration (dialogue SFT + later bridge stages)
+
+paths:
+  dataset_jsonl: data/edit_dataset/manifests/dataset.jsonl
+  dialogues_jsonl: data/edit_dataset/dialogues/dialogues.jsonl
+  bridge_cache_dir: data/edit_dataset/bridge_cache
+  ckpt_root: ckpts/edit_agent
+
+tokens:
+  k_edit: 8
+
+qwen:
+  model_id: Qwen/Qwen2.5-Omni-7B
+  attn: sdpa
+  dtype: bf16
+
+sft:
+  lora: {r: 16, alpha: 32, dropout: 0.05}
+  lr: 1.0e-4
+  batch: 1
+  grad_accum: 16
+  epochs: 2
+  warmup: 300
+  log_every: 10
+  save_every: 500
+
+# Stage A (bridge) — filled in when that stage is implemented
+bridge:
+  d_llm: 3584
+  dit_dim: 2048
+  cross_attn_layers: [0, 2, 4, 6, 8, 10, 12, 14]
+  lora: {r: 16, alpha: 32}
+  drops: {edit_cond: 0.1, src_cond: 0.1}
+  loss_w: {flow: 1.0, align: 1.0, nce: 0.5}

thinker/adapter_config.json

@@ -0,0 +1,41 @@
+{
+  "alora_invocation_tokens": null,
+  "alpha_pattern": {},
+  "arrow_config": null,
+  "auto_mapping": null,
+  "base_model_name_or_path": "Qwen/Qwen2.5-Omni-7B",
+  "bias": "none",
+  "corda_config": null,
+  "ensure_weight_tying": false,
+  "eva_config": null,
+  "exclude_modules": null,
+  "fan_in_fan_out": false,
+  "inference_mode": true,
+  "init_lora_weights": true,
+  "layer_replication": null,
+  "layers_pattern": null,
+  "layers_to_transform": null,
+  "loftq_config": {},
+  "lora_alpha": 32,
+  "lora_bias": false,
+  "lora_dropout": 0.05,
+  "megatron_config": null,
+  "megatron_core": "megatron.core",
+  "modules_to_save": [
+    "embed_tokens",
+    "lm_head"
+  ],
+  "peft_type": "LORA",
+  "peft_version": "0.18.1",
+  "qalora_group_size": 16,
+  "r": 16,
+  "rank_pattern": {},
+  "revision": null,
+  "target_modules": "model\\.layers\\.\\d+\\.(self_attn\\.(q_proj|k_proj|v_proj|o_proj)|mlp\\.(gate_proj|up_proj|down_proj))",
+  "target_parameters": null,
+  "task_type": "CAUSAL_LM",
+  "trainable_token_indices": null,
+  "use_dora": false,
+  "use_qalora": false,
+  "use_rslora": false
+}

thinker/adapter_model.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:54f10ba5b9a01da952544335847e8e0d78aef8c90960c2c391e01fd7469a9ca9
+size 2341522936