"""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 ( # noqa: E402 _memory_pressure, _physical_memory_footprint_bytes, _physical_memory_bytes, _process_memory, _swap_used_bytes, ) from generation.phased_pipeline import ( # noqa: E402 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: # noqa: BLE001 - preserve a useful failure artifact 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())