#!/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()