app.py

import os
import warnings
import torch
import librosa
import huggingface_hub

# Apply monkey-patch to gradio_client to prevent the bool TypeError in JSON schema parsing
try:
    import gradio_client.utils
    _orig_get_type = gradio_client.utils.get_type
    def _patched_get_type(schema):
        if isinstance(schema, bool):
            schema = {}
        return _orig_get_type(schema)
    gradio_client.utils.get_type = _patched_get_type

    _orig_json_schema_to_python_type = gradio_client.utils._json_schema_to_python_type
    def _patched_json_schema_to_python_type(schema, defs):
        if isinstance(schema, bool):
            schema = {}
        return _orig_json_schema_to_python_type(schema, defs)
    gradio_client.utils._json_schema_to_python_type = _patched_json_schema_to_python_type
    print("Successfully monkey-patched gradio_client.utils to handle boolean schemas.")
except Exception as e:
    print(f"Failed to apply gradio_client monkey-patch: {e}")

import gradio as gr
from piano_transcription_inference import PianoTranscription, sample_rate

# Suppress warnings to keep logs clean
warnings.filterwarnings("ignore")

# Download the model weights eagerly
print("Downloading/fetching model weights...")
WEIGHTS_PATH = huggingface_hub.snapshot_download(
    "Genius-Society/piano_trans",
) + "/CRNN_note_F1=0.9677_pedal_F1=0.9186.pth"

# Instantiate transcriptor on CPU
print(f"Initializing PianoTranscription on CPU using weights from: {WEIGHTS_PATH}")
transcriptor = PianoTranscription(
    device="cpu",
    checkpoint_path=WEIGHTS_PATH,
)
print("PianoTranscription initialized successfully!")

def transcribe_audio_to_midi(audio_path: str):
    """
    Transcribes piano audio file to MIDI.
    """
    if not audio_path:
        return "Please upload an audio file first.", None

    try:
        print(f"Loading audio from {audio_path}")
        # Load audio at the required sampling rate mono
        audio, _ = librosa.load(audio_path, sr=sample_rate, mono=True)
        print(f"Audio loaded successfully. Duration: {len(audio) / sample_rate:.2f} seconds")

        # Generate output path
        output_dir = "output"
        os.makedirs(output_dir, exist_ok=True)
        
        # Get base name without extension
        base_name = os.path.splitext(os.path.basename(audio_path))[0]
        midi_path = os.path.join(output_dir, f"{base_name}.mid")
        
        print(f"Transcribing audio to MIDI...")
        transcriptor.transcribe(audio, midi_path)
        print(f"Transcription complete! Saved to {midi_path}")
        
        return "Transcription complete! You can download your MIDI file below.", midi_path

    except Exception as e:
        import traceback
        error_msg = f"Error during transcription: {str(e)}\n\n{traceback.format_exc()}"
        print(error_msg)
        return f"Error: {str(e)}", None

# Build clean Gradio interface
with gr.Blocks(theme=gr.themes.Soft(), css="#col-container { max-width: 800px; margin: 0 auto; }") as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("""
        # 🎹 High-Resolution Piano Transcription
        Convert audio recordings of piano performances into MIDI files with high accuracy. 
        Powered by ByteDance's High-Resolution Piano Transcription System of Qiuqiang Kong.
        """)
        
        with gr.Row():
            audio_input = gr.Audio(
                label="Upload Piano Audio (MP3, WAV, etc.)", 
                type="filepath"
            )
            
        with gr.Row():
            submit_btn = gr.Button("Transcribe to MIDI", variant="primary")
            
        with gr.Row():
            status_output = gr.Textbox(label="Status", interactive=False)
            
        with gr.Row():
            midi_output = gr.File(label="Download MIDI File")
            
        gr.HTML("""
        <div style="margin-top: 20px; padding: 15px; background-color: var(--block-background-fill); border-radius: 8px;">
            <h3>🎵 Recommended Playback</h3>
            <p>For the best MIDI playback experience with a virtual 3D piano, sheet music representation, and high-quality sound:</p>
            <ol>
                <li>Download your transcribed MIDI file using the download button above.</li>
                <li>Go to <a href="https://app.midiano.com/" target="_blank" style="color: var(--primary-500); text-decoration: underline;">MidiAno (app.midiano.com)</a>.</li>
                <li>Drag and drop your MIDI file into the MidiAno browser tab to watch/listen to your performance!</li>
            </ol>
        </div>
        """)

    submit_btn.click(
        fn=transcribe_audio_to_midi,
        inputs=audio_input,
        outputs=[status_output, midi_output]
    )

if __name__ == "__main__":
    demo.launch()