aaronrene / gabriel-muse public
midi_shard.py python
174 lines 6.5 KB
Raw
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 13 days ago
1 """muse shard — partition a MIDI composition into bar-range shards for parallel agents.
2
3 Splits a track into N non-overlapping bar-range segments and writes each
4 shard as a separate MIDI file. An agent swarm can then work on shards in
5 parallel with zero risk of note-level conflicts, merging the shards back
6 together with ``muse mix``.
7
8 Usage::
9
10 muse shard tracks/full.mid --shards 4
11 muse shard tracks/full.mid --shards 8 --output-dir shards/
12 muse shard tracks/full.mid --bars-per-shard 16
13 muse shard tracks/full.mid --shards 4 --dry-run
14
15 Output::
16
17 Shard plan: tracks/full.mid → 4 shards
18 Total bars: 32 · ~8 bars per shard
19
20 Shard 0 bars 1– 8 → shards/full_shard_0.mid (28 notes)
21 Shard 1 bars 9–16 → shards/full_shard_1.mid (31 notes)
22 Shard 2 bars 17–24 → shards/full_shard_2.mid (24 notes)
23 Shard 3 bars 25–32 → shards/full_shard_3.mid (19 notes)
24
25 ✅ 4 shards written to shards/
26 """
27
28 from __future__ import annotations
29
30 import argparse
31 import logging
32 import pathlib
33 import sys
34
35 from muse.core.errors import ExitCode
36 from muse.core.repo import require_repo
37 from muse.core.validation import contain_path
38 from muse.plugins.midi._query import NoteInfo, load_track_from_workdir, notes_by_bar, notes_to_midi_bytes
39
40 logger = logging.getLogger(__name__)
41
42
43 def _shard_notes(
44 notes: list[NoteInfo],
45 bar_ranges: list[tuple[int, int]],
46 ) -> list[list[NoteInfo]]:
47 """Partition notes into groups by bar range, rebasing start ticks to 0."""
48 bars = notes_by_bar(notes)
49 if not notes:
50 return [[] for _ in bar_ranges]
51
52 tpb = notes[0].ticks_per_beat
53
54 shards: list[list[NoteInfo]] = []
55 for lo_bar, hi_bar in bar_ranges:
56 shard_notes: list[NoteInfo] = []
57 bar_offset = (lo_bar - 1) * 4 * tpb
58 for bar_num in range(lo_bar, hi_bar + 1):
59 for note in bars.get(bar_num, []):
60 rebased_tick = max(0, note.start_tick - bar_offset)
61 shard_notes.append(NoteInfo(
62 pitch=note.pitch,
63 velocity=note.velocity,
64 start_tick=rebased_tick,
65 duration_ticks=note.duration_ticks,
66 channel=note.channel,
67 ticks_per_beat=note.ticks_per_beat,
68 ))
69 shards.append(shard_notes)
70 return shards
71
72
73 def register(subparsers: "argparse._SubParsersAction[argparse.ArgumentParser]") -> None:
74 """Register the shard subcommand."""
75 parser = subparsers.add_parser("shard", help="Split a MIDI track into N bar-range shards for parallel agent work.", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
76 parser.add_argument("track", metavar="TRACK", help="Workspace-relative path to a .mid file.")
77 parser.add_argument("--shards", "-n", metavar="N", type=int, default=None, dest="num_shards", help="Number of shards to split into (mutually exclusive with --bars-per-shard).")
78 parser.add_argument("--bars-per-shard", "-b", metavar="N", type=int, default=None, help="Bars per shard (mutually exclusive with --shards).")
79 parser.add_argument("--output-dir", "-o", metavar="DIR", default="shards", help="Directory to write shard files (default: shards/).")
80 parser.add_argument("--dry-run", action="store_true", help="Preview shard plan without writing.")
81 parser.set_defaults(func=run)
82
83
84 def run(args: argparse.Namespace) -> None:
85 """Split a MIDI track into N bar-range shards for parallel agent work.
86
87 ``muse shard`` is the musical equivalent of partitioning a codebase for
88 a parallelised agent swarm. Each shard is a valid MIDI file covering a
89 non-overlapping bar range. Agents work on shards independently, then
90 the shards are recombined with ``muse mix``.
91
92 Specify either ``--shards N`` (divide evenly) or ``--bars-per-shard N``
93 (fixed shard size with a remainder shard at the end).
94 """
95 track: str = args.track
96 num_shards: int | None = args.num_shards
97 bars_per_shard: int | None = args.bars_per_shard
98 output_dir: str = args.output_dir
99 dry_run: bool = args.dry_run
100
101 if num_shards is not None and bars_per_shard is not None:
102 print("❌ --shards and --bars-per-shard are mutually exclusive.", file=sys.stderr)
103 raise SystemExit(ExitCode.USER_ERROR)
104 if num_shards is None and bars_per_shard is None:
105 num_shards = 4
106
107 root = require_repo()
108 result = load_track_from_workdir(root, track)
109 if result is None:
110 print(f"❌ Track '{track}' not found or not a valid MIDI file.", file=sys.stderr)
111 raise SystemExit(ExitCode.USER_ERROR)
112
113 notes, tpb = result
114 if not notes:
115 print(f" (track '{track}' contains no notes — nothing to shard)")
116 return
117
118 bars = notes_by_bar(notes)
119 all_bars = sorted(bars.keys())
120 total_bars = len(all_bars)
121 if total_bars == 0:
122 print(" (no bars detected)")
123 return
124
125 first_bar = all_bars[0]
126 last_bar = all_bars[-1]
127 bar_span = last_bar - first_bar + 1
128
129 # Determine bar-range splits
130 if num_shards is not None:
131 n = max(1, num_shards)
132 bps = max(1, bar_span // n)
133 else:
134 bps = max(1, bars_per_shard or 1)
135 n = (bar_span + bps - 1) // bps
136
137 bar_ranges: list[tuple[int, int]] = []
138 cur = first_bar
139 for i in range(n):
140 lo = cur
141 hi = lo + bps - 1 if i < n - 1 else last_bar
142 bar_ranges.append((lo, hi))
143 cur = hi + 1
144 if cur > last_bar:
145 break
146
147 track_stem = pathlib.Path(track).stem
148
149 print(f"\nShard plan: {track} → {len(bar_ranges)} shards")
150 print(f"Total bars: {total_bars} · ~{bps} bars per shard\n")
151
152 shard_notes_list = _shard_notes(notes, bar_ranges)
153
154 try:
155 out_dir = contain_path(root, output_dir)
156 except ValueError as exc:
157 print(f"❌ Invalid --output-dir: {exc}")
158 raise SystemExit(ExitCode.USER_ERROR)
159 for idx, ((lo, hi), shard_notes) in enumerate(zip(bar_ranges, shard_notes_list)):
160 out_name = f"{track_stem}_shard_{idx}.mid"
161 out_path = out_dir / out_name
162 print(
163 f" Shard {idx} bars {lo:>3}–{hi:>3} → {output_dir}/{out_name}"
164 f" ({len(shard_notes)} notes)"
165 )
166 if not dry_run:
167 out_dir.mkdir(parents=True, exist_ok=True)
168 midi_bytes = notes_to_midi_bytes(shard_notes, tpb) if shard_notes else notes_to_midi_bytes([], tpb)
169 out_path.write_bytes(midi_bytes)
170
171 if dry_run:
172 print("\n No files written (--dry-run).")
173 else:
174 print(f"\n✅ {len(bar_ranges)} shards written to {output_dir}/")
File History 1 commit
sha256:cf6265cea8c21d9228d90dec13ef6ec2dab5103d466db9cc4590681832de4bf8 docs(KD-STAGING): sync governance after KD-6b DONE Human 13 days ago