{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "accelerator": "GPU",
  "colab": {
   "provenance": [],
   "gpuType": "L4"
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# 🎵 ACE-Step v1.5 **base / XL base** — Entraînement LoRA / LoKR avec Side-Step (Colab)\n",
    "\n",
    "[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/walidb27/ACE-Step-1.5/blob/claude/nice-volta-wvnuc6/notebooks/ACE_Step_v15_Base_SideStep_Colab.ipynb)\n",
    "\n",
    "L'interface Gradio d'ACE-Step n'entraîne que le modèle **turbo** (timesteps discrets « shift=3 »).\n",
    "Ce notebook utilise **Side-Step** (`train.py`, mode `fixed`, intégré au dépôt) qui reproduit le\n",
    "vrai schéma d'entraînement des modèles **base / sft** :\n",
    "\n",
    "- timesteps **continus** (logit-normal, paramètres lus dans le `config.json` du modèle) ;\n",
    "- **CFG dropout** (le modèle base a été entraîné avec la guidance sans classifieur) ;\n",
    "- au choix le modèle **base** (2B) ou **XL base** (4B, `MODEL_VARIANT` en section 1) : l'XL est\n",
    "  environ 2x plus lent à entraîner (~13-14 h pour 500 époques sur A100), ses tenseurs et ses LoRA\n",
    "  sont distincts de ceux du modèle base ;\n",
    "- LoRA (PEFT) ou LoKR (LyCORIS), gradient checkpointing, optimiseurs 8-bit / Prodigy.\n",
    "\n",
    "**Étapes :** 1) GPU & Drive → 2) installation → 3) téléchargement des poids → 4) dataset →\n",
    "5) prétraitement → 6) entraînement → 7) récupération de l'adaptateur.\n",
    "\n",
    "> 💡 GPU conseillé : **L4 ou A100** (bf16). Le T4 fonctionne (fp16, ~7 Go de VRAM avec gradient\n",
    "> checkpointing) mais il est plus lent et le fp16 est moins stable."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1. GPU et Google Drive\n",
    "\n",
    "Tout est rangé dans votre dossier ACE-Step du Drive, comme le ferait l'interface Gradio :\n",
    "\n",
    "```\n",
    "MyDrive/ACE-Step/\n",
    "├── checkpoints/                    poids (réutilisés, seul le manquant est téléchargé)\n",
    "├── datasets/maes2/                 vos audios (+ JSON éventuel)\n",
    "├── datasets/preprocessed_tensors/  tenseurs Side-Step (suffixe _base)\n",
    "└── lora_output/<RUN_NAME>/\n",
    "    ├── checkpoints/epoch_N/        checkpoints intermédiaires\n",
    "    └── final/                      adaptateur final\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Vérifier le GPU, monter le Drive et définir les dossiers\n",
    "ACESTEP_DRIVE_DIR = \"/content/drive/MyDrive/ACE-Step\"  #@param {type:\"string\"}\n",
    "#@markdown Votre installation ACE-Step sur le Drive.\n",
    "MODEL_VARIANT = \"xl_base\"  #@param [\"base\", \"xl_base\"]\n",
    "#@markdown `base` = modèle 2B. `xl_base` = modèle XL (4B) : plus de détail possible, ~2x plus lent\n",
    "#@markdown à entraîner, ~9 Go de poids en plus. Les tenseurs et les LoRA ne sont pas interchangeables.\n",
    "\n",
    "import os\n",
    "import subprocess\n",
    "\n",
    "print(subprocess.run([\"nvidia-smi\"], capture_output=True, text=True).stdout or \"⚠️ Aucun GPU détecté : Exécution > Modifier le type d'exécution > GPU\")\n",
    "\n",
    "from google.colab import drive\n",
    "drive.mount(\"/content/drive\")\n",
    "\n",
    "if not os.path.isdir(ACESTEP_DRIVE_DIR):\n",
    "    raise FileNotFoundError(f\"Dossier introuvable : {ACESTEP_DRIVE_DIR}\")\n",
    "WORKDIR = ACESTEP_DRIVE_DIR\n",
    "CKPT_DIR = os.path.join(WORKDIR, \"checkpoints\")\n",
    "TENSORS_ROOT = os.path.join(WORKDIR, \"datasets\", \"preprocessed_tensors\")\n",
    "OUTPUTS_ROOT = os.path.join(WORKDIR, \"lora_output\")\n",
    "for d in (CKPT_DIR, TENSORS_ROOT, OUTPUTS_ROOT):\n",
    "    os.makedirs(d, exist_ok=True)\n",
    "\n",
    "MODEL_DIR = {\"base\": \"acestep-v15-base\", \"xl_base\": \"acestep-v15-xl-base\"}[MODEL_VARIANT]\n",
    "print(\"ACE-Step (Drive) :\", WORKDIR)\n",
    "print(\"Modèle entraîné  :\", MODEL_DIR)\n",
    "print(\"Contenu de checkpoints/ :\", sorted(os.listdir(CKPT_DIR)) or \"(vide)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2. Installation (ACE-Step + Side-Step)\n",
    "\n",
    "Pip affichera probablement `ERROR: pip's dependency resolver does not currently take into account…`\n",
    "suivi de conflits sur `cudf`, `cuml`, `cuda-python`, `google-adk`, `pytensor`, `fasthtml`…\n",
    "**C'est sans conséquence** : ce sont des paquets préinstallés par Colab que ce notebook n'utilise pas.\n",
    "Ce qui compte, ce sont les dernières lignes : `torch … | CUDA True` puis `✅ Décodage audio OK`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Récupérer le code et installer les dépendances (~5 min)\n",
    "CODE_SOURCE = \"clone\"  #@param [\"clone\", \"drive\"]\n",
    "#@markdown `clone` (recommandé) : copie fraîche du code dans la VM, rapide, avec Side-Step à jour.\n",
    "#@markdown `drive` : utilise le code de `MyDrive/ACE-Step` (il doit contenir `train.py` et\n",
    "#@markdown `acestep/training_v2/`). Plus lent : Python lit le code à travers le Drive.\n",
    "REPO_URL = \"https://github.com/walidb27/ACE-Step-1.5.git\"  #@param {type:\"string\"}\n",
    "REPO_BRANCH = \"claude/nice-volta-wvnuc6\"  #@param {type:\"string\"}\n",
    "\n",
    "import os\n",
    "\n",
    "if CODE_SOURCE == \"drive\":\n",
    "    REPO_DIR = WORKDIR\n",
    "    for needed in (\"train.py\", \"acestep/training_v2/cli/train_fixed.py\", \"requirements.txt\"):\n",
    "        if not os.path.exists(os.path.join(REPO_DIR, needed)):\n",
    "            raise FileNotFoundError(\n",
    "                f\"{needed} absent de {REPO_DIR} : version trop ancienne, utilisez CODE_SOURCE = clone\"\n",
    "            )\n",
    "else:\n",
    "    REPO_DIR = \"/content/ACE-Step-1.5\"\n",
    "    if not os.path.isdir(REPO_DIR):\n",
    "        !git clone --depth 1 -b {REPO_BRANCH} {REPO_URL} {REPO_DIR}\n",
    "%cd {REPO_DIR}\n",
    "\n",
    "# requirements.txt du dépôt, sans flash-attn (compilation de plusieurs dizaines de minutes ;\n",
    "# Side-Step bascule automatiquement sur SDPA).\n",
    "!grep -v -i -E \"flash-attn|torchcodec\" requirements.txt > /content/requirements-colab.txt\n",
    "!pip install -q -r /content/requirements-colab.txt\n",
    "# torchcodec (décodage mp3/wav) doit correspondre à torch 2.10 : la dernière version\n",
    "# ne se charge pas (\"undefined symbol\") et tout le prétraitement échoue.\n",
    "!pip install -q \"torchcodec==0.10.0\"\n",
    "# Dépendances Side-Step (CLI) + optimiseurs optionnels (AdamW8bit, Prodigy)\n",
    "!pip install -q rich bitsandbytes prodigyopt\n",
    "\n",
    "# Vérification finale : doit afficher \"CUDA True\" puis \"Décodage audio OK\"\n",
    "!python -c \"import torch, torchaudio; print('torch', torch.__version__, '| CUDA', torch.cuda.is_available(), '| bf16', torch.cuda.is_available() and torch.cuda.is_bf16_supported())\"\n",
    "check = !python -c \"from torchcodec.decoders import AudioDecoder; import torchcodec; print('OK', torchcodec.__version__)\" 2>&1\n",
    "if check and check[-1].startswith(\"OK\"):\n",
    "    print(\"✅ Décodage audio OK (torchcodec\", check[-1].split()[-1] + \")\")\n",
    "else:\n",
    "    print(\"\\n\".join(check[-15:]))\n",
    "    raise RuntimeError(\"torchcodec ne se charge pas : le prétraitement échouerait. Envoyez la sortie ci-dessus.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3. Vérifier les poids du modèle\n",
    "\n",
    "Cette cellule regarde ce que contient `MyDrive/ACE-Step/checkpoints/` et ne télécharge **que ce qui\n",
    "manque**. Si ACE-Step et le modèle choisi (`MODEL_DIR`) sont déjà installés, elle ne télécharge rien.\n",
    "\n",
    "| Élément | Rôle |\n",
    "| :--- | :--- |\n",
    "| `vae/` | convertit l'audio en données pour le modèle |\n",
    "| `Qwen3-Embedding-0.6B/` | encode le texte (légendes, paroles) |\n",
    "| `silence_latent.pt` | petit fichier de référence (dans le dossier turbo) |\n",
    "| `acestep-v15-base/` ou `acestep-v15-xl-base/` | le modèle à entraîner (XL : ~9 Go) |\n",
    "\n",
    "Elle remet aussi à jour les fichiers de code `.py` du dossier du modèle (comme Gradio).\n",
    "\n",
    "**`HF_TOKEN` : laissez vide.** Ce n'est utile que pour un modèle privé ; ceux d'ACE-Step sont publics."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Vérifier / télécharger les poids manquants\n",
    "HF_TOKEN = \"\"  #@param {type:\"string\"}\n",
    "#@markdown Laisser vide (modèles publics).\n",
    "\n",
    "import os\n",
    "from huggingface_hub import snapshot_download\n",
    "\n",
    "token = HF_TOKEN or None\n",
    "WEIGHT_FILES = (\"model.safetensors\", \"model.safetensors.index.json\",\n",
    "                \"diffusion_pytorch_model.safetensors\", \"pytorch_model.bin\")\n",
    "\n",
    "\n",
    "def _has_weights(sub):\n",
    "    \"\"\"Return True when checkpoints/<sub> already holds model weights.\"\"\"\n",
    "    d = os.path.join(CKPT_DIR, sub)\n",
    "    return os.path.isdir(d) and any(os.path.exists(os.path.join(d, f)) for f in WEIGHT_FILES)\n",
    "\n",
    "\n",
    "def _has_silence_latent():\n",
    "    \"\"\"Return True when a silence_latent.pt exists in any official DiT folder.\"\"\"\n",
    "    return any(os.path.isfile(os.path.join(CKPT_DIR, d, \"silence_latent.pt\"))\n",
    "               for d in os.listdir(CKPT_DIR) if d.startswith(\"acestep-v15-\"))\n",
    "\n",
    "\n",
    "patterns = [f\"{sub}/*\" for sub in (\"vae\", \"Qwen3-Embedding-0.6B\") if not _has_weights(sub)]\n",
    "if not _has_silence_latent():\n",
    "    patterns.append(\"acestep-v15-turbo/silence_latent.pt\")\n",
    "if patterns:\n",
    "    print(\"Téléchargement depuis ACE-Step/Ace-Step1.5 :\", patterns)\n",
    "    snapshot_download(repo_id=\"ACE-Step/Ace-Step1.5\", local_dir=CKPT_DIR,\n",
    "                      allow_patterns=patterns, token=token)\n",
    "else:\n",
    "    print(\"✅ VAE, encodeur texte et silence_latent déjà présents\")\n",
    "\n",
    "if _has_weights(MODEL_DIR):\n",
    "    print(f\"✅ {MODEL_DIR} déjà présent :\", os.path.join(CKPT_DIR, MODEL_DIR))\n",
    "    # Side-Step charge le code Python rangé dans le dossier du modèle (trust_remote_code).\n",
    "    # On l'aligne sur le code du dépôt, comme le fait Gradio à chaque démarrage.\n",
    "    import sys\n",
    "    sys.path.insert(0, REPO_DIR)\n",
    "    from acestep.model_downloader import _check_code_mismatch, _sync_model_code_files\n",
    "    if _check_code_mismatch(MODEL_DIR, CKPT_DIR):\n",
    "        print(\"   Code du modèle mis à jour :\", _sync_model_code_files(MODEL_DIR, CKPT_DIR))\n",
    "else:\n",
    "    print(f\"Téléchargement de {MODEL_DIR}…\")\n",
    "    cmd = f\"python -m acestep.model_downloader --model {MODEL_DIR} --skip-main --dir {CKPT_DIR}\"\n",
    "    if token:\n",
    "        cmd += f\" --token {token}\"\n",
    "    !{cmd}\n",
    "\n",
    "!ls {CKPT_DIR} {CKPT_DIR}/{MODEL_DIR}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4. Dataset\n",
    "\n",
    "Placez vos fichiers audio (`.wav .mp3 .flac .ogg .opus .m4a`, idéalement 30 s à 4 min, **5 morceaux\n",
    "minimum**) dans `AUDIO_DIR` (sur le Drive) ou envoyez un `.zip` avec la cellule ci-dessous.\n",
    "\n",
    "Deux façons de fournir les métadonnées :\n",
    "- **Construire le JSON ici** (cellule 4b) : mot déclencheur (`custom_tag`) en\n",
    "  **prepend / append / replace**, `genre_ratio`, légendes et paroles par fichier ;\n",
    "- **Utiliser un JSON existant**, par exemple celui exporté par l'onglet *Dataset* de Gradio\n",
    "  (qui sait aussi légender automatiquement avec le LM). La cellule 4b peut alors juste\n",
    "  réappliquer le tag et sa position.\n",
    "\n",
    "Voir `docs/sidestep/Dataset Preparation.md`."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Choisir le dossier audio (et éventuellement le JSON)\n",
    "DATASET_NAME = \"maes2\"  #@param {type:\"string\"}\n",
    "#@markdown Sous-dossier de `MyDrive/ACE-Step/datasets/`.\n",
    "DATASET_JSON = \"\"  #@param {type:\"string\"}\n",
    "#@markdown Optionnel. Vide = détection automatique d'un JSON ACE-Step (`datasets/maes2.json`\n",
    "#@markdown ou un `.json` avec une liste `samples` dans le dossier).\n",
    "UPLOAD_ZIP = False  #@param {type:\"boolean\"}\n",
    "#@markdown Cocher pour envoyer un .zip d'audios depuis votre ordinateur.\n",
    "\n",
    "import json\n",
    "import os\n",
    "import zipfile\n",
    "\n",
    "AUDIO_DIR = os.path.join(WORKDIR, \"datasets\", DATASET_NAME)\n",
    "os.makedirs(AUDIO_DIR, exist_ok=True)\n",
    "\n",
    "\n",
    "def _is_dataset_json(path):\n",
    "    # True for an ACE-Step dataset JSON (a list, or a dict with \"samples\").\n",
    "    try:\n",
    "        with open(path, encoding=\"utf-8\") as fh:\n",
    "            data = json.load(fh)\n",
    "    except (OSError, ValueError):\n",
    "        return False\n",
    "    return isinstance(data, list) or (isinstance(data, dict) and \"samples\" in data)\n",
    "\n",
    "\n",
    "if not DATASET_JSON:\n",
    "    candidates = [os.path.join(WORKDIR, \"datasets\", f\"{DATASET_NAME}.json\")]\n",
    "    candidates += sorted(\n",
    "        os.path.join(AUDIO_DIR, f) for f in os.listdir(AUDIO_DIR)\n",
    "        if f.endswith(\".json\") and not f.endswith(\".sidestep.json\")\n",
    "    )\n",
    "    found_json = [c for c in candidates if os.path.isfile(c) and _is_dataset_json(c)]\n",
    "    if found_json:\n",
    "        DATASET_JSON = found_json[0]\n",
    "        print(\"JSON détecté :\", DATASET_JSON)\n",
    "\n",
    "if UPLOAD_ZIP:\n",
    "    from google.colab import files\n",
    "    for name in files.upload():\n",
    "        if name.lower().endswith(\".zip\"):\n",
    "            with zipfile.ZipFile(name) as zf:\n",
    "                zf.extractall(AUDIO_DIR)\n",
    "            os.remove(name)\n",
    "\n",
    "EXTS = (\".wav\", \".mp3\", \".flac\", \".ogg\", \".opus\", \".m4a\")\n",
    "found = [os.path.join(r, f) for r, _, fs in os.walk(AUDIO_DIR) for f in fs if f.lower().endswith(EXTS)]\n",
    "print(f\"{len(found)} fichier(s) audio dans {AUDIO_DIR}\")\n",
    "for f in found[:20]:\n",
    "    print(\" -\", os.path.relpath(f, AUDIO_DIR))\n",
    "if DATASET_JSON and not os.path.isfile(DATASET_JSON):\n",
    "    raise FileNotFoundError(f\"JSON introuvable : {DATASET_JSON}\")\n",
    "if not found and not DATASET_JSON:\n",
    "    print(\"⚠️ Aucun audio trouvé : ajoutez des fichiers avant le prétraitement.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4b. Construire le JSON du dataset (tag en prepend / append / replace)\n",
    "\n",
    "Pour chaque audio `chanson.mp3`, la cellule lit, s'ils existent, des fichiers posés à côté :\n",
    "\n",
    "| Fichier | Contenu |\n",
    "| :--- | :--- |\n",
    "| `chanson.caption.txt` | légende (description du style, des instruments…) |\n",
    "| `chanson.lyrics.txt` (ou `chanson.txt`) | paroles avec `[Verse]`, `[Chorus]`… |\n",
    "| `chanson.json` | champs libres : `genre`, `bpm`, `keyscale`, `timesignature`, `language`, `caption`, `lyrics`, `prompt_override`… |\n",
    "\n",
    "Sans légende, `DEFAULT_CAPTION` est utilisée (ou le nom du fichier si elle est vide).\n",
    "Sans paroles : `[Instrumental]`. La durée est mesurée automatiquement.\n",
    "\n",
    "**Position du tag** (ex. tag `mystyle`, légende `dark synthwave`) :\n",
    "`prepend` → `mystyle, dark synthwave` · `append` → `dark synthwave, mystyle` ·\n",
    "`replace` → `mystyle` (la légende est ignorée).\n",
    "\n",
    "À l'inférence, mettez le tag dans le prompt pour activer le style. Toute modification ici\n",
    "relance automatiquement le prétraitement à l'étape 5."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Générer le JSON du dataset\n",
    "BUILD_JSON = True  #@param {type:\"boolean\"}\n",
    "#@markdown Décocher pour utiliser tel quel le `DATASET_JSON` choisi plus haut.\n",
    "CUSTOM_TAG = \"maes2\"  #@param {type:\"string\"}\n",
    "#@markdown Mot déclencheur, ex. `mystyle` (vide = pas de tag).\n",
    "TAG_POSITION = \"replace\"  #@param [\"prepend\", \"append\", \"replace\"]\n",
    "GENRE_RATIO = 0  #@param {type:\"slider\", min:0, max:100, step:5}\n",
    "#@markdown % des morceaux qui utilisent le champ `genre` au lieu de la légende (s'il est renseigné).\n",
    "DEFAULT_CAPTION = \"\"  #@param {type:\"string\"}\n",
    "DEFAULT_GENRE = \"\"  #@param {type:\"string\"}\n",
    "ALL_INSTRUMENTAL = False  #@param {type:\"boolean\"}\n",
    "#@markdown Cocher pour forcer `[Instrumental]` partout (paroles ignorées).\n",
    "\n",
    "import hashlib\n",
    "import json\n",
    "import os\n",
    "import subprocess\n",
    "import sys\n",
    "\n",
    "import soundfile as sf\n",
    "\n",
    "sys.path.insert(0, REPO_DIR)\n",
    "from acestep.training_v2.preprocess_prompt import build_simple_prompt\n",
    "\n",
    "\n",
    "def _read_text(path):\n",
    "    \"\"\"Return stripped file content, or \"\" when the file is missing.\"\"\"\n",
    "    if os.path.isfile(path):\n",
    "        with open(path, encoding=\"utf-8\") as fh:\n",
    "            return fh.read().strip()\n",
    "    return \"\"\n",
    "\n",
    "\n",
    "def _duration(path):\n",
    "    \"\"\"Measure audio duration in seconds (soundfile, then ffprobe).\"\"\"\n",
    "    try:\n",
    "        return int(sf.info(path).duration)\n",
    "    except Exception:\n",
    "        out = subprocess.run(\n",
    "            [\"ffprobe\", \"-v\", \"error\", \"-show_entries\", \"format=duration\",\n",
    "             \"-of\", \"default=nw=1:nk=1\", path],\n",
    "            capture_output=True, text=True,\n",
    "        ).stdout.strip()\n",
    "        return int(float(out)) if out else 0\n",
    "\n",
    "\n",
    "def _sample_from_files(audio_path):\n",
    "    \"\"\"Build one sample dict from an audio file and its sidecar files.\"\"\"\n",
    "    base = os.path.splitext(audio_path)[0]\n",
    "    extra = {}\n",
    "    if os.path.isfile(base + \".json\"):\n",
    "        with open(base + \".json\", encoding=\"utf-8\") as fh:\n",
    "            extra = json.load(fh)\n",
    "    stem = os.path.basename(base)\n",
    "    caption = (_read_text(base + \".caption.txt\") or extra.get(\"caption\")\n",
    "               or DEFAULT_CAPTION or stem.replace(\"_\", \" \").replace(\"-\", \" \"))\n",
    "    lyrics = (_read_text(base + \".lyrics.txt\") or _read_text(base + \".txt\")\n",
    "              or extra.get(\"lyrics\") or \"[Instrumental]\")\n",
    "    sample = {\n",
    "        # Stable id: an unchanged dataset must give an identical JSON (no re-preprocessing).\n",
    "        \"id\": hashlib.sha1(os.path.basename(audio_path).encode()).hexdigest()[:8],\n",
    "        \"audio_path\": os.path.relpath(audio_path, AUDIO_DIR),\n",
    "        \"filename\": os.path.basename(audio_path),\n",
    "        \"caption\": caption,\n",
    "        \"genre\": extra.get(\"genre\") or DEFAULT_GENRE,\n",
    "        \"lyrics\": lyrics,\n",
    "        \"bpm\": extra.get(\"bpm\"),\n",
    "        \"keyscale\": extra.get(\"keyscale\", \"\"),\n",
    "        \"timesignature\": extra.get(\"timesignature\", \"\"),\n",
    "        \"duration\": extra.get(\"duration\") or _duration(audio_path),\n",
    "        \"language\": extra.get(\"language\", \"unknown\"),\n",
    "        \"prompt_override\": extra.get(\"prompt_override\"),\n",
    "    }\n",
    "    return sample\n",
    "\n",
    "\n",
    "def _finalize(sample):\n",
    "    \"\"\"Apply instrumental flag and trigger tag to a sample (in place).\"\"\"\n",
    "    if ALL_INSTRUMENTAL:\n",
    "        sample[\"lyrics\"] = \"[Instrumental]\"\n",
    "    sample[\"is_instrumental\"] = sample.get(\"lyrics\", \"\").strip() == \"[Instrumental]\"\n",
    "    if CUSTOM_TAG:\n",
    "        # Per-sample tags win over the dataset-level one, so set both.\n",
    "        sample[\"custom_tag\"] = CUSTOM_TAG\n",
    "    sample[\"labeled\"] = True\n",
    "    return sample\n",
    "\n",
    "\n",
    "if BUILD_JSON:\n",
    "    if DATASET_JSON:\n",
    "        # Re-tag an existing JSON (e.g. exported by the Gradio Dataset Builder)\n",
    "        with open(DATASET_JSON, encoding=\"utf-8\") as fh:\n",
    "            src = json.load(fh)\n",
    "        samples = src if isinstance(src, list) else src.get(\"samples\", [])\n",
    "        src_dir = os.path.dirname(os.path.abspath(DATASET_JSON))\n",
    "        by_name = {os.path.basename(f): f for f in found}\n",
    "        missing = []\n",
    "        for s in samples:\n",
    "            ap = (s.get(\"audio_path\") or s.get(\"filename\", \"\")).replace(\"\\\\\", \"/\")\n",
    "            full = ap if os.path.isabs(ap) else os.path.join(src_dir, ap)\n",
    "            if not os.path.isfile(full):\n",
    "                # JSON built on another machine (e.g. C:/...): match by file name\n",
    "                full = by_name.get(os.path.basename(ap), \"\")\n",
    "            if full:\n",
    "                s[\"audio_path\"] = os.path.relpath(full, AUDIO_DIR)\n",
    "                s[\"filename\"] = os.path.basename(full)\n",
    "            else:\n",
    "                missing.append(ap)\n",
    "        if missing:\n",
    "            print(f\"⚠️ {len(missing)} audio(s) du JSON introuvable(s) dans {AUDIO_DIR} :\", missing[:5])\n",
    "        samples = [s for s in samples if s.get(\"filename\") in by_name]\n",
    "    else:\n",
    "        samples = [_sample_from_files(p) for p in sorted(found)]\n",
    "    samples = [_finalize(s) for s in samples]\n",
    "\n",
    "    dataset = {\n",
    "        \"metadata\": {\n",
    "            \"name\": DATASET_NAME,\n",
    "            \"custom_tag\": CUSTOM_TAG,\n",
    "            \"tag_position\": TAG_POSITION,\n",
    "            \"genre_ratio\": int(GENRE_RATIO),\n",
    "            \"all_instrumental\": all(s[\"is_instrumental\"] for s in samples),\n",
    "            \"num_samples\": len(samples),\n",
    "        },\n",
    "        \"samples\": samples,\n",
    "    }\n",
    "    out_json = os.path.join(AUDIO_DIR, f\"{DATASET_NAME}.sidestep.json\")\n",
    "    with open(out_json, \"w\", encoding=\"utf-8\") as fh:\n",
    "        json.dump(dataset, fh, ensure_ascii=False, indent=2)\n",
    "    DATASET_JSON = out_json\n",
    "    print(f\"✅ {len(samples)} morceau(x) -> {out_json}\")\n",
    "    if GENRE_RATIO and not any(s.get(\"genre\") for s in samples):\n",
    "        print(\"⚠️ GENRE_RATIO > 0 mais aucun genre renseigné : la légende sera utilisée.\")\n",
    "\n",
    "    # Preview of the exact text prompts Side-Step will encode\n",
    "    for s in samples[:3]:\n",
    "        caption_block = build_simple_prompt(s, tag_position=TAG_POSITION).split(\"# Caption\")[1]\n",
    "        print(f\"\\n--- {s['filename']} ---{caption_block.split('# Metas')[0].rstrip()}\")\n",
    "else:\n",
    "    print(\"JSON utilisé tel quel :\", DATASET_JSON or \"(aucun : légendes = noms de fichiers)\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5. Prétraitement (audio → tenseurs `.pt`)\n",
    "\n",
    "Deux passes à faible VRAM : VAE + encodeur texte, puis encodeur de conditionnement du modèle choisi.\n",
    "Les tenseurs doivent être produits **avec le même modèle que celui qu'on entraîne** (ici `base`).\n",
    "À ne refaire que si le dataset change."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Lancer le prétraitement\n",
    "MAX_DURATION = 240  #@param {type:\"number\"}\n",
    "FORCE_REPROCESS = False  #@param {type:\"boolean\"}\n",
    "\n",
    "%cd {REPO_DIR}\n",
    "import hashlib\n",
    "import os\n",
    "import shlex\n",
    "import shutil\n",
    "\n",
    "# Un dossier par modèle : des tenseurs faits pour base ne conviennent pas à XL (ni à turbo).\n",
    "TENSOR_DIR = os.path.join(TENSORS_ROOT, f\"{DATASET_NAME}_{MODEL_VARIANT}\")\n",
    "STAMP = os.path.join(TENSOR_DIR, \"_notebook_source.sha1\")\n",
    "\n",
    "# Fingerprint of what the tensors are built from: JSON content + audio list + max duration.\n",
    "# Side-Step skips existing .pt files, so a changed tag/caption needs a clean folder.\n",
    "h = hashlib.sha1(f\"{AUDIO_DIR}|{MAX_DURATION}|{sorted(found)}\".encode())\n",
    "if DATASET_JSON:\n",
    "    with open(DATASET_JSON, \"rb\") as fh:\n",
    "        h.update(fh.read())\n",
    "fingerprint = h.hexdigest()\n",
    "previous = open(STAMP).read().strip() if os.path.isfile(STAMP) else None\n",
    "\n",
    "\n",
    "def _tensor_count():\n",
    "    # Finished .pt tensors only (Side-Step also writes .tmp.pt intermediates).\n",
    "    if not os.path.isdir(TENSOR_DIR):\n",
    "        return 0\n",
    "    return sum(f.endswith(\".pt\") and not f.endswith(\".tmp.pt\") for f in os.listdir(TENSOR_DIR))\n",
    "\n",
    "\n",
    "up_to_date = _tensor_count() > 0 and previous == fingerprint\n",
    "\n",
    "if os.path.isdir(TENSOR_DIR) and (FORCE_REPROCESS or not up_to_date):\n",
    "    if previous and previous != fingerprint:\n",
    "        print(\"Dataset modifié (tag, légendes, fichiers…) : nouveau prétraitement.\")\n",
    "    shutil.rmtree(TENSOR_DIR)\n",
    "    up_to_date = False\n",
    "\n",
    "if up_to_date:\n",
    "    print(\"Tenseurs à jour :\", TENSOR_DIR, \"(cocher FORCE_REPROCESS pour refaire)\")\n",
    "else:\n",
    "    args = [\n",
    "        \"python\", \"-m\", \"acestep.training_v2.cli.train_fixed\", \"--plain\",\n",
    "        \"--checkpoint-dir\", CKPT_DIR,\n",
    "        \"--model-variant\", MODEL_VARIANT,\n",
    "        \"--preprocess\",\n",
    "        \"--audio-dir\", AUDIO_DIR,\n",
    "        \"--tensor-output\", TENSOR_DIR,\n",
    "        \"--max-duration\", str(MAX_DURATION),\n",
    "    ]\n",
    "    if DATASET_JSON:\n",
    "        args += [\"--dataset-json\", DATASET_JSON]\n",
    "    cmd = shlex.join(args)\n",
    "    print(cmd)\n",
    "    !{cmd}\n",
    "    if _tensor_count() > 0:\n",
    "        with open(STAMP, \"w\") as fh:\n",
    "            fh.write(fingerprint)\n",
    "\n",
    "print(f\"{_tensor_count()} tenseur(s) dans {TENSOR_DIR}\")"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title 🔍 Vérifier le dataset prétraité (légendes, tag, paroles)\n",
    "#@markdown Relit les tenseurs réellement utilisés par l'entraînement et affiche, pour chaque morceau,\n",
    "#@markdown le texte exact vu par le modèle. À lancer après le prétraitement, avant l'entraînement.\n",
    "SHOW_ALL = False  #@param {type:\"boolean\"}\n",
    "#@markdown Cocher pour afficher les 84 morceaux (sinon les 5 premiers + les anomalies).\n",
    "\n",
    "import collections\n",
    "import glob\n",
    "import json\n",
    "import os\n",
    "import sys\n",
    "\n",
    "import torch\n",
    "\n",
    "sys.path.insert(0, REPO_DIR)\n",
    "from acestep.training_v2.preprocess_prompt import build_simple_prompt\n",
    "\n",
    "tag_position = \"prepend\"\n",
    "if DATASET_JSON and os.path.isfile(DATASET_JSON):\n",
    "    with open(DATASET_JSON, encoding=\"utf-8\") as fh:\n",
    "        _ds = json.load(fh)\n",
    "    if isinstance(_ds, dict):\n",
    "        tag_position = _ds.get(\"metadata\", {}).get(\"tag_position\", \"prepend\")\n",
    "\n",
    "pts = sorted(p for p in glob.glob(os.path.join(TENSOR_DIR, \"*.pt\")) if not p.endswith(\".tmp.pt\"))\n",
    "print(f\"Dossier : {TENSOR_DIR}\")\n",
    "print(f\"JSON    : {DATASET_JSON or '(aucun)'}  |  tag_position = {tag_position}\")\n",
    "print(f\"{len(pts)} tenseur(s) trouvés\\n\")\n",
    "\n",
    "rows, problems = [], []\n",
    "for p in pts:\n",
    "    meta = torch.load(p, map_location=\"cpu\", weights_only=False, mmap=True)[\"metadata\"]\n",
    "    prompt = build_simple_prompt(meta, tag_position=tag_position)\n",
    "    caption_line = prompt.split(\"# Caption\")[1].split(\"# Metas\")[0].strip()\n",
    "    lyrics = (meta.get(\"lyrics\") or \"\").strip()\n",
    "    rows.append((meta.get(\"filename\"), caption_line, meta.get(\"custom_tag\", \"\"), lyrics))\n",
    "    if not meta.get(\"custom_tag\"):\n",
    "        problems.append(f\"{meta.get('filename')} : pas de tag\")\n",
    "    if lyrics == \"[Instrumental]\":\n",
    "        problems.append(f\"{meta.get('filename')} : pas de paroles ([Instrumental])\")\n",
    "\n",
    "for name, caption_line, tag, lyrics in (rows if SHOW_ALL else rows[:5]):\n",
    "    first = lyrics.splitlines()[:2]\n",
    "    print(f\"• {name}\\n    texte vu : {caption_line!r}\\n    paroles  : {' / '.join(first)[:90]}\")\n",
    "\n",
    "print(\"\\n=== Résumé ===\")\n",
    "print(\"Textes distincts vus par le modèle :\", dict(collections.Counter(r[1] for r in rows).most_common(5)))\n",
    "print(\"Tags :\", dict(collections.Counter(r[2] for r in rows)))\n",
    "print(\"Morceaux avec paroles :\", sum(r[3] != \"[Instrumental]\" for r in rows), \"/\", len(rows))\n",
    "if problems:\n",
    "    print(f\"\\n⚠️ {len(problems)} point(s) à vérifier :\")\n",
    "    for line in problems[:20]:\n",
    "        print(\"   -\", line)\n",
    "else:\n",
    "    print(\"\\n✅ Tous les morceaux ont un tag et des paroles.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6. Entraînement Side-Step (modèle base ou XL)\n",
    "\n",
    "Valeurs par défaut = preset `recommended` de Side-Step. `shift=1.0` et `50` steps sont enregistrés\n",
    "comme métadonnées d'inférence (valeurs correctes pour base) ; ils n'influencent pas l'entraînement.\n",
    "\n",
    "Repères :\n",
    "- **Nombre d'époques** : pour reproduire fidèlement une voix, compter **~500 époques pour ~100 morceaux**\n",
    "  (tutoriel officiel `docs/en/LoRA_Training_Tutorial.md`, confirmé en pratique). Le guide Side-Step\n",
    "  (50-100 époques au-delà de 50 morceaux) sous-estime pour cet usage. Avec `cosine`, une reprise\n",
    "  après coupure (`RESUME_FROM`) marche si `EPOCHS` reste identique ; pour prolonger au-delà, utiliser `constant`.\n",
    "- **VRAM juste ?** `OFFLOAD_ENCODER`, `OPTIMIZER = adamw8bit`, rang plus bas.\n",
    "- **Sur-apprentissage** (la sortie recopie le dataset) : moins d'époques, rang plus bas, plus de données.\n",
    "- **Pas d'effet audible** : plus d'époques ou rang plus élevé."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Paramètres d'entraînement\n",
    "RUN_NAME = \"maes2_xl_r256\"  #@param {type:\"string\"}\n",
    "#@markdown Un nom par entraînement (protégé contre l'écrasement).\n",
    "ADAPTER_TYPE = \"lora\"  #@param [\"lora\", \"lokr\"]\n",
    "RANK = 256  #@param {type:\"integer\"}\n",
    "ALPHA = 512  #@param {type:\"integer\"}\n",
    "DROPOUT = 0.1  #@param {type:\"number\"}\n",
    "ATTENTION_TYPE = \"both\"  #@param [\"both\", \"self\", \"cross\"]\n",
    "TARGET_LAYERS = \"attention + mlp\"  #@param [\"attention\", \"attention + mlp\"]\n",
    "#@markdown `attention + mlp` adapte aussi les couches feed-forward du DiT : plus de capacité\n",
    "#@markdown (timbre, détails), adaptateur ~2-3x plus gros, un peu plus de VRAM. Expérimental.\n",
    "LEARNING_RATE = 2e-4  #@param {type:\"number\"}\n",
    "EPOCHS = 500  #@param {type:\"integer\"}\n",
    "BATCH_SIZE = 2  #@param {type:\"integer\"}\n",
    "GRADIENT_ACCUMULATION = 2  #@param {type:\"integer\"}\n",
    "WARMUP_STEPS = 100  #@param {type:\"integer\"}\n",
    "OPTIMIZER = \"adamw\"  #@param [\"adamw\", \"adamw8bit\", \"adafactor\", \"prodigy\"]\n",
    "SCHEDULER = \"cosine\"  #@param [\"cosine\", \"cosine_restarts\", \"linear\", \"constant\", \"constant_with_warmup\"]\n",
    "CFG_RATIO = 0.15  #@param {type:\"number\"}\n",
    "TIMESTEP_MU = \"auto\"  #@param [\"auto\", \"-0.6\", \"-0.8\", \"-1.0\"] {allow-input: true}\n",
    "#@markdown Répartition des niveaux de bruit vus à l'entraînement. `auto` = valeur du modèle (-0.4).\n",
    "#@markdown Plus bas (-0.8) = plus d'entraînement sur les étapes peu bruitées, celles qui\n",
    "#@markdown portent les **détails fins** (grain de voix, texture). Expérimental.\n",
    "SAVE_EVERY = 25  #@param {type:\"integer\"}\n",
    "PRECISION = \"auto\"  #@param [\"auto\", \"bf16\", \"fp16\", \"fp32\"]\n",
    "#@markdown `auto` = bf16 sur A100 / L4 / H100, fp16 sur T4.\n",
    "OFFLOAD_ENCODER = False  #@param {type:\"boolean\"}\n",
    "SEED = 42  #@param {type:\"integer\"}\n",
    "#@markdown ---\n",
    "#@markdown **Reprise** : chemin d'un checkpoint, ex.\n",
    "#@markdown `/content/drive/MyDrive/ACE-Step/lora_output/maes2_base/checkpoints/epoch_50_loss_0.1234`\n",
    "#@markdown (la section 7 liste les checkpoints disponibles).\n",
    "RESUME_FROM = \"\"  #@param {type:\"string\"}\n",
    "OVERWRITE_RUN = False  #@param {type:\"boolean\"}\n",
    "#@markdown Cocher seulement pour réécrire volontairement un RUN_NAME déjà utilisé.\n",
    "\n",
    "import os\n",
    "import shlex\n",
    "import subprocess\n",
    "\n",
    "# Side-Step \"auto\" choisit toujours bf16 sur CUDA, que le T4 (ou plus ancien) ne gère pas\n",
    "# nativement : on bascule alors sur fp16.\n",
    "if PRECISION == \"auto\":\n",
    "    bf16_ok = subprocess.run(\n",
    "        [\"python\", \"-c\", \"import torch; print(torch.cuda.is_available() and torch.cuda.is_bf16_supported(including_emulation=False))\"],\n",
    "        capture_output=True, text=True,\n",
    "    ).stdout.strip() == \"True\"\n",
    "    PRECISION = \"bf16\" if bf16_ok else \"fp16\"\n",
    "    print(\"Précision choisie automatiquement :\", PRECISION)\n",
    "\n",
    "# Un espace collé par erreur dans un champ créerait un autre dossier.\n",
    "RUN_NAME = RUN_NAME.strip()\n",
    "RESUME_FROM = RESUME_FROM.strip()\n",
    "if RESUME_FROM and not os.path.isdir(RESUME_FROM):\n",
    "    raise FileNotFoundError(f\"RESUME_FROM introuvable : {RESUME_FROM!r}\")\n",
    "OUTPUT_DIR = os.path.join(OUTPUTS_ROOT, RUN_NAME)\n",
    "# Reprise sans relancer les sections 4/4b/5 : retrouver les tenseurs déjà prétraités.\n",
    "if \"TENSOR_DIR\" not in globals():\n",
    "    TENSOR_DIR = os.path.join(TENSORS_ROOT, f\"{globals().get('DATASET_NAME', 'maes2')}_{MODEL_VARIANT}\")\n",
    "    print(\"TENSOR_DIR non défini (prétraitement sauté) -> utilisation de\", TENSOR_DIR)\n",
    "if not os.path.isdir(TENSOR_DIR) or not any(f.endswith(\".pt\") for f in os.listdir(TENSOR_DIR)):\n",
    "    raise FileNotFoundError(f\"Aucun tenseur dans {TENSOR_DIR} : lancez le prétraitement (section 5).\")\n",
    "if MODEL_VARIANT not in os.path.basename(TENSOR_DIR):\n",
    "    raise RuntimeError(\n",
    "        f\"TENSOR_DIR ({TENSOR_DIR}) ne correspond pas au modèle {MODEL_VARIANT} : \"\n",
    "        \"relancez le prétraitement avec ce modèle (les tenseurs ne sont pas interchangeables).\"\n",
    "    )\n",
    "print(\"Modèle entraîné :\", MODEL_VARIANT, \"| tenseurs :\", TENSOR_DIR)\n",
    "# Ne jamais mélanger deux entraînements dans le même dossier (le nouveau final/ écraserait l'ancien).\n",
    "_previous_run = [d for d in (\"final\", \"checkpoints\") if os.path.isdir(os.path.join(OUTPUT_DIR, d))]\n",
    "if _previous_run and not RESUME_FROM and not OVERWRITE_RUN:\n",
    "    raise RuntimeError(\n",
    "        f\"{OUTPUT_DIR} contient déjà un entraînement ({', '.join(_previous_run)}). \"\n",
    "        \"Changez RUN_NAME, ou renseignez RESUME_FROM pour le reprendre, \"\n",
    "        \"ou cochez OVERWRITE_RUN pour écrire par-dessus.\"\n",
    "    )\n",
    "print(\"Checkpoints et adaptateur final ->\", OUTPUT_DIR)\n",
    "if OPTIMIZER == \"prodigy\" and LEARNING_RATE < 0.1:\n",
    "    print(\"ℹ️ Prodigy ajuste le LR lui-même : un LR de 1.0 est conseillé.\")\n",
    "\n",
    "train_args = [\n",
    "    \"python\", os.path.join(REPO_DIR, \"train.py\"), \"--plain\", \"--yes\", \"fixed\",\n",
    "    \"--checkpoint-dir\", CKPT_DIR,\n",
    "    \"--model-variant\", MODEL_VARIANT,\n",
    "    \"--dataset-dir\", TENSOR_DIR,\n",
    "    \"--output-dir\", OUTPUT_DIR,\n",
    "    \"--adapter-type\", ADAPTER_TYPE,\n",
    "    \"--attention-type\", ATTENTION_TYPE,\n",
    "    \"--lr\", str(LEARNING_RATE),\n",
    "    \"--epochs\", str(EPOCHS),\n",
    "    \"--batch-size\", str(BATCH_SIZE),\n",
    "    \"--gradient-accumulation\", str(GRADIENT_ACCUMULATION),\n",
    "    \"--warmup-steps\", str(WARMUP_STEPS),\n",
    "    \"--optimizer-type\", OPTIMIZER,\n",
    "    \"--scheduler-type\", SCHEDULER,\n",
    "    \"--cfg-ratio\", str(CFG_RATIO),\n",
    "    \"--save-every\", str(SAVE_EVERY),\n",
    "    \"--precision\", PRECISION,\n",
    "    \"--seed\", str(SEED),\n",
    "    \"--shift\", \"1.0\",\n",
    "    \"--num-inference-steps\", \"50\",\n",
    "]\n",
    "modules = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\"]\n",
    "if TARGET_LAYERS == \"attention + mlp\":\n",
    "    modules += [\"gate_proj\", \"up_proj\", \"down_proj\"]\n",
    "train_args += [\"--target-modules\", *modules]\n",
    "if ADAPTER_TYPE == \"lora\":\n",
    "    train_args += [\"--rank\", str(RANK), \"--alpha\", str(ALPHA), \"--dropout\", str(DROPOUT)]\n",
    "else:\n",
    "    train_args += [\"--lokr-linear-dim\", str(RANK), \"--lokr-linear-alpha\", str(ALPHA)]\n",
    "if str(TIMESTEP_MU).strip() not in (\"\", \"auto\"):\n",
    "    train_args += [\"--timestep-mu\", str(float(TIMESTEP_MU))]\n",
    "if OFFLOAD_ENCODER:\n",
    "    train_args.append(\"--offload-encoder\")\n",
    "if RESUME_FROM:\n",
    "    train_args += [\"--resume-from\", RESUME_FROM]\n",
    "\n",
    "TRAIN_CMD = shlex.join(train_args)\n",
    "print(TRAIN_CMD)"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title (Optionnel) TensorBoard — lancer avant l'entraînement pour suivre la loss\n",
    "%load_ext tensorboard\n",
    "%tensorboard --logdir \"{OUTPUT_DIR}/runs\""
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title 🚀 Lancer l'entraînement\n",
    "# Side-Step n'accepte que des chemins situés sous le dossier courant (protection\n",
    "# \"safe root\"). On lance donc depuis /content, qui contient le code ET le Drive.\n",
    "%cd /content\n",
    "!{TRAIN_CMD}"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 7. Adaptateur et checkpoints sur le Drive\n",
    "\n",
    "Tout est écrit directement sur le Drive pendant l'entraînement, rien n'est à copier :\n",
    "\n",
    "- checkpoints intermédiaires (tous les `SAVE_EVERY` époques, utilisables en inférence) :\n",
    "  `MyDrive/ACE-Step/lora_output/<RUN_NAME>/checkpoints/epoch_N_loss_X/`\n",
    "- adaptateur final : `MyDrive/ACE-Step/lora_output/<RUN_NAME>/final/`\n",
    "- courbes TensorBoard : `MyDrive/ACE-Step/lora_output/<RUN_NAME>/runs/`\n",
    "\n",
    "Si la session Colab coupe, relancez les sections 1 à 3, puis l'entraînement avec\n",
    "`RESUME_FROM` = le dernier checkpoint listé ci-dessous."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Lister les checkpoints sauvegardés et forcer l'écriture sur le Drive\n",
    "import os\n",
    "\n",
    "ckpt_root = os.path.join(OUTPUT_DIR, \"checkpoints\")\n",
    "if os.path.isdir(ckpt_root):\n",
    "    ckpts = sorted(os.listdir(ckpt_root), key=lambda n: int(n.split(\"_\")[1]) if n.split(\"_\")[1].isdigit() else 0)\n",
    "    for name in ckpts:\n",
    "        print(\"  \", os.path.join(ckpt_root, name))\n",
    "    if ckpts:\n",
    "        print(\"\\nDernier checkpoint (pour RESUME_FROM) :\", os.path.join(ckpt_root, ckpts[-1]))\n",
    "else:\n",
    "    print(\"Aucun checkpoint pour l'instant dans\", ckpt_root)\n",
    "final_dir = os.path.join(OUTPUT_DIR, \"final\")\n",
    "print(\"Adaptateur final :\", final_dir if os.path.isdir(final_dir) else \"(pas encore créé)\")\n",
    "\n",
    "# Le Drive se synchronise en arrière-plan : os.sync() pousse les écritures en attente.\n",
    "os.sync()\n",
    "print(\"✅ Écritures envoyées au Drive\")"
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title (Optionnel) Télécharger l'adaptateur final en .zip sur votre ordinateur\n",
    "import os\n",
    "import shutil\n",
    "\n",
    "from google.colab import files\n",
    "\n",
    "final_dir = os.path.join(OUTPUT_DIR, \"final\")\n",
    "if not os.path.isdir(final_dir):\n",
    "    print(\"Pas encore de dossier final :\", final_dir)\n",
    "else:\n",
    "    archive = shutil.make_archive(os.path.join(\"/content\", f\"{RUN_NAME}_final\"), \"zip\", final_dir)\n",
    "    files.download(archive)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 8. Tester l'adaptateur dans l'interface Gradio d'ACE-Step\n",
    "\n",
    "L'interface Gradio sert ici seulement à **générer** (pas à entraîner) : elle charge les\n",
    "adaptateurs Side-Step tels quels (LoRA PEFT ou LoKR `lokr_weights.safetensors`).\n",
    "\n",
    "1. Lancez la cellule ci-dessous. Elle démarre Gradio avec le modèle entraîné (`MODEL_DIR`) et affiche un lien\n",
    "   public `https://….gradio.live`.\n",
    "   Au premier lancement, ACE-Step télécharge les composants manquants (LM 5Hz, DiT turbo, ~6 Go).\n",
    "2. Dans l'interface : *LoRA Adapter* → collez le chemin affiché → *Load LoRA* → cochez *Use LoRA*,\n",
    "   réglez *LoRA Scale* (1.0 = pleine force).\n",
    "3. Générez avec **shift = 1.0** et **~50 steps** (réglages base / XL base), en mettant votre\n",
    "   `CUSTOM_TAG` dans le prompt.\n",
    "\n",
    "La cellule tourne tant que l'interface est ouverte : l'interrompre (■) arrête Gradio."
   ]
  },
  {
   "cell_type": "code",
   "metadata": {
    "cellView": "form"
   },
   "execution_count": null,
   "outputs": [],
   "source": [
    "#@title Lancer Gradio (modèle entraîné + lien public)\n",
    "ADAPTER_TO_TEST = \"\"  #@param {type:\"string\"}\n",
    "#@markdown Vide = adaptateur final du run courant. Sinon, ex. `.../checkpoints/epoch_50`.\n",
    "INIT_LLM = \"auto\"  #@param [\"auto\", \"true\", \"false\"]\n",
    "#@markdown Le LM 5Hz (Thinking / génération de paroles) est optionnel ; `false` économise la VRAM.\n",
    "\n",
    "%cd {REPO_DIR}\n",
    "import os\n",
    "import shlex\n",
    "\n",
    "adapter = ADAPTER_TO_TEST or os.path.join(OUTPUT_DIR, \"final\")\n",
    "print(\"Chemin à coller dans « LoRA Adapter » :\")\n",
    "print(\"   \", adapter, \"\" if os.path.isdir(adapter) else \"(⚠️ introuvable pour l'instant)\")\n",
    "\n",
    "# Gradio cherche les poids ici (mêmes checkpoints que l'entraînement)\n",
    "os.environ[\"ACESTEP_CHECKPOINTS_DIR\"] = CKPT_DIR\n",
    "# Au démarrage, ACE-Step charge le LM depuis <code>/checkpoints sans tenir compte de\n",
    "# ACESTEP_CHECKPOINTS_DIR : on y place un lien vers les checkpoints du Drive.\n",
    "repo_ckpt = os.path.join(REPO_DIR, \"checkpoints\")\n",
    "if os.path.realpath(repo_ckpt) != os.path.realpath(CKPT_DIR):\n",
    "    if os.path.islink(repo_ckpt) or (os.path.isdir(repo_ckpt) and not os.listdir(repo_ckpt)):\n",
    "        os.remove(repo_ckpt) if os.path.islink(repo_ckpt) else os.rmdir(repo_ckpt)\n",
    "    if not os.path.exists(repo_ckpt):\n",
    "        os.symlink(CKPT_DIR, repo_ckpt)\n",
    "        print(\"Lien créé :\", repo_ckpt, \"->\", CKPT_DIR)\n",
    "    else:\n",
    "        print(\"⚠️\", repo_ckpt, \"existe déjà et n'est pas vide : le LM y sera cherché.\")\n",
    "gradio_args = [\n",
    "    \"python\", \"-m\", \"acestep.acestep_v15_pipeline\",\n",
    "    \"--share\",\n",
    "    \"--init_service\", \"true\",\n",
    "    \"--config_path\", MODEL_DIR,\n",
    "    \"--backend\", \"pt\",\n",
    "    \"--allowed-path\", WORKDIR,\n",
    "]\n",
    "if INIT_LLM != \"auto\":\n",
    "    gradio_args += [\"--init_llm\", INIT_LLM]\n",
    "!{shlex.join(gradio_args)}"
   ]
  }
 ]
}