datasets/raw/my-set/
├── 01. First Song.flac
├── 01. First Song.txt
├── 02. Second Song.flac
├── 02. Second Song.txt
└── styles.json

[Verse]
Hundred guns and thousand shots
Our ship is on fire

[Chorus]
Rising flame, your desire
Watch out now

{
  "_trigger": "myset_tr1gg3r",
  "styles": {
    "01. First Song.flac": "power metal, fast, 172 bpm, ...",
    "02. Second Song.flac": "power metal, 140 bpm, ..."
  }
}

myset_tr1gg3r, <style line>
[Lyrics]
[Verse]
...

import json, shutil, sys
from pathlib import Path

name, styles_file = sys.argv[1], Path(sys.argv[2])
prep = Path.home() / "Ruach_Studio/datasets/prepared" / name
spec = json.loads(styles_file.read_text(encoding="utf-8"))
trigger, styles = spec["_trigger"], spec["styles"]
meta = json.loads((prep / "dataset.json").read_text("utf-8"))
tracks = meta["tracks"]
backup = Path.home() / f"{name}_captions_backup"
backup.mkdir(exist_ok=True)

done = 0
for t in tracks:
    cap, style = prep / f"{t['stem']}.txt", styles.get(t["file"])
    if not cap.is_file() or not style:
        print("skipped:", t["stem"], t["file"])
        continue
    if not (backup / cap.name).exists():
        shutil.copy2(cap, backup / cap.name)
    lines = cap.read_text(encoding="utf-8").split("\n")
    lines[0] = f"{trigger}, {style}"
    cap.write_text("\n".join(lines), encoding="utf-8")
    done += 1
print(f"{done} of {len(tracks)} replaced (originals: {backup})")

cd ~/Ruach_Studio
.venv/bin/python ~/apply_styles.py my-set  datasets/raw/my-set/styles.json
head -n 3 datasets/prepared/my-set/001.txt

mkdir ~/modal-ruach && cd ~/modal-ruach
python3 -m venv .venv && source .venv/bin/activate
pip install modal && modal setup

modal volume create ruach-data
modal volume put ruach-data  ~/Ruach_Studio/checkpoints/comfy/checkpoints/yue2_3b_bf16.safetensors  /models/checkpoints/yue2_3b_bf16.safetensors

"""Ruach Studio's YuE2 LoRA trainer on Modal.
Latent cache built at home; Modal gets the trainer code,
the bf16 checkpoint and the prepared set.
RUACH_GPU=L40S (default) | H100 | H200 | B200 | RTX-PRO-6000
RUACH_NGPU=2 for two cards."""
import glob, json, os, shutil, subprocess, threading, time
import modal

app = modal.App("ruach-lora")
vol = modal.Volume.from_name("ruach-data", create_if_missing=True)

LOCAL_TRAINER = os.path.expanduser(
    os.environ.get("RUACH_TRAINER", "~/Ruach_Studio/lab/trainer"))
KIT = "/workspace/ruach"
TRAINER = f"{KIT}/lab/trainer"
WORK = "/tmp/ruach"
CHECKPOINT = "/data/models/checkpoints/yue2_3b_bf16.safetensors"
CU128 = "https://download.pytorch.org/whl/cu128"

GPU = os.environ.get("RUACH_GPU", "L40S")
NGPU = int(os.environ.get("RUACH_NGPU", "1"))
GPU_SPEC = GPU if NGPU == 1 else f"{GPU}:{NGPU}"

image = (
    modal.Image.debian_slim(python_version="3.12")
    .apt_install("ffmpeg", "libsndfile1", "rsync")
    .pip_install("torch==2.11.0", "torchaudio==2.11.0", index_url=CU128)
    .pip_install("safetensors==0.8.0", "transformers==4.57.6",
                 "einops==0.8.2", "rotary-embedding-torch==0.6.5",
                 "numpy", "soundfile", "librosa", "sentencepiece",
                 "pyarrow")
    .env({"PYTHONUNBUFFERED": "1",
          "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True"})
    .add_local_dir(LOCAL_TRAINER, TRAINER, copy=True,
                   ignore=["__pycache__", "*.pyc"])
    .run_commands(f"cd {TRAINER} && python -c \"import sys; "
                  "sys.path.insert(0, '.'); import train; "
                  "import pyarrow.parquet; print('trainer imports ok')\"")
)

# every key spelled out so a change in train.py's defaults can't move it
RECIPE = {
    "epochs": 100, "save_epochs": 10,
    "rank": 64, "rank_ar": 0, "rank_nar": 0,
    "lr": 1e-4, "lr_scale": "sqrt",
    "lr_tail_epochs": 0.0, "lr_tail_floor": 0.0,
    "ar_loss_weight": 0.5, "ar_kl_weight": 0.2, "ar_lr_multiplier": 0.5,
    "train_window_frames": 1500, "chunk": 512,
    "ar_max_tokens": 0, "ar_cap_tokens": 0,
    "cot": "off", "abc_dropout": 0.5, "do_separation": False,
    "seed": 42, "max_grad_norm": 1.0, "accum": 1,
    "checkpointing": "auto",  # off while memory allows, on at an OOM
    "ema": False,
    "val_fraction": 0.0, "patience": 0, "nar_phase_epochs": 0,
    "resume_from": 0, "lazy_cache": "auto",
    "vram_meter": False, "long_off_card0": False,
}


