remove description markdown, add GPU support

app.py

@@ -30,7 +30,6 @@ def process_fn(xml_path: str | None, use_example: bool, seed: int) -> str:
 
 
 with gr.Blocks(title="DExter", delete_cache=(3600, 86400)) as demo:
-    gr.Markdown("Generate expressive piano MIDI from a MusicXML score. The official model samples 999 diffusion steps, so generation can take several minutes on CPU.")
     inputs = [
         gr.File(type="filepath", file_types=[".xml", ".musicxml"], label="Piano MusicXML score").set_info(
             "Uncompressed MusicXML, up to 2 MB and 400 notes. Upload takes priority over the example."

runtime.py

@@ -25,6 +25,7 @@ from renderer import Renderer
 CHECKPOINT = ROOT / "weights" / "last-v1.ckpt"
 EXAMPLE = ROOT / "examples" / "Schubert_D783_no15.musicxml"
 MAX_NOTES = 400
+DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
 
 @lru_cache(maxsize=1)
@@ -38,14 +39,13 @@ def load_model():
         cfg = compose(config_name="inference")
     args = OmegaConf.to_container(cfg.model.model.args, resolve=True)
     task = OmegaConf.to_container(cfg.task, resolve=True)
-    # Preserve the normalization statistics embedded in the pretrained checkpoint.
     for key in ("dataset_means", "dataset_stds"):
         if task.get(key) is None:
             task.pop(key)
     model = getattr(dexter_model, cfg.model.model.name).load_from_checkpoint(
-        checkpoint_path=str(CHECKPOINT), map_location="cpu", **args, **task
+        checkpoint_path=str(CHECKPOINT), map_location=DEVICE, **args, **task
     )
-    model.eval()
+    model.to(DEVICE).eval()
     return model, int(cfg.seg_len), int(cfg.overlap)
 
 
@@ -76,11 +76,11 @@ def render(score_path: str | Path, output_path: str | Path, seed: int = 13) -> P
         window = codec[start : start + seg_len]
         padded = np.pad(window, ((0, seg_len - len(window)), (0, 0)))
         segments.append(padded)
-    score_batch = torch.from_numpy(np.stack(segments))
+    score_batch = torch.from_numpy(np.stack(segments)).to(DEVICE)
     batch = {
-        "p_codec": torch.zeros(len(segments), seg_len, 5),
+        "p_codec": torch.zeros(len(segments), seg_len, 5, device=DEVICE),
         "s_codec": score_batch,
-        "c_codec": torch.zeros(len(segments), seg_len, 7),
+        "c_codec": torch.zeros(len(segments), seg_len, 7, device=DEVICE),
     }
 
     with torch.inference_mode():