diff --git a/README.md b/README.md index a18e6a0..1cc529d 100644 --- 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 index fa73e87..97aa82c 100644 --- 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 index fe13bd6..ad6034b 100644 --- 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": [],