cli.py

"""cli.py — 项目入口 (PM v0.2 §10)。

命令:
  python cli.py baseline --seconds 5      # Milestone 1: 跑通并测量
  python cli.py generate "描述" --seconds 30  # 直接生成 (阶段驻留管线)
  python cli.py idea "想法" --seconds 30  # 完整 Idea→Music (重写 caption + 生成)
  python cli.py webui                     # 启动 WebUI (浏览器测试入口)
"""
from __future__ import annotations

import sys
from pathlib import Path

import click

ROOT = Path(__file__).parent


@click.group()
def main() -> None:
    """Music3-Mac16: MiniMax Music3 on 16GB Apple Silicon Mac."""


@main.command()
@click.option("--host", default="127.0.0.1", help="监听地址")
@click.option("--port", default=8642, type=int, help="端口")
def webui(host: str, port: int) -> None:
    """启动 WebUI (浏览器测试入口)。

    用法: music-minimax webui  → 打开 http://127.0.0.1:8642
    """
    import webbrowser

    from webui.app import start

    url = f"http://{host}:{port}"
    click.echo(f"启动 WebUI: {url}")
    webbrowser.open(url)
    start(host=host, port=port)


@main.command()
@click.option("--seconds", default=5, type=click.IntRange(1, 60), help="目标生成秒数")
@click.option("--label", default="baseline", help="本次运行的标签（写进 profile json）")
def baseline(seconds: int, label: str) -> None:
    """Milestone 1: Baseline — 跑通现有 MLX Music3 并记录峰值内存/速度/正确性。"""
    from generation.pipeline import BaselineRunner

    runner = BaselineRunner(label=label)
    runner.run(target_seconds=seconds)


@main.command()
@click.argument("prompt")
@click.option("--seconds", default=15, type=click.IntRange(1, 240), help="目标生成秒数")
@click.option("--lyrics", default=None, help="结构化歌词 (可选, 缺省为 [instrumental])")
@click.option("--steps", default=30, type=click.IntRange(1, 50), help="Flow-matching 步数")
@click.option("--seed", default=0, type=int, help="随机种子")
@click.option("--label", default="idea", help="输出目录标签")
@click.option("--model", default="models/runtime/MiniMax-Music3-4bit",
              help="模型路径 (4bit / mxfp8 / 8bit 等 mlx-community 变体)")
@click.option("--experimental-long-form", is_flag=True,
              help="显式开启 120/180/240 秒实验路径，固定使用 flush=32")
def generate(prompt: str, seconds: int, lyrics: str | None, steps: int, seed: int,
             label: str, model: str, experimental_long_form: bool) -> None:
    """Idea→Music: 从自然语言描述生成 WAV（阶段驻留管线）。

    用法: music-minimax generate "dark melodic techno, 128 BPM" --seconds 30
          music-minimax generate "pop ballad" --model models/runtime/MiniMax-Music3-mxfp8
    """
    from generation.phased_pipeline import PhasedPipeline

    pipeline = PhasedPipeline(label=label, model_path=model)
    output = None
    if experimental_long_form:
        long_form_dir = ROOT / "outputs" / "long_form_test" / f"{seconds}s"
        output = str(long_form_dir / "output.wav")
    report = pipeline.generate(
        caption=prompt,
        lyrics=lyrics or "[instrumental]",
        duration=seconds,
        steps=steps,
        seed=seed,
        output=output,
        experimental_long_form=experimental_long_form,
    )
    click.echo(f"✓ 生成完成: {report['output_wav']}")
    click.echo(f"  峰值内存: {report['max_stage_peak_gb']} GB (预算 11.5GB)")
    click.echo(f"  AR 阶段  : {report['ar_peak_gb']} GB | Flow 阶段: {report['flow_peak_gb']} GB")


@main.command()
@click.argument("idea")
@click.option("--lyrics", default=None, help="可选歌词 (带段落标签)")
@click.option("--seconds", default=30, type=click.IntRange(1, 60), help="目标生成秒数")
@click.option("--rewriter-model", default=None,
              help="本地重写模型 (默认 mlx-community/Qwen3-8B-4bit, 首次自动下载)")
@click.option("--label", default="idea_to_music", help="输出目录标签")
@click.option("--dry-run", is_flag=True, help="只重写 caption 不生成音乐")
def idea(idea: str, lyrics: str | None, seconds: int, rewriter_model: str | None,
         label: str, dry_run: bool) -> None:
    """Idea→Music 完整工作流: 想法 → caption 重写 → 本地生成 WAV。

    用法: music-minimax idea "dark melodic techno with female airy vocal, lonely"
    """
    from caption.rewriter import rewrite_caption
    from generation.phased_pipeline import PhasedPipeline

    click.echo("① 重写想法 → 结构化 caption (本地模型)...")
    caption = rewrite_caption(
        idea=idea,
        lyrics=lyrics or "",
        model=rewriter_model or "mlx-community/Qwen3-8B-4bit",
    )
    click.echo(f"\n── 生成的 caption ──\n{caption}\n─────────────────────")

    if dry_run:
        click.echo("(--dry-run: 只输出 caption, 不生成音乐)")
        return

    click.echo(f"② 生成 {seconds}s 音乐 (阶段驻留管线)...")
    pipeline = PhasedPipeline(label=label)
    report = pipeline.generate(
        caption=caption,
        lyrics=lyrics or "[instrumental]",
        duration=seconds,
    )
    click.echo(f"✓ 完整工作流完成: {report['output_wav']}")
    click.echo(f"  峰值内存: {report['max_stage_peak_gb']} GB (预算 11.5GB)")


if __name__ == "__main__":
    main()