martyyz-ai · moonwhaler · Aug 15, 2026
diff --git a/README.md b/README.md
@@ -61,19 +61,32 @@ The model weights are hosted on [HuggingFace](https://huggingface.co/MuScriptor)
 
 ### 🎛️ ComfyUI Node Documentation
 
-**Node Category**: `MuScriptor`  
-**Node Name**: `MuScriptor Transcribe Audio to MIDI`
+**Node Category**: `MuScriptor`
+
+#### `MuScriptor Load Model`
+
+Loads (and caches in RAM/VRAM) a MuScriptor model, either an auto-downloaded published variant or a local checkpoint of your own. Outputs a `MUSCRIPTOR_MODEL` that feeds into `MuScriptor Transcribe Audio to MIDI`. This decouples loading from transcription, so the same loaded model can feed multiple Transcribe nodes.
+
+| Parameter      | Type     | Description                                                                                                                                                                                   |
+| -------------- | -------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `model_source` | `ENUM`   | Either `auto (download): small/medium/large` (fetched from HuggingFace and cached), or a local `.safetensors` file placed under `ComfyUI/models/muscriptor/`.                               |
+| `device`       | `ENUM`   | Device target: `auto`, `cuda`, `cpu`, or `mps`.                                                                                                                                              |
+| `dtype`        | `ENUM`   | Computation precision: `auto`, `float32`, `float16`, `bfloat16`.                                                                                                                             |
+
+To use your own model: drop a `.safetensors` file into `ComfyUI/models/muscriptor/` (create the folder if it doesn't exist) and optionally place a matching `config.json` (with `dim`, `num_heads`, `num_layers`, `card`) next to it so the architecture is detected correctly. Restart ComfyUI (or refresh node definitions) so it shows up in the `model_source` dropdown.
+
+**Output**: `muscriptor_model` (`MUSCRIPTOR_MODEL`)
+
+#### `MuScriptor Transcribe Audio to MIDI`
 
 #### Inputs:
 
-| Parameter         | Type                  | Description                                                                              |
-| ----------------- | --------------------- | ---------------------------------------------------------------------------------------- |
-| `audio`           | `AUDIO` _(optional)_  | ComfyUI standard `AUDIO` port. Takes priority if connected.                              |
-| `audio_path`      | `STRING` _(optional)_ | Local file path to an audio file (WAV, MP3, FLAC, etc.).                                 |
-| `model_size`      | `ENUM`                | Model size: `small`, `medium` (default), or `large`.                                     |
-| `device`          | `ENUM`                | Device target: `auto`, `cuda`, `cpu`, or `mps`.                                          |
-| `dtype`           | `ENUM`                | Computation precision: `auto`, `float32`, `float16`, `bfloat16`.                         |
-| `use_sampling`    | `BOOLEAN`             | Use temperature sampling instead of greedy decoding (default: `False`).                  |
+| Parameter          | Type                            | Description                                                                              |
+| ------------------ | -------------------------------- | ---------------------------------------------------------------------------------------- |
+| `muscriptor_model`  | `MUSCRIPTOR_MODEL`               | Model loaded from a `MuScriptor Load Model` node.                                        |
+| `audio`             | `AUDIO` _(optional)_             | ComfyUI standard `AUDIO` port. Takes priority if connected.                              |
+| `audio_path`        | `STRING` _(optional)_            | Local file path to an audio file (WAV, MP3, FLAC, etc.).                                 |
+| `use_sampling`      | `BOOLEAN`                        | Use temperature sampling instead of greedy decoding (default: `False`).                  |
 | `temperature`     | `FLOAT`               | Sampling temperature (0.0 to 2.0).                                                       |
 | `cfg_coef`        | `FLOAT`               | Classifier-free guidance coefficient (default: `1.0`).                                   |
 | `beam_size`       | `INT`                 | Beam search width (`1` = greedy/sampling, `≥2` = beam search).                           |
@@ -146,19 +159,32 @@ Os pesos dos modelos estão hospedados no [HuggingFace](https://huggingface.co/M
 
 ### 🎛️ Documentação do Nó no ComfyUI
 
-**Categoria do Nó**: `MuScriptor`  
-**Nome do Nó**: `MuScriptor Transcribe Audio to MIDI`
+**Categoria do Nó**: `MuScriptor`
+
+#### `MuScriptor Load Model`
+
+Carrega (e mantém em cache na RAM/VRAM) um modelo MuScriptor, seja uma variante publicada com download automático, seja um checkpoint local seu. Gera um `MUSCRIPTOR_MODEL` que alimenta o nó `MuScriptor Transcribe Audio to MIDI`. Isso separa o carregamento da transcrição, permitindo que o mesmo modelo carregado alimente vários nós Transcribe.
+
+| Parâmetro      | Tipo   | Descrição                                                                                                                                                              |
+| -------------- | ------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------- |
+| `model_source` | `ENUM` | `auto (download): small/medium/large` (baixado do HuggingFace e armazenado em cache) ou um arquivo `.safetensors` local em `ComfyUI/models/muscriptor/`.             |
+| `device`       | `ENUM` | Dispositivo de execução: `auto`, `cuda`, `cpu` ou `mps`.                                                                                                             |
+| `dtype`        | `ENUM` | Precisão numérica: `auto`, `float32`, `float16`, `bfloat16`.                                                                                                         |
+
+Para usar seu próprio modelo: coloque um arquivo `.safetensors` em `ComfyUI/models/muscriptor/` (crie a pasta se não existir) e, opcionalmente, um `config.json` correspondente (com `dim`, `num_heads`, `num_layers`, `card`) ao lado dele para que a arquitetura seja detectada corretamente. Reinicie o ComfyUI (ou atualize as definições de nós) para que apareça no menu `model_source`.
+
+**Saída**: `muscriptor_model` (`MUSCRIPTOR_MODEL`)
+
+#### `MuScriptor Transcribe Audio to MIDI`
 
 #### Entradas (Inputs):
 
-| Parâmetro         | Tipo                  | Descrição                                                                             |
-| ----------------- | --------------------- | ------------------------------------------------------------------------------------- |
-| `audio`           | `AUDIO` _(opcional)_  | Porta padrão de áudio do ComfyUI. Tem prioridade se conectada.                        |
-| `audio_path`      | `STRING` _(opcional)_ | Caminho local para um arquivo de áudio (WAV, MP3, FLAC, etc.).                        |
-| `model_size`      | `ENUM`                | Tamanho do modelo: `small`, `medium` (padrão) ou `large`.                             |
-| `device`          | `ENUM`                | Dispositivo de execução: `auto`, `cuda`, `cpu` ou `mps`.                              |
-| `dtype`           | `ENUM`                | Precisão numérica: `auto`, `float32`, `float16`, `bfloat16`.                          |
-| `use_sampling`    | `BOOLEAN`             | Usar amostragem por temperatura em vez de decodificação gulosa (padrão: `False`).     |
+| Parâmetro          | Tipo                   | Descrição                                                                             |
+| ------------------- | ---------------------- | ------------------------------------------------------------------------------------- |
+| `muscriptor_model`  | `MUSCRIPTOR_MODEL`      | Modelo carregado por um nó `MuScriptor Load Model`.                                   |
+| `audio`             | `AUDIO` _(opcional)_    | Porta padrão de áudio do ComfyUI. Tem prioridade se conectada.                        |
+| `audio_path`        | `STRING` _(opcional)_   | Caminho local para um arquivo de áudio (WAV, MP3, FLAC, etc.).                        |
+| `use_sampling`      | `BOOLEAN`               | Usar amostragem por temperatura em vez de decodificação gulosa (padrão: `False`).     |
 | `temperature`     | `FLOAT`               | Temperatura da amostragem (0.0 a 2.0).                                                |
 | `cfg_coef`        | `FLOAT`               | Coeficiente do Classifier-Free Guidance (padrão: `1.0`).                              |
 | `beam_size`       | `INT`                 | Largura da busca por feixe (Beam Search) (`1` = guloso, `≥2` = busca por feixe).      |

diff --git a/comfy_node.py b/comfy_node.py
@@ -9,13 +9,58 @@
 from muscriptor.events import NoteStartEvent, NoteEndEvent, ProgressEvent
 from muscriptor.tokenizer.mt3 import resolve_instrument_names
 
+# Local model files live under ComfyUI/models/muscriptor/ (a dedicated folder,
+# since MuScriptor checkpoints use a bespoke architecture incompatible with
+# the standard checkpoint/unet loaders). Registered here so get_filename_list
+# / get_full_path can find them like any other ComfyUI model folder.
+try:
+    import folder_paths
+    if "muscriptor" not in folder_paths.folder_names_and_paths:
+        _muscriptor_models_dir = os.path.join(folder_paths.models_dir, "muscriptor")
+        os.makedirs(_muscriptor_models_dir, exist_ok=True)
+        folder_paths.folder_names_and_paths["muscriptor"] = (
+            [_muscriptor_models_dir],
+            {".safetensors"},
+        )
+except ImportError:
+    folder_paths = None
+
+# Bare size keywords resolve to the published HuggingFace variant and are
+# downloaded/cached automatically (see TranscriptionModel.load_model).
+_AUTO_DOWNLOAD_SOURCES = {
+    "auto (download): small": "small",
+    "auto (download): medium": "medium",
+    "auto (download): large": "large",
+}
+
+
+def _list_model_sources() -> list[str]:
+    """Combo box choices: auto-download variants plus local files found in
+    ComfyUI/models/muscriptor/."""
+    sources = list(_AUTO_DOWNLOAD_SOURCES.keys())
+    if folder_paths is not None:
+        sources += folder_paths.get_filename_list("muscriptor")
+    return sources
+
+
+def _resolve_weights_path(model_source: str) -> str:
+    """Map a combo box choice to what TranscriptionModel.load_model expects."""
+    if model_source in _AUTO_DOWNLOAD_SOURCES:
+        return _AUTO_DOWNLOAD_SOURCES[model_source]
+    if folder_paths is not None:
+        full_path = folder_paths.get_full_path("muscriptor", model_source)
+        if full_path is not None:
+            return full_path
+    return model_source
+
+
 # Cache dictionary to keep loaded model instances in memory
 _MODEL_CACHE = {}
 
-def get_model(model_size: str, device: str, dtype: str) -> TranscriptionModel:
+def load_muscriptor_model(weights_path: str, device: str, dtype: str) -> TranscriptionModel:
     """Helper to load and cache the transcription model."""
     import muscriptor.accelerator
-    
+
     # Resolve torch device
     if device == "auto":
         model_device = (
@@ -25,27 +70,27 @@ def get_model(model_size: str, device: str, dtype: str) -> TranscriptionModel:
         )
     else:
         model_device = torch.device(device)
-        
+
     # Resolve dtype
     if dtype == "auto":
         model_dtype = torch.float16 if model_device.type == "mps" else torch.float32
     else:
         model_dtype = getattr(torch, dtype)
 
-    cache_key = (model_size, str(model_device), str(model_dtype))
-    
+    cache_key = (weights_path, str(model_device), str(model_dtype))
+
     if cache_key not in _MODEL_CACHE:
         # Clear other cached models to avoid CUDA OOM / memory bloat
         _MODEL_CACHE.clear()
-        
-        print(f"[MuScriptor] Loading model (size={model_size}, device={model_device}, dtype={model_dtype})...")
+
+        print(f"[MuScriptor] Loading model (source={weights_path}, device={model_device}, dtype={model_dtype})...")
         model = TranscriptionModel.load_model(
-            weights_path=model_size,
+            weights_path=weights_path,
             device=model_device,
             dtype=model_dtype
         )
         _MODEL_CACHE[cache_key] = model
-        
+
     return _MODEL_CACHE[cache_key]
 
 
@@ -73,14 +118,34 @@ def events_to_notes_list(events) -> list[dict]:
     return notes
 
 
-class MuScriptorTranscribe:
+class MuScriptorLoadModel:
     @classmethod
     def INPUT_TYPES(cls):
         return {
             "required": {
-                "model_size": (["small", "medium", "large"], {"default": "medium"}),
+                "model_source": (_list_model_sources(), {"default": "auto (download): medium"}),
                 "device": (["auto", "cpu", "cuda", "mps"], {"default": "auto"}),
                 "dtype": (["auto", "float32", "float16", "bfloat16"], {"default": "auto"}),
+            }
+        }
+
+    RETURN_TYPES = ("MUSCRIPTOR_MODEL",)
+    RETURN_NAMES = ("muscriptor_model",)
+    FUNCTION = "load_model"
+    CATEGORY = "MuScriptor"
+
+    def load_model(self, model_source, device, dtype):
+        weights_path = _resolve_weights_path(model_source)
+        model = load_muscriptor_model(weights_path, device, dtype)
+        return (model,)
+
+
+class MuScriptorTranscribe:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "muscriptor_model": ("MUSCRIPTOR_MODEL",),
                 "use_sampling": ("BOOLEAN", {"default": False}),
                 "temperature": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0, "step": 0.05}),
                 "cfg_coef": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}),
