| """Run one real long-form Music3 generation in a fresh process. |
| |
| This is intentionally a small end-to-end harness. It invokes the normal |
| phased pipeline with the explicit experimental long-form flag, fixed |
| flush=32, and the existing chunked Flow path. It does not run the M9-A |
| AR-only benchmark or compare multiple modes in one process. |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import json |
| import platform |
| import subprocess |
| import sys |
| import threading |
| import time |
| import traceback |
| import wave |
| from dataclasses import dataclass |
| from pathlib import Path |
| from typing import Any, Sequence |
|
|
| import mlx.core as mx |
|
|
| ROOT = Path(__file__).resolve().parent.parent |
| if str(ROOT) not in sys.path: |
| sys.path.insert(0, str(ROOT)) |
|
|
| from generation.ar_profiler import ( |
| _memory_pressure, |
| _physical_memory_footprint_bytes, |
| _physical_memory_bytes, |
| _process_memory, |
| _swap_used_bytes, |
| ) |
| from generation.phased_pipeline import ( |
| DEFAULT_MODEL, |
| PhasedPipeline, |
| ) |
|
|
|
|
| ALLOWED_DURATIONS = (120, 180, 240) |
| DEFAULT_OUTPUT_ROOT = ROOT / "outputs" / "long_form_test" |
| DEFAULT_CAPTION = ( |
| "dark melodic techno, 128 BPM, female airy vocal, lonely but not sad, " |
| "slow atmospheric intro, heavy bass enters after 16 bars, " |
| "long continuous arrangement" |
| ) |
| DEFAULT_LYRICS = ( |
| "[verse]\n" |
| "Under a dead sky, I count the lights\n" |
| "Every quiet window holds a different flame\n" |
| "I walk the avenue where all the old trains fade\n" |
| "The night keeps turning softly, calling out my name\n" |
| "[pre-chorus]\n" |
| "The city breathes in neon tides\n" |
| "The wires sing above me, restless in the rain\n" |
| "I feel the distant thunder moving through the lane\n" |
| "One more step toward morning, one more pulse again\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[verse]\n" |
| "Every shadow turns into a sign\n" |
| "Every empty station leaves a trace behind\n" |
| "I gather up the echoes that I thought I lost in time\n" |
| "The road is wide and open, and the dark is mine\n" |
| "[pre-chorus]\n" |
| "The city breathes in neon tides\n" |
| "The rain is on the rooftops, silver in the light\n" |
| "The beat beneath the pavement keeps us synchronized\n" |
| "One more step together, one more turn tonight\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[bridge]\n" |
| "Hold the silence, let the pulse return\n" |
| "Leave behind the lessons that were never ours to learn\n" |
| "There is still a signal underneath the burn\n" |
| "Follow where it leads us, watch the colors turn\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[verse]\n" |
| "Morning finds us moving through a field of blue\n" |
| "All the roads we wandered still come back to you\n" |
| "I can hear the open sky beginning to renew\n" |
| "Every distant heartbeat says we will get through\n" |
| "[pre-chorus]\n" |
| "The city breathes in neon tides\n" |
| "The last stars disappear but leave the skyline bright\n" |
| "We carry all the motion of a thousand nights\n" |
| "One more step together, one more turn in flight\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[instrumental]\n" |
| "[verse]\n" |
| "Through the open doorway, warmer air arrives\n" |
| "All the hidden colors wake and come alive\n" |
| "We were never waiting, we were learning how to drive\n" |
| "With the sound around us, we remember why\n" |
| "[bridge]\n" |
| "Hold the silence, let the pulse return\n" |
| "Let the old horizons lose the shape they used to wear\n" |
| "Every new direction is a promise in the air\n" |
| "We can keep on walking, we can meet it there\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[instrumental]\n" |
| "[verse]\n" |
| "Every shadow turns into a sign\n" |
| "Every moving light becomes a thread we can align\n" |
| "I can feel the future in the rhythm of the line\n" |
| "Keep the signal open, keep your hand in mine\n" |
| "[pre-chorus]\n" |
| "The city breathes in neon tides\n" |
| "The beat is getting stronger as the hours slide\n" |
| "Nothing in the distance has to stay outside\n" |
| "One more step together, let the current guide\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| "[verse]\n" |
| "When the morning opens, there is more to say\n" |
| "We will take the long road, we will find a way\n" |
| "Every beat remembers everything we made\n" |
| "And the song keeps rising with the light of day\n" |
| "[chorus]\n" |
| "Rise above the noise, carry me through the night\n" |
| "Hold on to the rhythm, keep the fire bright\n" |
| "When the world is breaking, we can make it right\n" |
| "Rise above the noise, carry me through the night\n" |
| ) |
| DEFAULT_SEED = 0 |
| DEFAULT_STEPS = 30 |
| PRESSURE_ORDER = {None: -1, "normal": 0, "warning": 1, "critical": 2} |
|
|
|
|
| def validate_long_form_seconds(seconds: int) -> int: |
| if seconds not in ALLOWED_DURATIONS: |
| raise ValueError(f"seconds must be one of {ALLOWED_DURATIONS}") |
| return seconds |
|
|
|
|
| def long_form_output_dir(root: str | Path, seconds: int) -> Path: |
| validate_long_form_seconds(seconds) |
| return Path(root) / f"{seconds}s" |
|
|
|
|
| def read_wav_metadata(path: Path) -> dict[str, Any]: |
| with wave.open(str(path), "rb") as handle: |
| frames = handle.getnframes() |
| sample_rate = handle.getframerate() |
| channels = handle.getnchannels() |
| return { |
| "frames": frames, |
| "sample_rate": sample_rate, |
| "channels": channels, |
| "duration_seconds": frames / sample_rate if sample_rate else 0.0, |
| } |
|
|
|
|
| def _event_time(events: Sequence[dict[str, Any]], name: str) -> float | None: |
| for event in events: |
| if event.get("name") == name: |
| value = event.get("t_s") |
| return float(value) if value is not None else None |
| return None |
|
|
|
|
| def _duration_between( |
| events: Sequence[dict[str, Any]], start_name: str, end_name: str |
| ) -> float | None: |
| start = _event_time(events, start_name) |
| end = _event_time(events, end_name) |
| if start is None or end is None or end < start: |
| return None |
| return round(end - start, 3) |
|
|
|
|
| def stage_timings(events: Sequence[dict[str, Any]]) -> dict[str, float | None]: |
| return { |
| "ar_stage_seconds": _duration_between(events, "stage_a_start", "stage_a_done"), |
| "flow_vocoder_seconds": _duration_between( |
| events, "stage_b_weights", "stage_b_done" |
| ), |
| "flow_stage_seconds": _duration_between( |
| events, "stage_b_start", "stage_b_done" |
| ), |
| "wav_write_seconds": _duration_between(events, "stage_b_done", "write_done"), |
| } |
|
|
|
|
| def _run(command: list[str]) -> str | None: |
| try: |
| result = subprocess.run(command, capture_output=True, text=True, check=False) |
| except OSError: |
| return None |
| return result.stdout.strip() if result.returncode == 0 else None |
|
|
|
|
| def _environment(model: str) -> dict[str, Any]: |
| import importlib.metadata |
|
|
| def version(name: str) -> str | None: |
| try: |
| return importlib.metadata.version(name) |
| except importlib.metadata.PackageNotFoundError: |
| return None |
|
|
| return { |
| "commit_hash": _run(["git", "rev-parse", "HEAD"]), |
| "mac_model": _run(["sysctl", "-n", "hw.model"]), |
| "chip": "Apple M4" if _run(["sysctl", "-n", "hw.model"]) == "Mac16,10" else None, |
| "apple_silicon": platform.machine(), |
| "ram_bytes": _physical_memory_bytes(), |
| "macos": platform.mac_ver()[0], |
| "python": sys.version.split()[0], |
| "mlx": version("mlx"), |
| "mlx_audio": version("mlx-audio"), |
| "model": model, |
| "quantization": "4bit" if "4bit" in model.lower() else "unknown", |
| } |
|
|
|
|
| @dataclass |
| class SystemMonitor: |
| interval_seconds: float = 60.0 |
|
|
| def __post_init__(self) -> None: |
| self.initial_swap_bytes = _swap_used_bytes() |
| self.latest_swap_bytes = self.initial_swap_bytes |
| self.max_swap_delta_bytes = 0 |
| self.max_rss_bytes = 0 |
| self.max_footprint_bytes = 0 |
| self.max_mlx_active_bytes = 0 |
| self.max_mlx_peak_bytes = 0 |
| self.highest_pressure: str | None = None |
| self.sample_count = 0 |
| self.critical_event = threading.Event() |
| self._stop_event = threading.Event() |
| self._thread: threading.Thread | None = None |
|
|
| def sample(self) -> None: |
| rss, _vm = _process_memory() |
| footprint = _physical_memory_footprint_bytes() or 0 |
| swap = _swap_used_bytes() |
| pressure = _memory_pressure() |
| self.max_rss_bytes = max(self.max_rss_bytes, rss) |
| self.max_footprint_bytes = max(self.max_footprint_bytes, footprint) |
| self.max_mlx_active_bytes = max(self.max_mlx_active_bytes, int(mx.get_active_memory())) |
| self.max_mlx_peak_bytes = max(self.max_mlx_peak_bytes, int(mx.get_peak_memory())) |
| self.latest_swap_bytes = swap |
| if swap is not None and self.initial_swap_bytes is not None: |
| self.max_swap_delta_bytes = max( |
| self.max_swap_delta_bytes, swap - self.initial_swap_bytes |
| ) |
| if PRESSURE_ORDER.get(pressure, -1) > PRESSURE_ORDER.get(self.highest_pressure, -1): |
| self.highest_pressure = pressure |
| if pressure == "critical": |
| self.critical_event.set() |
| self.sample_count += 1 |
|
|
| def _run(self) -> None: |
| while not self._stop_event.wait(self.interval_seconds): |
| self.sample() |
|
|
| def start(self) -> None: |
| self.sample() |
| self._thread = threading.Thread(target=self._run, name="long-form-monitor") |
| self._thread.daemon = True |
| self._thread.start() |
|
|
| def stop(self) -> None: |
| self._stop_event.set() |
| if self._thread is not None: |
| self._thread.join(timeout=max(5.0, self.interval_seconds + 1.0)) |
| self.sample() |
|
|
|
|
| def _gb(value: int | None) -> float | None: |
| return round(value / 1024**3, 3) if value is not None else None |
|
|
|
|
| def run_long_form_test( |
| seconds: int, |
| *, |
| output_root: str | Path = DEFAULT_OUTPUT_ROOT, |
| ) -> dict[str, Any]: |
| validate_long_form_seconds(seconds) |
| run_dir = long_form_output_dir(output_root, seconds) |
| run_dir.mkdir(parents=True, exist_ok=True) |
| output_path = run_dir / "output.wav" |
| config_path = run_dir / "config.json" |
| summary_path = run_dir / "summary.json" |
| log_path = run_dir / "benchmark.log" |
| environment = _environment(DEFAULT_MODEL) |
| config = { |
| "target_seconds": seconds, |
| "caption": DEFAULT_CAPTION, |
| "lyrics": DEFAULT_LYRICS, |
| "seed": DEFAULT_SEED, |
| "steps": DEFAULT_STEPS, |
| "model": DEFAULT_MODEL, |
| "quantization": "4bit", |
| "experimental_long_form": True, |
| "flush_interval": 32, |
| "output": str(output_path), |
| "environment": environment, |
| } |
| config_path.write_text(json.dumps(config, indent=2, ensure_ascii=False) + "\n") |
| log_lines = [ |
| f"start target_seconds={seconds} seed={DEFAULT_SEED} steps={DEFAULT_STEPS}", |
| "mode=experimental_long_form flush_interval=32 phased_pipeline=true", |
| f"output={output_path}", |
| ] |
|
|
| monitor = SystemMonitor() |
| pipeline_report: dict[str, Any] | None = None |
| error: str | None = None |
| started = time.perf_counter() |
| monitor.start() |
| try: |
| pipeline = PhasedPipeline( |
| label=f"long_form_test/{seconds}s", model_path=DEFAULT_MODEL |
| ) |
| pipeline_report = pipeline.generate( |
| caption=DEFAULT_CAPTION, |
| lyrics=DEFAULT_LYRICS, |
| duration=seconds, |
| steps=DEFAULT_STEPS, |
| seed=DEFAULT_SEED, |
| output=str(output_path), |
| experimental_long_form=True, |
| stop_event=monitor.critical_event, |
| ) |
| except BaseException as exc: |
| error = ( |
| "KeyboardInterrupt: manually stopped after sustained swap growth" |
| if isinstance(exc, KeyboardInterrupt) |
| else f"{type(exc).__name__}: {exc}" |
| ) |
| log_lines.append(traceback.format_exc()) |
| finally: |
| monitor.stop() |
| elapsed = time.perf_counter() - started |
|
|
| wav_metadata = None |
| if output_path.is_file(): |
| try: |
| wav_metadata = read_wav_metadata(output_path) |
| except (OSError, wave.Error) as exc: |
| error = error or f"WAVReadError: {exc}" |
|
|
| events = pipeline_report.get("events", []) if pipeline_report else [] |
| timings = stage_timings(events) |
| actual_seconds = wav_metadata["duration_seconds"] if wav_metadata else None |
| duration_matches = ( |
| actual_seconds is not None |
| and abs(actual_seconds - seconds) <= max(1.0, seconds * 0.02) |
| ) |
| if error: |
| status = "failed" |
| elif monitor.critical_event.is_set(): |
| status = "safety_stopped" |
| error = "critical memory pressure reported by monitor" |
| elif not output_path.is_file() or not duration_matches: |
| status = "incomplete" |
| else: |
| status = "completed" |
|
|
| pipeline_peak = None |
| if pipeline_report is not None: |
| pipeline_peak = float(pipeline_report.get("max_stage_peak_gb", 0.0)) |
| monitor_peak = _gb(monitor.max_mlx_peak_bytes) or 0.0 |
| end_to_end = round(elapsed, 3) |
| summary = { |
| "status": status, |
| "success": status == "completed", |
| "error": error, |
| "target_seconds": seconds, |
| "actual_wav_seconds": round(actual_seconds, 3) if actual_seconds is not None else None, |
| "duration_matches_target": duration_matches, |
| "ar_stage_seconds": timings["ar_stage_seconds"], |
| "flow_vocoder_seconds": timings["flow_vocoder_seconds"], |
| "flow_stage_seconds": timings["flow_stage_seconds"], |
| "wav_write_seconds": timings["wav_write_seconds"], |
| "end_to_end_seconds": end_to_end, |
| "rtf": round(elapsed / actual_seconds, 3) if actual_seconds else None, |
| "mlx_peak_gb": round(max(pipeline_peak or 0.0, monitor_peak), 3), |
| "process_peak_rss_gb": _gb(monitor.max_rss_bytes), |
| "process_peak_footprint_gb": _gb(monitor.max_footprint_bytes), |
| "system_swap_delta_gb": _gb(monitor.max_swap_delta_bytes), |
| "highest_memory_pressure": monitor.highest_pressure, |
| "system_responsiveness": "not_directly_measured_by_runner", |
| "monitor_sample_count": monitor.sample_count, |
| "wav": wav_metadata, |
| "pipeline_report": pipeline_report, |
| "config": config, |
| } |
| log_lines.extend( |
| [ |
| f"status={status}", |
| f"actual_wav_seconds={summary['actual_wav_seconds']}", |
| f"end_to_end_seconds={end_to_end}", |
| f"ar_stage_seconds={timings['ar_stage_seconds']}", |
| f"flow_vocoder_seconds={timings['flow_vocoder_seconds']}", |
| f"mlx_peak_gb={summary['mlx_peak_gb']}", |
| f"process_peak_rss_gb={summary['process_peak_rss_gb']}", |
| f"process_peak_footprint_gb={summary['process_peak_footprint_gb']}", |
| f"system_swap_delta_gb={summary['system_swap_delta_gb']}", |
| f"highest_memory_pressure={summary['highest_memory_pressure']}", |
| ] |
| ) |
| for event in events: |
| log_lines.append( |
| f"event={event.get('name')} t_s={event.get('t_s')} note={event.get('note', '')}" |
| ) |
| log_path.write_text("\n".join(log_lines) + "\n") |
| summary_path.write_text(json.dumps(summary, indent=2, ensure_ascii=False) + "\n") |
| print(json.dumps(summary, indent=2, ensure_ascii=False)) |
| return summary |
|
|
|
|
| def _parser() -> argparse.ArgumentParser: |
| parser = argparse.ArgumentParser(description="Music3 M4 16GB long-form real-world test") |
| parser.add_argument("--seconds", type=int, choices=ALLOWED_DURATIONS, required=True) |
| parser.add_argument("--output-root", type=Path, default=DEFAULT_OUTPUT_ROOT) |
| return parser |
|
|
|
|
| def main(argv: Sequence[str] | None = None) -> int: |
| args = _parser().parse_args(argv) |
| summary = run_long_form_test(args.seconds, output_root=args.output_root) |
| return 0 if summary["status"] == "completed" else 2 |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
