generate.py

#!/usr/bin/env python
"""Generation MiniMax-Music3 en local (Apple Silicon / MPS, CUDA, ou CPU).

Le modele est officiellement CUDA-only. Ce script tente MPS avec fallback CPU
sur les ops non supportees. Voir README.md.
"""

import argparse
import os
import sys
import time
from pathlib import Path

ROOT = Path(os.environ.get("MM3_ROOT", Path(__file__).resolve().parent.parent))
DEFAULT_MODEL_DIR = ROOT / "models" / "MiniMax-Music3"
DEFAULT_OUT_DIR = ROOT / "outputs"
MLX_LM_DIR = ROOT / "models" / "lm-mlx"

# Doivent etre poses avant l'import de torch.
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
os.environ.setdefault("PYTORCH_MPS_HIGH_WATERMARK_RATIO", "0.0")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")

import soundfile as sf  # noqa: E402
import torch  # noqa: E402
from diffusers import ModularPipeline  # noqa: E402


def pick_device(requested):
    if requested != "auto":
        return requested
    if torch.cuda.is_available():
        return "cuda"
    if torch.backends.mps.is_available():
        return "mps"
    return "cpu"


def pick_dtype(name, device):
    if name != "auto":
        return {"bfloat16": torch.bfloat16, "float16": torch.float16, "float32": torch.float32}[name]
    if device == "cpu":
        return torch.float32
    return torch.bfloat16


def enable_block_profiling():
    """Chronometre chaque bloc du pipeline (AR, denoise, decode)."""
    import time as _time
    from collections import defaultdict

    from diffusers.modular_pipelines import ModularPipelineBlocks
    from diffusers.modular_pipelines.minimax_music3 import before_denoise, decoders, denoise, encoders

    stats = defaultdict(float)
    for module in (encoders, before_denoise, denoise, decoders):
        for name, obj in vars(module).items():
            if not (isinstance(obj, type) and issubclass(obj, ModularPipelineBlocks)):
                continue
            if obj.__module__ != module.__name__ or "__call__" not in obj.__dict__:
                continue

            def wrap(cls, original):
                def timed_call(self, components, state, **kwargs):
                    torch.mps.synchronize()
                    start = _time.perf_counter()
                    result = original(self, components, state, **kwargs)
                    torch.mps.synchronize()
                    stats[cls.__name__] += _time.perf_counter() - start
                    return result

                cls.__call__ = timed_call

            wrap(obj, obj.__dict__["__call__"])

    # Fonctions internes de la boucle autoregressive.
    for fname in ("_sample_top_k", "_generate_depth_codes", "_embed_audio_frame"):
        original_fn = getattr(encoders, fname)

        def wrap_fn(name, fn):
            def timed_fn(*a, **kw):
                torch.mps.synchronize()
                start = _time.perf_counter()
                out = fn(*a, **kw)
                torch.mps.synchronize()
                stats[name] += _time.perf_counter() - start
                return out

            setattr(encoders, name, timed_fn)

        wrap_fn(fname, original_fn)
    return stats


def ensure_local_paths(model_dir):
    """Le modular_model_index.json livre pointe chaque composant sur le repo id du Hub,
    ce qui fait re-telecharger 28 Go malgre les poids locaux. On le repointe sur le dossier."""
    import json

    index = model_dir / "modular_model_index.json"
    data = json.loads(index.read_text(encoding="utf-8"))
    changed = False
    for value in data.values():
        if isinstance(value, list) and len(value) == 3 and isinstance(value[2], dict):
            spec = value[2]
            if spec.get("pretrained_model_name_or_path") != str(model_dir):
                spec["pretrained_model_name_or_path"] = str(model_dir)
                changed = True
    if changed:
        index.write_text(json.dumps(data, indent=2), encoding="utf-8")
        print(f"modular_model_index.json repointe sur {model_dir}")


def read_text(value):
    """Accepte un chemin de fichier ou du texte direct.

    Le test de chemin est garde : une chaine multiligne ou trop longue ne peut pas
    etre un chemin, et la passer a Path.exists() leve OSError 63 (nom trop long).
    """
    if value is None:
        return None
    if "\n" not in value and len(value) < 1024:
        try:
            path = Path(value)
            if path.is_file():
                return path.read_text(encoding="utf-8")
        except OSError:
            pass
    return value