@@ -103,9 +168,7 @@ def INPUT_TYPES(cls):
 
     def transcribe_audio(
         self,
-        model_size,
-        device,
-        dtype,
+        muscriptor_model,
         use_sampling,
         temperature,
         cfg_coef,
@@ -150,8 +213,8 @@ def transcribe_audio(
                 except ValueError as e:
                     raise ValueError(f"[MuScriptor] Error resolving instrument names: {e}")
 
-        # 3. Load Model
-        model = get_model(model_size, device, dtype)
+        # 3. Model comes preloaded from MuScriptorLoadModel
+        model = muscriptor_model
 
         # 4. Perform Transcription
         print("[MuScriptor] Starting music transcription...")
@@ -202,9 +265,11 @@ def transcribe_audio(
 
 
 NODE_CLASS_MAPPINGS = {
+    "MuScriptorLoadModel": MuScriptorLoadModel,
     "MuScriptorTranscribe": MuScriptorTranscribe,
 }
 
 NODE_DISPLAY_NAME_MAPPINGS = {
+    "MuScriptorLoadModel": "MuScriptor Load Model",
     "MuScriptorTranscribe": "MuScriptor Transcribe Audio to MIDI",
 }
diff --git a/example/MuScriptor Example.json b/example/MuScriptor Example.json
@@ -1,8 +1,8 @@
 {
   "id": "e803a4eb-858b-46b8-93e4-d37300742aba",
   "revision": 0,
-  "last_node_id": 4,
-  "last_link_id": 4,
+  "last_node_id": 5,
+  "last_link_id": 5,
   "nodes": [
     {
       "id": 2,
@@ -56,7 +56,7 @@
         88
       ],
       "flags": {},
-      "order": 2,
+      "order": 3,
       "mode": 0,
       "inputs": [
         {
@@ -99,7 +99,7 @@
         88
       ],
       "flags": {},
-      "order": 3,
+      "order": 4,
       "mode": 0,
       "inputs": [
         {
@@ -140,9 +140,14 @@
         372
       ],
       "flags": {},
-      "order": 1,
+      "order": 2,
       "mode": 0,
       "inputs": [
+        {
+          "name": "muscriptor_model",
+          "type": "MUSCRIPTOR_MODEL",
+          "link": 5
+        },
         {
           "name": "audio",
           "shape": 7,
@@ -177,9 +182,6 @@
         }
       },
       "widgets_values": [
-        "large",
-        "cuda",
-        "auto",
         false,
         1,
         1,
@@ -190,6 +192,46 @@
         "",
         "drums"
       ]
+    },
+    {
+      "id": 5,
+      "type": "MuScriptorLoadModel",
+      "pos": [
+        1366.2126923678743,
+        150
+      ],
+      "size": [
+        270,
+        106
+      ],
+      "flags": {},
+      "order": 1,
+      "mode": 0,
+      "inputs": [],
+      "outputs": [
+        {
+          "name": "muscriptor_model",
+          "type": "MUSCRIPTOR_MODEL",
+          "links": [
+            5
+          ]
+        }
+      ],
+      "properties": {
+        "cnr_id": "comfyui_fearnworksnodes",
+        "ver": "0.1.2",
+        "Node name for S&R": "MuScriptorLoadModel",
+        "ue_properties": {
+          "widget_ue_connectable": {},
+          "input_ue_unconnectable": {},
+          "version": "7.8"
+        }
+      },
+      "widgets_values": [
+        "auto (download): large",
+        "cuda",
+        "auto"
+      ]
     }
   ],
   "links": [
@@ -216,6 +258,14 @@
       4,
       0,
       "STRING"
+    ],
+    [
+      5,
+      5,
+      0,
+      1,
+      0,
+      "MUSCRIPTOR_MODEL"
     ]
   ],
   "groups": [],