app.py

"""LilyScript — a symbolic-music AIGC app built on the LilyletNotaGen model.

Left column:
  (1) generation parameter panel   (2) streaming run log
  (3) .lyl file list (session outputs + built-in examples)   (4) editable lyl editor
Right column:
  sheet-music panel (placeholder for now; later a Lilylet music score renderer
  reusing the lilylet-live-editor pipeline).

Generation streams patch-by-patch: raw decoded text (with `[r:x/y]` stream
markers) goes to the run log, while the measure-segmented postprocessed text
fills the editor. The backend is the int8 + two-level KV-cache ONNX generator
(see lilyscript/generator.py); models load from a local dir for now.
"""

import os
import re
import time
import json

import gradio as gr

from lilyscript.generator import StreamingLilyletGenerator

HERE = os.path.dirname(os.path.abspath(__file__))
# TODO: swap for huggingface_hub.snapshot_download(repo_id=...) to pull the int8
# ONNX weights from the hub instead of a local dir.
MODEL_DIR = os.environ.get('LILYSCRIPT_MODEL_DIR', os.path.join(HERE, 'models'))
ASSET_DIR = os.path.join(HERE, 'assets')
EXAMPLES_DIR = os.path.join(HERE, 'examples')
OUTPUT_DIR = os.path.join(HERE, 'outputs')

EXAMPLE_PREFIX = '\U0001F4C4 '		# 📄 examples
OUTPUT_PREFIX = '✨ '			# ✨ session outputs

# Suggested metadata values (editable — the dropdowns allow custom input), loaded
# from assets/styles.json. Drawn from the NotaGenX period/instrumentation
# vocabulary + values seen in examples.
_STYLES = json.load(open(os.path.join(ASSET_DIR, 'styles.json'), encoding='utf-8'))
COMPOSERS = _STYLES['composers']
GENRES = _STYLES['genres']
INSTRUMENTS = _STYLES['instruments']

_GEN = None


def get_generator ():
	'''Lazily build the (heavy) ONNX generator on first use.'''
	global _GEN
	if _GEN is None:
		_GEN = StreamingLilyletGenerator(MODEL_DIR, ASSET_DIR)
	return _GEN


def load_examples ():
	'''Read built-in example .lyl files into a {label: text} dict.'''
	store = {}
	if os.path.isdir(EXAMPLES_DIR):
		for name in sorted(os.listdir(EXAMPLES_DIR)):
			if name.endswith('.lyl'):
				with open(os.path.join(EXAMPLES_DIR, name), encoding='utf-8') as f:
					store[EXAMPLE_PREFIX + name] = f.read()
	return store


def load_outputs ():
	'''Read previously-generated .lyl files from the outputs dir into a
	{label: text} dict, so past session outputs survive a server restart.'''
	store = {}
	if os.path.isdir(OUTPUT_DIR):
		for name in sorted(os.listdir(OUTPUT_DIR)):
			if name.endswith('.lyl'):
				with open(os.path.join(OUTPUT_DIR, name), encoding='utf-8') as f:
					store[OUTPUT_PREFIX + name[:-4]] = f.read()
	return store


def load_library ():
	'''Initial file list: built-in examples + any persisted session outputs.'''
	return {**load_examples(), **load_outputs()}


_STYLE_LINE_RE = re.compile(r'^\[(composer|genre|instrument)\s+".*"\]\s*$')


def sync_prompt (composer, genre, instrument, current):
	'''Rewrite the metadata-prompt text from the three style dropdowns.

	The `[composer/genre/instrument "..."]` lines are regenerated from the
	dropdowns and placed at the top; any other lines the user typed (e.g.
	`[key "..."]`) are preserved below in their original order.
	'''
	lines = []
	for field, value in (('composer', composer), ('genre', genre), ('instrument', instrument)):
		value = (value or '').strip()
		if value:
			lines.append(f'[{field} "{value}"]')
	# keep every line that isn't one of the three managed style lines
	for ln in (current or '').splitlines():
		if not _STYLE_LINE_RE.match(ln.strip()):
			if ln.strip():
				lines.append(ln)
	return '\n'.join(lines)


def run_generation (prompt, measures, temperature, max_patches, seed, store, top_k=0, top_p=0.9):
	'''Streaming generate callback. Yields updates for (log, editor, file_list, store).

	store: {label: lyl_text} dict held in gr.State; the produced document is added
	to it under a timestamped label once generation finishes.

	top_k / top_p have fixed defaults (no UI controls); pass them explicitly to override.
	'''
	gen = get_generator()
	meas = int(measures) if measures and int(measures) > 0 else None
	store = dict(store or {})

	raw = pretty = ''
	for raw, pretty, done in gen.generate_stream(
		prompt_text=prompt or '', max_patches=int(max_patches), temperature=float(temperature),
		top_k=int(top_k), top_p=float(top_p), measures=meas, seed=int(seed)):
		# stream: update log + editor, keep file list / store unchanged
		yield raw, pretty, gr.update(), store

	# finished: persist the document, refresh the file list, select the new entry
	label = OUTPUT_PREFIX + time.strftime('%H%M%S') + ('_m%d' % meas if meas else '') + '_s%d' % int(seed)
	store[label] = pretty
	os.makedirs(OUTPUT_DIR, exist_ok=True)
	with open(os.path.join(OUTPUT_DIR, label.replace(OUTPUT_PREFIX, '') + '.lyl'), 'w', encoding='utf-8') as f:
		f.write(pretty)
	yield raw, pretty, gr.update(choices=list(store.keys()), value=label), store


