[Added] Demucs support and node

README.md

@@ -14,12 +14,14 @@ audio demixing, also known as audio separation.
 From an audio the objective is to separate the vocals, instruments, drums, bass, etc.
 from the rest of the sounds.
 
-To achieve this we use [MDX Net](https://arxiv.org/abs/2111.12203) neural networks (models).
-**AudioSeparation** currently supports [39 models](https://huggingface.co/set-soft/audio_separation)
-collected by the [UVR5](https://github.com/Anjok07/ultimatevocalremovergui) project.
+To achieve this we use [MDX Net](https://arxiv.org/abs/2111.12203) and
+[Demucs](https://github.com/facebookresearch/demucs) neural networks (models).
+**AudioSeparation** currently supports [46 models](https://huggingface.co/set-soft/audio_separation)
+mostly collected by the [UVR5](https://github.com/Anjok07/ultimatevocalremovergui) project.
 
-The models are small (from 21 MB to 65 MB), but really efficient.
-Models specialized on different stems are provided.
+The MDX models are small (from 21 MB to 65 MB), but really efficient, the Demucs models are bigger
+(from 84 MB to 870) MB, but slightly better, and supports 4 stems.
+MDX models specialized on different stems are provided.
 We support more than one model for each task because some times a model will perform better
 for a song and worst for others.
 
@@ -32,7 +34,7 @@ The objectives for these nodes are:
 - Easy of use
 - Clear download (with progress and known destination)
 - Support for all possible input audio formats (mono/stereo, any sample rate, any batch size)
-- Good quality vs size
+- Good quality vs size, or you can choose better quality using Demucs
 - Reduced dependencies
 
 ---
@@ -50,6 +52,7 @@ The objectives for these nodes are:
     *   [Command Line](#command-line)
 *   ✨ [Nodes](#-nodes)
     *   [Vocals using MDX](#vocals-using-mdx)
+    *   [Demucs Audio Separator](#demucs-audio-separator)
 *   📝 [Usage Notes](#-usage-notes)
 *   ⚖️ [License](#️-license)
 *   🙏 [Attributions](#-attributions)
@@ -111,7 +114,7 @@ A list of all the available tools can be found [here](tool/README.md).
 
 Models are automatically downloaded.
 
-When using ComfyUI they are downloaded to `ComfyUI/models/audio/MDX`.
+When using ComfyUI they are downloaded to `ComfyUI/models/audio/MDX` and `ComfyUI/models/audio/Demucs`.
 
 When using the command line the default is `../models` relative to the script, but you can specify another dir.
 
@@ -123,7 +126,7 @@ For the command line you can also download the ONNX files from other repos.
 
 ## 📦 Dependencies
 
-These nodes just uses `torchaudio` (part of PyTorch), `numpy` for math and `tqdm` for progress bars.
+These nodes just uses `torchaudio` (part of PyTorch), `numpy` for math, `safetensors` to load models and `tqdm` for progress bars.
 All of them are used by ComfyUI, so you don't need to install any additional dependency on a ComfyUI setup.
 
 The following are optional dependencies:
@@ -141,13 +144,14 @@ You can start using template workflows, go to the ComfyUI *Workflow* menu and th
 look for *Audio Separation*
 
 If you want to do it manually you'll find the nodes in the *audio/separation* category.
-Or you can use the search menu, double click in the canvas and then type **MDX**:
+Or you can use the search menu, double click in the canvas and then type **MDX** (or **Demucs**):
 
 ![ComfyUI Search](doc/node_search.png)
 
 Choose a node to extract what you want, i.e. *Vocals*. The complement output for it will
 be the instruments, but using a node for *Instrumental* separation you'll usually get a better result
-than using the *Complement*  output.
+than using the *Complement* output. In the case of *Demucs* models you get 4 or 6 stems at a time,
+the "UVR Demucs" is an exception, it just supports Vocals and Other.
 
 Then simply connect your audio input to the node (i.e. **LoadAudio** node from Comfy core) and
 connect its output to some audio node (i.e. **PreviewAudio** or **SaveAudio** nodes from Comfy core).
@@ -207,6 +211,44 @@ share the same structure, so here is the first:
      - **Input Batch Handling:** If `input_sound` is a batch the outputs will be batches. The process is sequential, not parallel.
      - **Missing Models:** They are downloaded and stored under `models/audio/MDX` of the ComfyUI installation
 
+And here is the Demucs node:
+
+### Demucs Audio Separator
+   - **Display Name:** `Demucs Audio Separator`
+   - **Internal Name:** `AudioSeparateDemucs`
+   - **Category:** `audio/separation`
+   - **Description:** Takes one audio input (which can be a batch) separates the vocals, drums and bass from the rest of the sounds.
+     The node has outputs for guitar and piano, which can be separated by the *Hybrid Transformer 6 sources* model, which is quite
+     experimental.
+   - **Inputs:**
+     - `input_sound` (AUDIO): The audio input. Can be a single audio item or a batch.
+     - `model` (COMBO): The name of the model to use. Choose one from the list.
+     - `shifts` (INT): Number of random shifts for equivariant stabilization.
+        It does extra passes using slightly shifted audio, which can produce better results.
+        Higher values improve quality but are slower. 0 disables it.
+     - `overlap` (FLOAT): Amount of overlap between audio chunks.
+        This is expressed as a portion of the total chunk, i.e. 0.25 is 25%.
+        Higher values can reduce stitching artifacts but are slower.
+     - `custom_segment` (BOOLEAN): Enable to override the model's default segment length.
+        Disabling uses the recommended length from the model file.
+        Useful for HDemucs and Demucs models, not much for HTDemucs.
+     - segment (INT): Length of audio chunks to process at a time (in seconds).
+        Higher values need more VRAM but can improve quality.
+     - `taget_device` (COMBO): The device where we will run the neural network.
+   - **Output:**
+     - `Vocals` (AUDIO): The separated vocals
+     - `Drums` (AUDIO): The separated drums. Not for "UVR" version.
+     - `Bass` (AUDIO): The separated bass. Not for "UVR" version.
+     - `Other` (AUDIO): The separated stuff that doesn't fit in the other outputs
+     - `Guitar` (AUDIO): The separated guitar, only for *Hybrid Transformer 6 sources* model, which is quite
+     - `Piano` (AUDIO): The separated piano, only for *Hybrid Transformer 6 sources* model, which is quite
+   - **Behavior Details:**
+     - **Sample Rate:** The sample rate of the input is adjusted to 44.1 kHz
+     - **Channels:** Mono audios are converted to fake stereo (left == right)
+     - **Input Batch Handling:** If `input_sound` is a batch the outputs will be batches.
+     - **Missing Models:** They are downloaded and stored under `models/audio/Demucs` of the ComfyUI installation
+     - **Models:** Note that most models are a *bag of models*, this is four models working together.
+
 
 ## 📝 Usage Notes
 
@@ -217,6 +259,8 @@ share the same structure, so here is the first:
   You can control log verbosity through ComfyUI's startup arguments (e.g., `--preview-method auto --verbose DEBUG` for more detailed ComfyUI logs
   which might also affect custom node loggers if they are configured to inherit levels). The logger name used is "AudioSeparation".
   You can force debugging level for these nodes defining the `AUDIOSEPARATION_NODES_DEBUG` environment variable to `1`.
+- **Models format:** We use safetensors because this format is safer than PyTorch files (.pth, .th, etc.) and doesn't need an extra runtime (like ONNX does)
+- **No quantized Demucs:** These models just save download time, but pulls extra dependency (diffq), they are just lower quality versions of their non-quantized counterparts.
 
 
 ## ⚖️ License
@@ -237,6 +281,7 @@ ______
 - Models collected by the [UVR5 project](https://github.com/Anjok07/ultimatevocalremovergui) and
   found in the [UVR Resources](https://huggingface.co/Politrees/UVR_resources) by
   [Artyom Bebroy](https://github.com/Politrees)
+- Demucs models are from Meta Platforms, Inc. Except for the "UVR" version
 - The logo image was created using text generated using [Text Studio](https://www.textstudio.com/) and
   resources from [Vecteezy](https://www.vecteezy.com/) by:
      - [Titima Ongkantong](https://www.vecteezy.com/members/titima157)

nodes.py

@@ -3,23 +3,23 @@
 # License: GPLv3
 # Project: ComfyUI-AudioSeparation
 import os
-import torch
 from typing import Dict
 # ComfyUI imports
 import folder_paths  # ComfyUI's way to access model paths
 # Local imports
 from .source.utils.logger import main_logger
 from .source.utils.load_audio import audio_get_channels, force_stereo, force_sample_rate
-from .source.utils.torch import get_torch_device_options
+from .source.utils.torch import get_torch_device_options, get_canonical_device
 from .source.utils.comfy_node_action import send_node_action
-from .source.inference.demixer import get_demixer
 from .source.db.models_db import ModelsDB
-
+from .source.inference.demixer import get_demixer
 
 DEF_MODEL = 'Kim_Vocal_2.safetensors'
 DEF_ENTRY = 'Default'
 MODELS_DIR = os.path.join(folder_paths.models_dir, "audio", "MDX")
-models_db = ModelsDB(MODELS_DIR)
+DEMUCS_DIR = os.path.join(folder_paths.models_dir, "audio", "Demucs")
+models_db_mdx = ModelsDB(MODELS_DIR)
+models_db_demucs = ModelsDB(DEMUCS_DIR)
 logger = main_logger
 
 
@@ -28,15 +28,15 @@ class AudioSeparateVocals:
     MODEL_T = 'MDX'
     FILE_T = 'safetensors'
     DEFAULT_MODEL = "Kim_Vocal_2.safetensors"
+    models_db = models_db_mdx
 
     @classmethod
     def _get_available_audio_models(cls):
-        global models_db
         # Refresh the database
-        models_db.refresh()
+        cls.models_db.refresh()
         # Filter the models this node can handle
-        cls.models_filtered = models_db.get_filtered(primary_stem=cls.PRIMARY_STEM, model_t=cls.MODEL_T, file_t=cls.FILE_T,
-                                                     default=cls.DEFAULT_MODEL, repeat_dl=True)
+        cls.models_filtered = cls.models_db.get_filtered(primary_stem=cls.PRIMARY_STEM, model_t=cls.MODEL_T, file_t=cls.FILE_T,
+                                                         default=cls.DEFAULT_MODEL, repeat_dl=True)
         # We add any model downloaded and memorized by the GUI
         return cls.models_filtered.get_display_names()
 
@@ -81,12 +81,15 @@ class AudioSeparateVocals:
         model_path = model_data.get('model_path')
 
         # Create or recycle a demixer
-        device = torch.device(target_device)
+        device = get_canonical_device(target_device)
         if self.demixer is None or self.demixer.d['hash'] != model_data['hash']:
             # New demixer
             logger.debug("Creating a new demixer object")
             # This will load the model, optionally downloading it
             self.demixer = get_demixer(model_data, device, MODELS_DIR)
+        else:
+            # Update the device
+            self.demixer.set_device(device)
 
         # Handle a change in the icon of the model name
         if model_path is None:
@@ -141,3 +144,93 @@ class AudioSeparateVarious(AudioSeparateVocals):
     UNIQUE_NAME = "AudioSeparateVarious"
     DISPLAY_NAME = "Various using MDX"
     RETURN_NAMES = ("Main", "Complement",)
+
+
+class AudioSeparateDemucs(AudioSeparateVocals):
+    PRIMARY_STEM = None
+    MODEL_T = 'Demucs'
+    FILE_T = 'safetensors'
+    DEFAULT_MODEL = "htdemucs_ft.safetensors"
+    models_db = models_db_demucs
+
+    @classmethod
+    def INPUT_TYPES(cls):
+        device_options, default_device = get_torch_device_options()
+        return {
+            "required": {
+                "input_sound": ("AUDIO",),
+                "model": (cls._get_available_audio_models(),),  # Dropdown for model selection
+                "shifts": ("INT", {
+                    "default": 0, "min": 0, "max": 16, "step": 1, "display": "slider",
+                    "tooltip": "Number of random shifts for equivariant stabilization.\n"
+                               "Higher values improve quality but are slower. 0 disables it."
+                }),
+                "overlap": ("FLOAT", {
+                    "default": 0.25, "min": 0.0, "max": 0.99, "step": 0.01, "display": "slider",
+                    "tooltip": "Amount of overlap between audio chunks.\n"
+                               "Higher values can reduce stitching artifacts but are slower."
+                }),
+                "custom_segment": ("BOOLEAN", {
+                    "default": False,
+                    "label_on": "enabled",
+                    "label_off": "disabled",
+                    "tooltip": "Enable to override the model's default segment length.\n"
+                               "Disabling uses the recommended length from the model file.\n"
+                               "Useful for HDemucs and Demucs models, not much for HTDemucs."
+                }),
+                "segment": ("INT", {
+                    "default": 44, "min": 10, "max": 120, "step": 1, "display": "slider",
+                    "tooltip": "Length of audio chunks to process at a time (in seconds).\n"
+                               "Higher values need more VRAM but can improve quality."
+                }),
+                "target_device": (device_options, {
+                    "default": default_device,
+                    "tooltip": "The device (CPU or CUDA) to which the projection layer will be assigned for computation."}),
+            }
+        }
+
+    # Define the output names. We define all 6 possible stems.
+    # The execution logic will handle returning 'None' for unused outputs.
+    RETURN_NAMES = ("vocals", "drums", "bass", "other", "guitar", "piano")
+    RETURN_TYPES = ("AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO", "AUDIO")
+    DESCRIPTION = "Demucs Audio Separator (4/6 stems)"
+    UNIQUE_NAME = "AudioSeparateDemucs"
+    DISPLAY_NAME = "Demucs Audio Separator"
+
+    def execute(self, input_sound: Dict, model: str, shifts: int, overlap: float, custom_segment: bool, segment: int,
+                target_device: str):
+        # Get information for the selected model
+        main_logger.info(f"Selected model: {model}")
+        model_data = self.models_filtered.get_by_display_name(model)
+        if model_data is None:
+            raise ValueError("Unknown model selected, please refresh pressing `R` and select another")
+        model_path = model_data.get('model_path')
+
+        # Create or recycle a demixer
+        device = get_canonical_device(target_device)
+        if self.demixer is None or self.demixer.d['hash'] != model_data['hash']:
+            # New demixer
+            logger.debug("Creating a new demixer object")
+            # This will load the model, optionally downloading it
+            self.demixer = get_demixer(model_data, device, DEMUCS_DIR)
+        else:
+            # Update the device
+            self.demixer.set_device(device)
+
+        # Handle a change in the icon of the model name
+        if model_path is None:
+            # Was downloaded
+            send_node_action("change_widget", "model", model_data['indicator'] + model_data['filtered_name'])
+
+        # Match channels and S/R
+        waveform = input_sound['waveform']
+        sample_rate = input_sound['sample_rate']
+        if audio_get_channels(waveform) == 1 and self.demixer.ch == 2:
+            waveform = force_stereo(waveform)
+        if sample_rate != self.demixer.sr:
+            waveform = force_sample_rate(waveform, sample_rate, self.demixer.sr)
+
+        # Demix
+        wavs = self.demixer(waveform, shifts=shifts, overlap=overlap, segment=segment if custom_segment else None)
+
+        return tuple(wavs)

source/db/models_db.py

@@ -22,7 +22,8 @@ known_models_mtime = None
 ICON_REMOTE = "⬇️  "         # "\u2B07"      # ⬇️
 ICON_DOWNLOADED = "\U0001F4BE "  # 💾 Floppy Disk
 KNOWN_SOURCES = {'Politrees/MDXNet': 'https://huggingface.co/Politrees/UVR_resources/resolve/main/models/MDXNet',
-                 'Main/MDX': 'https://huggingface.co/set-soft/audio_separation/resolve/main/MDX'}
+                 'Main/MDX': 'https://huggingface.co/set-soft/audio_separation/resolve/main/MDX',
+                 'Main/Demucs': 'https://huggingface.co/set-soft/audio_separation/resolve/main/Demucs', }
 
 
 def get_db_filename(provided=None):
@@ -68,9 +69,10 @@ def load_known_models(json_path=None):
         logger.error("Error: The models database was not found at the expected location.")
         logger.error("Please check the directory structure.")
         return None
-    except json.JSONDecodeError:
-        logger.error(f"Error: The file at '{json_path}' is not a valid JSON file.")
-        return None
+    except json.JSONDecodeError as e:
+        msg = f"Error: The file at '{json_path}' is not a valid JSON file: {e}"
+        logger.error(msg)
+        raise ValueError(msg)
     except Exception as e:
         logger.error(f"An unexpected error occurred: {e}")
         return None
@@ -125,6 +127,8 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
     # Allow for multiple values in the filters
     if isinstance(primary_stem, str):
         primary_stem = {primary_stem}
+    elif isinstance(primary_stem, list):
+        primary_stem = set(primary_stem)
     if isinstance(model_t, str):
         model_t = {model_t}
     if isinstance(file_t, str):
@@ -153,8 +157,14 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
             continue
         # Check the stem
         try:
-            if primary_stem is not None and d['primary_stem'] not in primary_stem:
-                continue
+            if primary_stem is not None:
+                mps = d['primary_stem']
+                if isinstance(mps, str):
+                    if mps not in primary_stem:
+                        continue
+                else:  # A list
+                    if not (set(mps) & primary_stem):
+                        continue
         except KeyError:
             logger.error(f"Missing `primary_stem` for {name}")
             continue
@@ -201,6 +211,7 @@ def get_models_full(primary_stem=None, model_t=None, file_t=None, json_path=None
                     on_disk_as_down.append(ICON_REMOTE + filtered_name)
             else:
                 to_down.append(ICON_REMOTE + filtered_name)
+        d['hash'] = hash
         found[filtered_name] = d
         found_hashes[hash] = d
         if file_name is not None:
@@ -228,14 +239,19 @@ def get_download_url(data):
         return None
 
 
+def cli_add_db(parser, default_json_file=None):
+    default_json_file = default_json_file or get_db_filename()
+    parser.add_argument('--json_file', type=str, default=default_json_file,
+                        help="Path to the models database JSON file.")
+
+
 def cli_add_models_and_db(parser):
     # Compute the models dir assuming the script is run from a clone of the repo
     default_json_file = get_db_filename()
 
     parser.add_argument('--models_dir', type=str, default=os.path.dirname(default_json_file),
                         help="Path to the directory containing model files.")
-    parser.add_argument('--json_file', type=str, default=default_json_file,
-                        help="Path to the models database JSON file.")
+    cli_add_db(parser, default_json_file=default_json_file)
 
 
 def download_model(data, models_dir):
@@ -305,13 +321,42 @@ class ModelsDB(object):
     def __init__(self, models_dir: str, json_path: str = None):
         super().__init__()
         self.models_dir = models_dir
-        self.json_path = json_path
+        self.json_path = get_db_filename(json_path)
         self.refresh()
 
     def refresh(self):
         self.downloaded = hash_dir(self.models_dir)
         self.models = load_known_models(self.json_path)
 
+    def remove(self, data):
+        if data is None:
+            return
+        del self.models[data["hash"]]
+
+    def add(self, hash, data):
+        self.models[hash] = data
+
+    def save(self):
+        # Remove run-time information
+        for k, v in self.models.items():
+            try:
+                del v["hash"]
+            except KeyError:
+                pass
+            try:
+                del v["indicator"]
+            except KeyError:
+                pass
+            try:
+                del v["filtered_name"]
+            except KeyError:
+                pass
+            try:
+                del v["model_path"]
+            except KeyError:
+                pass
+        save_known_models(self.models, self.json_path)
+
     def get_filtered(self, primary_stem=None, model_t=None, file_t=None, default=None, repeat_dl=False):
         return FilteredModels(primary_stem=primary_stem, model_t=model_t, file_t=file_t, json_path=self.json_path,
                               downloaded=self.downloaded, default=default, repeat_dl=repeat_dl)

source/inference/demixer.py

@@ -5,6 +5,7 @@
 #
 # Wrappers for the model and inference
 import logging
+import math
 import torch
 # ComfyUI imports
 try:
@@ -16,22 +17,42 @@ except Exception:
 from .stft import stft_chunk_process, stft_get_chunks
 from ..db.load_model import load_model
 from ..utils.misc import NODES_NAME
+from ..utils.torch import model_to_target, get_offload_device
+from .demucs_api import apply_model, BagOfModels
 
 logger = logging.getLogger(f"{NODES_NAME}.demixer")
 SAMPLE_RATE = 44100
 
 
+class DemixerGeneric(object):
+    def __init__(self, d, device, models_dir):
+        super().__init__()
+        self.d = d
+        self.model_run = load_model(d, device, models_dir)
+        self.device = device
+
+    def set_device(self, device):
+        if device == self.device:
+            return
+        self.device = device
+        if hasattr(self.model_run, "target_device"):
+            self.model_run = device
+
+
+# ############################################################################################################################
+# MDX-Net
+# ############################################################################################################################
+
+
 def show_inference_parameters(d):
     logger.debug("Using inference parameters:")
     logger.debug(f"  Frequency Bins (n_fft/2): {d['mdx_n_fft_scale_set']//2}")
     logger.debug(f"  Amplitude Compensation: {d['compensate']}")
 
 
-class DemixerMDX(object):
+class DemixerMDX(DemixerGeneric):
     def __init__(self, d, device, models_dir):
-        self.d = d
-        self.model_run = load_model(d, device, models_dir)
-        self.device = device
+        super().__init__(d, device, models_dir)
         show_inference_parameters(d)
         self.sr = SAMPLE_RATE
         self.ch = 2
@@ -99,6 +120,177 @@ class DemixerMDX(object):
             raise e
 
 
+# ############################################################################################################################
+# Demucs
+# ############################################################################################################################
+
+
+def get_steps_for_demucs(model, wav, segment, shifts, overlap):
+    segment = segment or model.segment
+    segment_length = int(model.samplerate * segment)
+    stride = int((1 - overlap) * segment_length)
+    return math.ceil(wav.shape[-1] / stride) * (shifts + 1)
+
+
+class DemixerDemucs(DemixerGeneric):
+    def __init__(self, d, device, models_dir):
+        super().__init__(d, device, models_dir)
+        self.sr = self.model_run.samplerate
+        # Demucs code will move the model to and from the device
+        # The advantage is that it will be do it for the sub_models
+        # So here we keep it offloaded and tell Demucs code to do the work
+        self.model_run.target_device = get_offload_device()
+
+    def get_steps(self, wav, segment, shifts, overlap):
+        """ Tries to figure out how much steps we will need for inference """
+        if isinstance(self.model_run, BagOfModels):
+            total = 0
+            for sub_model in self.model_run.models:
+                steps = get_steps_for_demucs(sub_model, wav, segment, shifts, overlap)
+                total += steps
+                logger.debug(f"- Steps {steps}")
+            logger.debug(f"Total steps {total}")
+            return total
+        steps = get_steps_for_demucs(self.model_run, wav, segment, shifts, overlap)
+        logger.debug(f"Steps {steps}")
+        return steps
+
+    def demucs_callback(self, v):
+        self.comfy_progress_bar.update(1)
+
+    def __call__(self, waveform_tensor, segment=None, shifts=0, overlap=0.25):
+        try:
+            # --- 1. Normalize input shape to handle both batched and non-batched data ---
+            if waveform_tensor.ndim == 2:
+                # Input is [C, samples], add a batch dimension to make it [1, C, samples]
+                logger.debug("Input is not batched. Adding a temporary batch dimension.")
+                waveform_tensor = waveform_tensor.unsqueeze(0)
+                input_was_batched = False
+            elif waveform_tensor.ndim == 3:
+                # Input is already batched [B, C, samples]
+                input_was_batched = True
+            else:
+                raise ValueError(f"Unsupported waveform shape: {waveform_tensor.shape}. Expected 2 or 3 dimensions.")
+
+            batch_size = waveform_tensor.shape[0]
+            logger.info("🎛️  Performing demix...")
+
+            model = self.model_run
+            input_tensor_on_device = waveform_tensor.to(self.device)
+            # Determine the segment size
+            # 1. User selection
+            # 2. Value in the config
+            # 3. Auto: from each model (when None)
+            forced_segment = segment
+            if forced_segment is None:
+                forced_segment = self.model_run.config_segment
+                if forced_segment:
+                    logger.debug(f"Using model provided segment size {forced_segment} s")
+                else:
+                    logger.debug("Using default segment size")
+                    forced_segment = None
+            else:
+                logger.debug(f"Using user provided segment size {forced_segment} s")
+            if with_comfy:
+                comfy_progress_bar = comfy.utils.ProgressBar(self.get_steps(waveform_tensor, forced_segment, shifts, overlap))
+
+            with model_to_target(model):
+                separated_tensors = apply_model(
+                    model,
+                    input_tensor_on_device,
+                    device=self.device,
+                    segment=forced_segment,
+                    shifts=shifts + 1,  # Shifts 0 is disabled, 1 is just one pass, 2 is 2 passes
+                    overlap=overlap,
+                    split=True,       # Enable chunking
+                    progress=True,    # Show a progress bar in the console
+                    callback=lambda x: (x.get('state') == 'end') and comfy_progress_bar.update(1) if with_comfy else None,
+                )
+
+            # Move the final result tensor back to the CPU before creating the output dicts.
+            # This is good practice to free up VRAM for subsequent nodes.
+            separated_tensors = separated_tensors.cpu()
+            assert batch_size == separated_tensors.shape[0]
+
+            # The output is [batch, sources, channels, samples].
+            # But we will separate the stems and return a batch for each stem, so we need
+            # [sources, batch, channels, samples].
+            separated_sources = separated_tensors.permute(1, 0, 2, 3)
+
+            # The model object tells us the names of the stems it produced
+            model_stems = model.sources.copy()
+            # UVR model is a 2 stems model with vocals and non_vocals, map it gracefully
+            if model_stems[1] == 'non_vocals':
+                model_stems[1] = 'other'
+            logger.debug(f"Model produced {len(model_stems)} stems: {model_stems}")
+
+            # Create a dictionary mapping the stem name to the audio tensor
+            output_map = dict(zip(model_stems, separated_sources))
+
+            # --- Silent audio for missing stems
+            reference_tensor = None
+            # Find the first valid tensor from our output to use as a shape reference
+            for tensor in output_map.values():
+                if tensor is not None and isinstance(tensor, torch.Tensor):
+                    reference_tensor = tensor
+                    break
+
+            # Determine the shape and sample rate for our silent audio fallback
+            if reference_tensor is not None:
+                # If we have a successful stem, use its properties
+                ref_batch_size, ref_channels, ref_samples = reference_tensor.shape
+            else:
+                # EDGE CASE: All stems failed. Fall back to the input audio's properties.
+                logger.warning("All stems failed. Using input audio shape for silence.")
+                # input_audio['waveform'] has shape [batch, channels, samples]
+                ref_batch_size, ref_channels, ref_samples = waveform_tensor.shape
+            # ---
+
+            # --- Gracefully create the 6 outputs ---
+            # Iterate through our fixed RETURN_NAMES and get the corresponding tensor.
+            # If a stem name doesn't exist in our output_map (e.g., 'guitar' for a 4-stem model),
+            # the .get() method will return None, which ComfyUI handles correctly.
+            final_outputs = []
+            for stem_name in ("vocals", "drums", "bass", "other", "guitar", "piano"):
+                output_tensor = output_map.get(stem_name, None)
+
+                # If the stem was produced, wrap it back into an AUDIO dict.
+                # If not, append None. ComfyUI handles None outputs correctly.
+                if output_tensor is not None:
+                    if not input_was_batched:
+                        output_tensor = output_tensor.squeeze(0)
+                    output_dict = {
+                        "waveform": output_tensor,
+                        "sample_rate": model.samplerate,
+                        "stem": stem_name.capitalize()
+                    }
+                    final_outputs.append(output_dict)
+                else:
+                    logger.debug(f"Stem '{stem_name}' not produced by model. Returning silence.")
+                    # Create a silent tensor with the correct shape, device, and dtype
+                    silent_waveform = torch.zeros(ref_batch_size, ref_channels, ref_samples, dtype=torch.float32)
+
+                    # Wrap the silent tensor in the ComfyUI AUDIO dictionary format
+                    silent_dict = {
+                        "waveform": silent_waveform,
+                        "sample_rate": model.samplerate,
+                        "stem": stem_name.capitalize()
+                    }
+                    final_outputs.append(silent_dict)
+
+            return tuple(final_outputs)
+
+        except Exception as e:
+            logger.error(f"Error during separation: {str(e)}")
+            raise e
+
+
 def get_demixer(d, device, models_dir):
-    # Currently just MDX
-    return DemixerMDX(d, device, models_dir)
+    model_t = d['model_t'].lower()
+    if model_t == "mdx":
+        return DemixerMDX(d, device, models_dir)
+    elif model_t == "demucs":
+        return DemixerDemucs(d, device, models_dir)
+    msg = f"Unknown model type `{model_t}`"
+    logger.error(msg)
+    raise ValueError(msg)

source/inference/get_model.py

@@ -4,18 +4,114 @@
 # Project: ComfyUI-AudioSeparation
 #
 # Helper to get a model from the correct class
+import importlib
+import json
 import logging
+from safetensors import safe_open
 from .MDX_Net import MDX_Net
-from ..utils.misc import NODES_NAME
+from ..utils.misc import NODES_NAME, json_object_hook
+# Demucs class imports
+from .demucs_api import BagOfModels
 
 logger = logging.getLogger(f"{NODES_NAME}.get_model")
 
 
-# Currently we have just one type of networks, but this is a clean way to support more, or even test replacements
+def get_metadata(file_path):
+    """ Read the metadata from a safetensors file """
+    logger.debug(f"Reading metadata from {file_path}")
+    metadata = {}
+    with safe_open(file_path, framework="pt", device="cpu") as f:
+        metadata = f.metadata()
+    if not metadata:
+        raise ValueError(f"Could not read metadata from safetensors file: {file_path}")
+    return metadata
+
+
+def get_hyperparameter(metadata, parameter, as_type, default=None, warn_diff=True):
+    value = metadata.get(parameter)
+    if value is None:
+        if default is None:
+            raise ValueError(f"Missing `{parameter}` hyperparameter")
+        return default
+    if as_type == "int":
+        value = int(value)
+    if warn_diff and value != default:
+        logger.warning(f"Hyperparameter mismatch: database = {default}, metadata = {value}")
+    return value
+
+
+def get_mdx_model(d):
+    """ Create an MDX_Net object with the specified parameters """
+    # Check the file is consistent we our data base
+    metadata = get_metadata(d['model_path'])
+    dim_f = get_hyperparameter(metadata, 'mdx_dim_f_set', "int", d['mdx_dim_f_set'])
+    channels = get_hyperparameter(metadata, 'channels', "int", d['channels'])
+    stages = get_hyperparameter(metadata, 'stages', "int", d['stages'])
+    # Create a class with this parameters
+    return MDX_Net(dim_f=dim_f, ch=channels, num_stages=stages)
+
+
+def get_demucs_model(d):
+    """ Create a Demucs, HDemucs (Hybrid) or HTDemucs (Hybrid Transformer) object.
+        All metadata comes from the safetensors """
+    file_path = d['model_path']
+
+    # 1. First, open the file safely to read only the metadata header.
+    metadata = get_metadata(file_path)
+
+    is_bag = json.loads(metadata.get('is_bag_of_models', 'false'))
+    signatures = json.loads(metadata['signatures'])
+
+    sub_models = []
+    for sig in set(signatures):  # Use set to only instantiate each unique architecture once
+        model_meta_str = metadata.get(sig)
+        if not model_meta_str:
+            raise ValueError(f"Metadata for signature '{sig}' not found in safetensors file.")
+
+        # Use the object_hook here to reconstruct Fraction objects automatically
+        model_meta = json.loads(model_meta_str, object_hook=json_object_hook)
+
+        class_module, class_name = model_meta['class_module'], model_meta['class_name']
+        args, kwargs = model_meta['args'], model_meta['kwargs']
+
+        logger.debug(f"  - Reconstructing architecture for '{sig}': {class_module}.{class_name}")
+        assert '.' not in class_name, "Security check failed, won't import a file outside my directory"
+        # Import the class from a module in this dir with the same name as the class
+        local_class_module = '.'.join(__name__.split('.')[:-1]) + "." + class_name
+        logger.debug(f"  - Redirecting class: {class_module} -> {local_class_module}")
+        module = importlib.import_module(local_class_module)
+        klass = getattr(module, class_name)
+
+        sub_models.append({'sig': sig, 'model': klass(*args, **kwargs)})
+
+    # Create a mapping from signature to model instance
+    model_map = {m['sig']: m['model'] for m in sub_models}
+
+    # Re-order the models to match the YAML's signature list
+    ordered_models = [model_map[sig] for sig in signatures]
+    segment = float(metadata.get('segment', '0'))
+
+    if not is_bag:
+        final_model = ordered_models[0]
+        final_model.signatures = signatures
+    else:
+        logger.debug("Rebuilding BagOfModels container...")
+        weights = json.loads(metadata.get('weights', 'null'))
+
+        final_model = BagOfModels(ordered_models, weights=weights, segment=segment)
+        final_model.signatures = signatures
+
+    final_model.config_segment = segment
+
+    return final_model
+
+
 def get_model(d):
     model_t = d['model_t'].lower()
-    if model_t != "mdx":
-        msg = f"Unknown model type `{model_t}`"
-        logger.error(msg)
-        raise ValueError(msg)
-    return MDX_Net(dim_f=d['mdx_dim_f_set'], ch=d['channels'], num_stages=d['stages'])
+    if model_t == "mdx":
+        return get_mdx_model(d)
+    elif model_t == "demucs":
+        return get_demucs_model(d)
+    msg = f"Unknown model type `{model_t}`"
+    logger.error(msg)
+    raise ValueError(msg)

source/inference/stft.py

@@ -126,7 +126,6 @@ def stft_chunk_process(waveform, d, model_run, device, segment_size=256, hop_len
 
     total_chunks = 1 + (mixture.shape[1] - chunk_size + step - 1) // step
     logger.info(f"⚙️  Processing {total_chunks} chunks...")
-    model_run.target_device = device
 
     with model_to_target(model_run):
         for i in tqdm(range(0, mixture.shape[1] - chunk_size + 1, step)):