Transkun Model

.vscode/settings.json

@@ -0,0 +1,5 @@
+{
+    "python-envs.defaultEnvManager": "ms-python.python:conda",
+    "python-envs.defaultPackageManager": "ms-python.python:conda",
+    "python-envs.pythonProjects": []
+}

README.md

@@ -1,13 +1,13 @@
 ---
-title: Transkun
+title: Transkun Piano Transcription
 emoji: 🌍
 colorFrom: red
 colorTo: blue
 sdk: gradio
-sdk_version: 6.2.0
+sdk_version: 5.28.0
 app_file: app.py
 pinned: false
-short_description: Transkun Piano Transcription
+short_description: Transkun Piano Transcription by Yujia Yan and Zhiyao Duan
 ---
 
 Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

app.py

@@ -0,0 +1,96 @@
+import spaces
+from transkun.ModelTransformer import TransKun, writeMidi
+import numpy as np
+import torch
+from functools import lru_cache
+import librosa
+
+import pkg_resources
+import moduleconf
+
+from pyharp.core import ModelCard, build_endpoint
+from pyharp.media.audio import load_audio
+import gradio as gr
+
+model_card = ModelCard(
+    name="Transkun Piano Transcription",
+    decription=("Transcribes solo piano performance into MIDI notation"),
+    author="Yujia Yan, Zhiyao Duan",
+    tags=["transcription"]
+)
+
+@lru_cache(maxsize=2)
+def load_model(device: str):
+    defaultWeight =  (pkg_resources.resource_filename("transkun", "pretrained/2.0.pt"))
+    defaultConf =  (pkg_resources.resource_filename("transkun", "pretrained/2.0.conf"))
+
+    confManager = moduleconf.parseFromFile(defaultConf)
+    conf = confManager["Model"].config
+
+    checkpoint = torch.load(defaultWeight, map_location=device)
+    model = TransKun(conf=conf).to(device)
+
+    # Mirrors your checkpoint loading logic :contentReference[oaicite:3]{index=3}
+    if "best_state_dict" in checkpoint:
+        model.load_state_dict(checkpoint["best_state_dict"], strict=False)
+    else:
+        model.load_state_dict(checkpoint["state_dict"], strict=False)
+
+    model.eval()
+    return model
+
+@spaces.GPU
+def transcribe(input_file):
+    device = "cuda"
+    model = load_model(device)
+
+    signal = load_audio(input_file)
+    waveform = np.asarray(signal.audio_data)
+    sr = int(signal.sample_rate)
+
+    waveform = np.squeeze(waveform)
+
+    # If 2D, assume (channels, samples). Make mono.
+    if waveform.ndim == 2:
+        if waveform.shape[0] > 1:
+            waveform = waveform.mean(axis=0)
+        else:
+            waveform = waveform.reshape(-1)
+
+    if sr != model.fs:
+        waveform = librosa.resample(waveform.astype(np.float32), orig_sr=sr, target_sr=model.fs)
+        sr = model.fs
+
+    x = torch.from_numpy(waveform).to(device)
+
+    notesEst = model.transcribe(x, discardSecondHalf=False)
+
+    outputMidi = writeMidi(notesEst)
+    out_path = "transkun_out.mid"
+    outputMidi.write(out_path)
+    return out_path
+
+def process_fn(input_audio_path: str) -> str:
+    midi_path = transcribe(input_audio_path)
+    return midi_path
+
+with gr.Blocks() as demo:
+    input_audio = gr.Audio(label="Upload Solo Piano Audio", type="filepath").harp_required(True)
+    #output_midi = gr.File(label="Output MIDI File", file_types=[".mid"]).harp_required(True)
+
+    # ensure this is serialized as midi_track
+    output_midi = gr.File(
+        label="Output MIDI File",
+        file_types=[".mid", ".midi"],  
+        type="filepath"
+    ).harp_required(True)
+
+
+    app = build_endpoint(
+        model_card=model_card,
+        input_components=[input_audio],
+        output_components=[output_midi],
+        process_fn=process_fn
+    )
+
+demo.queue().launch(share=True, show_error=True)

requirements.txt

@@ -0,0 +1,3 @@
+git+https://github.com/TEAMuP-dev/pyharp.git@v0.3.0
+transkun
+librosa