def load_file (label, store):
	'''File-list selection -> load that document into the editor.'''
	return (store or {}).get(label, '')


SHEET_PLACEHOLDER = '''
<div style="height:100%;min-height:600px;display:flex;align-items:center;
	justify-content:center;border:1px dashed #c9c9c9;border-radius:8px;
	color:#999;font-family:sans-serif;text-align:center;">
	<div>
		<div style="font-size:42px;margin-bottom:8px;">&#127932;</div>
		<div>Sheet-music preview</div>
		<div style="font-size:12px;margin-top:4px;">
			(coming soon)
		</div>
	</div>
</div>
'''


CUSTOM_CSS = '''
/* Score List: truncate long file names to a single line with an ellipsis. */
.score-list label {
	max-width: 100%;
}
.score-list label > span {
	display: inline-block;
	max-width: 100%;
	overflow: hidden;
	text-overflow: ellipsis;
	white-space: nowrap;
	vertical-align: middle;
}
'''


def build_ui ():
	examples = load_library()

	with gr.Blocks(title='LilyScript') as demo:
		gr.Markdown('## 🎼 LilyScript — symbolic music generation with Lilylet')
		store = gr.State(examples)

		with gr.Row(equal_height=True):
			# ---------------- LEFT ----------------
			with gr.Column(scale=5):
				# (1) compose params, with (2) the collapsible run log stacked below
				with gr.Group():
					gr.Markdown('## Compose')
					with gr.Group():
						gr.Markdown('- Style Options')
						with gr.Row():
							composer = gr.Dropdown(label='composer', choices=COMPOSERS, value='',
								allow_custom_value=True)
							genre = gr.Dropdown(label='genre', choices=GENRES, value='',
								allow_custom_value=True)
							instrument = gr.Dropdown(label='instrument', choices=INSTRUMENTS, value='',
								allow_custom_value=True)
					prompt = gr.Textbox(label='Metadata prompt', lines=3, value='',
						placeholder='extra metadata lines, e.g.\n[key "C major"]\n(optional)')
					measures = gr.Number(label='Measures (0 = let model decide)', value=0, precision=0)
					with gr.Row():
						temperature = gr.Slider(0.0, 2.0, value=1.0, step=0.05, label='temperature')
						max_patches = gr.Number(label='max patches', value=1024, precision=0)
						seed = gr.Slider(0, 2147483647, value=42, step=1, label='seed')
					with gr.Row():
						gen_btn = gr.Button('Generate', variant='primary')
						stop_btn = gr.Button('Stop', variant='stop')

				with gr.Accordion('Logs', open=True):
					log = gr.Textbox(show_label=False, lines=10, max_lines=10,
						autoscroll=True, interactive=False, container=False)

				# bottom row: (3) file list | (4) editor
				with gr.Row(equal_height=True):
					with gr.Column(scale=2, min_width=160):
						with gr.Group():
							gr.Markdown('## Score List')
							file_list = gr.Radio(show_label=False, choices=list(examples.keys()),
								value=None, interactive=True, container=False,
								elem_classes=['score-list'])
					with gr.Column(scale=5):
						with gr.Group():
							gr.Markdown('## Lilylet editor')
							editor = gr.Code(show_label=False, language=None, lines=18,
								max_lines=18, interactive=True)

			# ---------------- RIGHT ----------------
			with gr.Column(scale=6):
				with gr.Group():
					gr.Markdown('## Sheet music')
					gr.HTML(SHEET_PLACEHOLDER)

		# ---- wiring ----
		# style dropdowns -> keep the metadata-prompt text box in sync
		for field in (composer, genre, instrument):
			field.change(sync_prompt, inputs=[composer, genre, instrument, prompt], outputs=[prompt])

		gen_event = gen_btn.click(
			run_generation,
			inputs=[prompt, measures, temperature, max_patches, seed, store],
			outputs=[log, editor, file_list, store],
		)
		stop_btn.click(None, None, None, cancels=[gen_event])
		file_list.select(load_file, inputs=[file_list, store], outputs=[editor])

	return demo


if __name__ == '__main__':
	build_ui().queue().launch(theme=gr.themes.Soft(), css=CUSTOM_CSS)