def _sync(src, dst):
    if os.path.isdir(src):
        os.makedirs(dst, exist_ok=True)
        subprocess.run(["rsync", "-a", src + "/", dst + "/"], check=False)
    try:
        vol.commit()
        print(f"[volume] synced {time.strftime('%H:%M:%S')}", flush=True)
    except Exception as e:  # noqa: BLE001
        print(f"[volume] commit failed: {e}", flush=True)


def _captions(ds):
    caps = sorted(glob.glob(f"{ds}/*.txt"))
    def cached(c):
        stem = os.path.splitext(os.path.basename(c))[0]
        return os.path.isfile(f"{ds}/_ruach_cache/{stem}_ruach.safetensors")
    return caps, [c for c in caps if not cached(c)]


def _trainer_lack(ds):
    import sys
    if TRAINER not in sys.path:
        sys.path.insert(0, TRAINER)
    import cache as latent_cache
    return latent_cache.missing(ds)


@app.function(image=image, gpu=GPU_SPEC, cpu=8.0, memory=65536,
              volumes={"/data": vol}, timeout=86400, max_containers=1)
def train(name: str, dataset: str, checkpoint: str = CHECKPOINT,
          overrides: str = "{}"):
    import torch
    cards = torch.cuda.device_count()
    print(f"cards: {cards} x {torch.cuda.get_device_name(0)}", flush=True)

    vol.reload()
    ds = f"/data/datasets/{dataset}"
    caps, lack = _captions(ds)
    if not caps:
        raise FileNotFoundError(f"no captions in {ds}")
    if lack or _trainer_lack(ds):
        raise RuntimeError("tracks without a latent cache: "
                           "build it at home and upload the set again")
    if not os.path.isfile(checkpoint):
        raise FileNotFoundError(checkpoint)

    r = dict(RECIPE, **json.loads(overrides))
    tracks = len(caps)
    cfg = {k: v for k, v in r.items() if k not in ("epochs", "save_epochs")}
    cfg.setdefault("steps", int(r["epochs"] * tracks))
    cfg.setdefault("save_every", int(r["save_epochs"] * tracks))
    cfg.update(name=name, out=WORK, dataset=ds, checkpoint=checkpoint)

    local_out, vol_out = f"{WORK}/output/{name}", f"/data/runs/{name}"
    if os.path.isdir(vol_out):  # same --name resumes from the last save
        shutil.copytree(vol_out, local_out, dirs_exist_ok=True)
        print(f"resuming from {vol_out}", flush=True)
    os.makedirs(local_out, exist_ok=True)
    cfg_path = f"{local_out}/config.json"
    with open(cfg_path, "w", encoding="utf-8") as f:
        json.dump(cfg, f, indent=1)
    print(json.dumps(cfg, indent=1), flush=True)
    print(f"{tracks} tracks: {cfg['steps']} steps, "
          f"a save every {cfg['save_every']}", flush=True)

    script = f"{TRAINER}/train.py"
    cmd = (["python", "-m", "torch.distributed.run", "--standalone",
            f"--nproc_per_node={cards}", script, cfg_path]
           if cards > 1 else ["python", script, cfg_path])
    stop = threading.Event()

    def loop():
        while not stop.wait(600):
            _sync(local_out, vol_out)
    threading.Thread(target=loop, daemon=True).start()
    try:
        code = subprocess.run(cmd, cwd=KIT).returncode
    finally:
        stop.set()
        _sync(local_out, vol_out)
    if code != 0:
        raise RuntimeError(f"train.py exited {code}: see the log above")


@app.function(image=image, volumes={"/data": vol}, timeout=1800)
def check(dataset: str = "", checkpoint: str = CHECKPOINT):
    vol.reload()
    print(checkpoint, "found" if os.path.isfile(checkpoint) else "MISSING")
    if dataset:
        ds = f"/data/datasets/{dataset}"
        caps, lack = _captions(ds)
        print(f"{len(caps)} captions, {len(caps) - len(lack)} cached"
              + (" (ready)" if not lack else ""))
        print(f"trainer's cache check: {len(_trainer_lack(ds))} "
              "tracks to cache")

cd ~/Ruach_Studio && .venv/bin/python - <<'PY'
import os, sys
sys.path.insert(0, "lab/trainer")
import cache
ds = "datasets/prepared/NAME"
ck = "checkpoints/comfy/checkpoints/yue2_3b_bf16.safetensors"
lack = cache.missing(ds)
print(len(lack), "tracks to cache")
if lack:
    cache.build(ds, ck, os.path.abspath("."), "cuda", lack)
PY

modal volume put ruach-data  ~/Ruach_Studio/datasets/prepared/NAME /datasets/NAME
modal run modal_ruach.py::check --dataset NAME

modal run --detach modal_ruach.py::train  --name RUN --dataset NAME

modal run --detach modal_ruach.py::train  --name RUN --dataset NAME  --overrides '{"rank": 32, "epochs": 50, "save_epochs": 5}'

modal volume ls ruach-data /runs/RUN
mkdir -p ~/Ruach_Studio/loras/RUN
T=15
for e in 010 020 030 040 050 060 070 080 090 100; do
  f=~/Ruach_Studio/loras/RUN/RUN-e${e}.safetensors
  s=$(printf "%09d" $((10#$e * T)))
  [ -f "$f" ] || modal volume get ruach-data    /runs/RUN/RUN_${s}.safetensors "$f"
done