def main():
    ap = argparse.ArgumentParser(description="Genere un morceau avec MiniMax-Music3.")
    ap.add_argument("-p", "--prompt", required=True,
                    help="Description musicale (texte ou chemin vers un .txt).")
    ap.add_argument("-l", "--lyrics", default=None,
                    help="Paroles avec balises [verse]/[chorus] (texte ou chemin vers un .txt).")
    ap.add_argument("-d", "--duration", type=float, default=60.0,
                    help="Duree cible en secondes (max ~300).")
    ap.add_argument("-s", "--steps", type=int, default=30, help="Pas de denoising.")
    ap.add_argument("--seed", type=int, default=7)
    ap.add_argument("-o", "--out", default=None, help="Fichier WAV de sortie.")
    ap.add_argument("--model-dir", default=str(DEFAULT_MODEL_DIR))
    ap.add_argument("--device", default="auto", choices=["auto", "cuda", "mps", "cpu"])
    ap.add_argument("--dtype", default="auto", choices=["auto", "bfloat16", "float16", "float32"])
    ap.add_argument("--lm", default="auto", choices=["auto", "mlx", "torch"],
                    help="Backend du language_model. mlx = quantifie 4 bits, indispensable sous 32 Go de RAM.")
    ap.add_argument("--profile", action="store_true", help="Chronometre chaque etage du pipeline.")
    ap.add_argument("--cpu-offload", action="store_true",
                    help="Charge les composants a la demande (utile si la RAM sature).")
    args = ap.parse_args()

    model_dir = Path(args.model_dir)
    if not (model_dir / "modular_model_index.json").exists():
        sys.exit(f"Poids introuvables dans {model_dir}. Lancer scripts/download.sh d'abord.")

    ensure_local_paths(model_dir)
    block_stats = enable_block_profiling() if args.profile else None

    device = pick_device(args.device)
    dtype = pick_dtype(args.dtype, device)
    prompt = read_text(args.prompt)
    lyrics = read_text(args.lyrics) or ""

    out = Path(args.out) if args.out else DEFAULT_OUT_DIR / f"song-{int(time.time())}.wav"
    out.parent.mkdir(parents=True, exist_ok=True)

    print(f"device={device} dtype={str(dtype).split('.')[-1]} duration={args.duration}s steps={args.steps}")

    lm_mode = args.lm
    if lm_mode == "auto":
        lm_mode = "mlx" if MLX_LM_DIR.exists() and device != "cuda" else "torch"
    if lm_mode == "mlx" and not MLX_LM_DIR.exists():
        sys.exit(f"LM MLX absent de {MLX_LM_DIR}. Lancer scripts/convert_lm_mlx.py.")
    print(f"language_model: {lm_mode}")

    t0 = time.time()
    if args.cpu_offload:
        from diffusers import ComponentsManager

        manager = ComponentsManager()
        manager.enable_auto_cpu_offload(device=device)
        pipe = ModularPipeline.from_pretrained(str(model_dir), components_manager=manager)
        pipe.load_components(dtype=dtype)
    else:
        pipe = ModularPipeline.from_pretrained(str(model_dir))
        if lm_mode == "mlx":
            sys.path.insert(0, str(Path(__file__).resolve().parent))
            from mlx_bridge import MlxLanguageModel

            names = [n for n in pipe.pretrained_component_names if n != "language_model"]
            pipe.load_components(names=names, dtype=dtype)
            pipe.to(device)
            pipe.update_components(language_model=MlxLanguageModel(MLX_LM_DIR, device, dtype))
        else:
            pipe.load_components(dtype=dtype)
            pipe.to(device)
    print(f"composants charges en {time.time() - t0:.1f}s")

    try:
        generator = torch.Generator(device=device).manual_seed(args.seed)
    except Exception:
        generator = torch.Generator().manual_seed(args.seed)

    t1 = time.time()
    audio = pipe(
        prompt=prompt,
        lyrics=lyrics,
        audio_duration=args.duration,
        num_inference_steps=args.steps,
        generator=generator,
        output="audios",
    )[0]
    print(f"generation en {time.time() - t1:.1f}s")

    if block_stats:
        width = max(len(k) for k in block_stats)
        print("\nprofil blocs:")
        for name in sorted(block_stats, key=block_stats.get, reverse=True):
            print(f"  {name:<{width}}  {block_stats[name]:7.1f}s")
        if lm_mode == "mlx":
            from mlx_bridge import report

            report()

    if isinstance(audio, torch.Tensor):
        audio = audio.float().cpu().numpy()
    # Le decoder sort (channels, samples), soundfile attend (samples, channels).
    audio = audio.T if audio.ndim == 2 and audio.shape[0] < audio.shape[1] else audio
    sf.write(str(out), audio, pipe.sampling_rate)
    seconds = audio.shape[0] / pipe.sampling_rate
    print(f"ecrit: {out}  ({pipe.sampling_rate} Hz, {seconds:.1f}s, {audio.shape[1]} canaux)")


if __name__ == "__main__":
    main()