Bumped version to 1.1.1

They download the audio from internet
In most cases just uses 10 seconds to make it faster

README.md

@@ -30,12 +30,19 @@ to keep the secondary vocals along with the instruments.
 
 The objectives for these nodes are:
 
-- Multiple stems (Vocals, Instruments, Drums, Bass, etc.)
-- 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, or you can choose better quality using Demucs
-- Reduced dependencies
+✅ Multiple stems (Vocals, Instruments, Drums, Bass, etc.)
+
+✅ 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, or you can choose better quality using Demucs
+
+✅ Reduced dependencies, no need to install extra Python modules when using ComfyUI
+
+✅ Multiple examples
 
 ---
 
@@ -53,7 +60,9 @@ The objectives for these nodes are:
 *   ✨ [Nodes](#-nodes)
     *   [Vocals using MDX](#vocals-using-mdx)
     *   [Demucs Audio Separator](#demucs-audio-separator)
+*   🖼️ [Examples](#️-examples)
 *   📝 [Usage Notes](#-usage-notes)
+*   📜 [Project History](#-project-history)
 *   ⚖️ [License](#️-license)
 *   🙏 [Attributions](#-attributions)
 
@@ -250,6 +259,21 @@ And here is the Demucs node:
      - **Models:** Note that most models are a *bag of models*, this is four models working together.
 
 
+## 🖼️ Examples
+
+Once installed the examples are available in the ComfyUI workflow templates, in the *audio-separation* section.
+
+Note that we have two versions, the regular and the *quick* version. The *quick* version is ideal for quick tests, the input files
+are downloaded and, in most cases, only 10 seconds of audio are processed.
+
+- [00_Vocals_simple.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/00_Vocals_simple.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/00_Vocals_quick.json): Example to get vocals using MDX
+- [01_Vocals_Drums_Bass.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/01_Vocals_Drums_Bass.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/01_Vocals_Drums_Bass_quick.json): Example to get vocals, drums, bass and others using MDX
+- [02_Batch.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/02_Batch.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/02_Batch_quick.json): Shows how to apply MDX demix to a batch of audios, using *Audio Batch* nodes.
+- [03_Instrumental_keep.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/03_Instrumental_keep.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/03_Instrumental_keep_quick.json): Shows how to extract vocals maintaining the same number of channels and sample rate, using *Audio Batch* nodes.
+- [04_Demucs.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/04_Demucs.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/04_Demucs_quick.json): Separates vocals, drums, bass and others using Demucs, better quality.
+- [05_Demix_and_Remix.json](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/05_Demix_and_Remix.json) [Quick](https://raw.githubusercontent.com/set-soft/AudioSeparation/refs/heads/main/example_workflows/05_Demix_and_Remix_quick.json): Example to separate vocals and others to the left channel and drums and bass to the right channel, using *Audio Batch* nodes.
+
+
 ## 📝 Usage Notes
 
 - **AUDIO Type:** These nodes work with ComfyUI's standard "AUDIO" data type, which is a Python dictionary containing:
@@ -263,6 +287,15 @@ And here is the Demucs node:
 - **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.
 
 
+## 📜 Project History
+
+- 1.0.0 2025-07-02: Initial release. MDX-Net models support
+
+- 1.1.0 2025-07-11: Demucs models support
+
+- 1.1.1 2025-07-21: More examples. One more Demucs model
+
+
 ## ⚖️ License
 
 [GPL-3.0](LICENSE)

example_workflows/00_Vocals_simple.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/01_Vocals_Drums_Bass_quick.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/02_Batch_quick.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/03_Instrumental_keep_quick.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/04_Demucs_quick.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/05_Demix_and_Remix.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

example_workflows/05_Demix_and_Remix_quick.jpg

@@ -0,0 +1 @@
+audioseparation_logo.jpg

models/uvr_model_data.json

@@ -543,6 +543,21 @@
         "primary_stem": "Vocals",
         "stages": 5
     },
+    "5ee7b421b6228a90e63855b108ad830b": {
+        "desc": "Hybrid Demucs MusDB High Train",
+        "download": "Main/Demucs",
+        "file_t": "safetensors",
+        "model_t": "Demucs",
+        "name": "hdemucs_high_trained.safetensors",
+        "params": "83639368",
+        "primary_stem": [
+            "Vocals",
+            "Drums",
+            "Bass",
+            "Other"
+        ],
+        "use_demucs_pt_process": true
+    },
     "5f6483271e1efb9bfb59e4a3e6d4d098": {
         "channels": 32,
         "compensate": 1.035,
@@ -660,6 +675,21 @@
         "is_roformer": true,
         "model_type": "SCNet"
     },
+    "6f82220b9fc8445331d3a246f4ab509a": {
+        "desc": "Hybrid Demucs MusDB High Train",
+        "download": "Torch/Demucs",
+        "file_t": "pytorch",
+        "model_t": "Demucs",
+        "name": "hdemucs_high_trained.yaml",
+        "params": 83639368,
+        "primary_stem": [
+            "Vocals",
+            "Drums",
+            "Bass",
+            "Other"
+        ],
+        "signatures": "[\"hdemucs_high_trained\"]"
+    },
     "7205f4e53ec3226681ecd0c38e2fa747": {
         "desc": "Hybrid Transformer Demucs",
         "download": "Main/Demucs",

pyproject.toml

@@ -1,7 +1,11 @@
 [project]
 name = "audio-separation"
-description = "Audio separation (aka demixing) nodes, for Vocals, Instruments, Bass, Drums and Others. Using MDX-Net, no extra dependencies, support for batch and resample."
-version = "1.0.0"
+description = """
+Audio separation (aka demixing) nodes, for Vocals, Instruments, Bass, Drums and Others (experimental Piano and Guitar).
+Using MDX-Net and Demucs, no extra dependencies, support for batch and resample.
+Choose between High Quality and Speed. All safetensor models (No ONNX, No PyTorch)
+"""
+version = "1.1.1"
 license = { file = "LICENSE" }
 dependencies = []
 

source/inference/Demucs.py

@@ -11,8 +11,8 @@ import typing as tp
 import torch
 from torch import nn
 from torch.nn import functional as F
+import torchaudio
 
-from .resample_frac import resample_frac   # From julius
 from .demucs_code import capture_init, center_trim, unfold
 from .CrossTransformerEncoder import LayerScale
 
@@ -373,6 +373,8 @@ class Demucs(nn.Module):
         if rescale:
             rescale_module(self, reference=rescale)
 
+        self.upsampler = self.downsampler = None
+
     def valid_length(self, length):
         """
         Return the nearest valid length to use with the model so that
@@ -413,7 +415,16 @@ class Demucs(nn.Module):
         x = F.pad(x, (delta // 2, delta - delta // 2))
 
         if self.resample:
-            x = resample_frac(x, 1, 2)
+            if self.upsampler is None:
+                # Create the resamplers as instance attributes.
+                # They will be automatically moved to the correct device when model.to(device) is called.
+                self.upsampler = torchaudio.transforms.Resample(orig_freq=self.samplerate,
+                                                                new_freq=2 * self.samplerate,
+                                                                lowpass_filter_width=24).to(x.device)
+                self.downsampler = torchaudio.transforms.Resample(orig_freq=2 * self.samplerate,
+                                                                  new_freq=self.samplerate,
+                                                                  lowpass_filter_width=24).to(x.device)
+            x = self.upsampler(x)
 
         saved = []
         for encode in self.encoder:
@@ -429,7 +440,7 @@ class Demucs(nn.Module):
             x = decode(x + skip)
 
         if self.resample:
-            x = resample_frac(x, 2, 1)
+            x = self.downsampler(x)
         x = x * std + mean
         x = center_trim(x, length)
         x = x.view(x.size(0), len(self.sources), self.audio_channels, x.size(-1))

source/inference/HDemucs.py

@@ -70,7 +70,7 @@ class ScaledEmbedding(nn.Module):
 class HEncLayer(nn.Module):
     def __init__(self, chin, chout, kernel_size=8, stride=4, norm_groups=1, empty=False,
                  freq=True, dconv=True, norm=True, context=0, dconv_kw={}, pad=True,
-                 rewrite=True):
+                 rewrite=True, force_norm_in_last=False):
         """Encoder layer. This used both by the time and the frequency branch.
 
         Args:
@@ -109,6 +109,9 @@ class HEncLayer(nn.Module):
             pad = [pad, 0]
             klass = nn.Conv2d
         self.conv = klass(chin, chout, kernel_size, stride, pad)
+        if force_norm_in_last:
+            # PyTorch Audio uses it on last (empty) layer
+            self.norm1 = norm_fn(chout)
         if self.empty:
             return
         self.norm1 = norm_fn(chout)
@@ -257,7 +260,7 @@ class MultiWrap(nn.Module):
 class HDecLayer(nn.Module):
     def __init__(self, chin, chout, last=False, kernel_size=8, stride=4, norm_groups=1, empty=False,
                  freq=True, dconv=True, norm=True, context=1, dconv_kw={}, pad=True,
-                 context_freq=True, rewrite=True):
+                 context_freq=True, rewrite=True, force_norm_in_last=False):
         """
         Same as HEncLayer but for decoder. See `HEncLayer` for documentation.
         """
@@ -397,6 +400,7 @@ class HDemucs(nn.Module):
                  # Normalization
                  norm_starts=4,
                  norm_groups=4,
+                 force_norm_in_last=False,
                  # DConv residual branch
                  dconv_mode=1,
                  dconv_depth=2,
@@ -533,6 +537,7 @@ class HDemucs(nn.Module):
             kwt['kernel_size'] = kernel_size
             kwt['stride'] = stride
             kwt['pad'] = True
+            kwt['force_norm_in_last'] = force_norm_in_last
             kw_dec = dict(kw)
             multi = False
             if multi_freqs and index < multi_freqs_depth:

source/inference/demixer.py

@@ -7,6 +7,7 @@
 import logging
 import math
 import torch
+import tqdm
 # ComfyUI imports
 try:
     import comfy.utils
@@ -20,6 +21,8 @@ from ..utils.misc import NODES_NAME
 from ..utils.torch import model_to_target, get_offload_device
 from .demucs_api import apply_model, BagOfModels
 
+from torchaudio.transforms import Fade
+
 logger = logging.getLogger(f"{NODES_NAME}.demixer")
 SAMPLE_RATE = 44100
 
@@ -135,6 +138,62 @@ def get_steps_for_demucs(model, wav, segment, shifts, overlap):
     return math.ceil(wav.shape[-1] / stride) * (shifts + 1)
 
 
+def separate_sources(
+    model: torch.nn.Module,
+    mix: torch.Tensor,
+    sample_rate: int,
+    segment: float = 10.0,
+    overlap: float = 0.1,
+    device: torch.device = None,
+    chunk_fade_shape: str = "linear",
+) -> torch.Tensor:
+    """
+    From: https://pytorch.org/audio/stable/tutorials/hybrid_demucs_tutorial.html
+
+    Apply model to a given mixture. Use fade, and add segments together in order to add model segment by segment.
+
+    Args:
+        segment (int): segment length in seconds
+        device (torch.device, str, or None): if provided, device on which to
+            execute the computation, otherwise `mix.device` is assumed.
+            When `device` is different from `mix.device`, only local computations will
+            be on `device`, while the entire tracks will be stored on `mix.device`.
+    """
+    batch, channels, length = mix.shape
+
+    chunk_len = int(sample_rate * segment * (1 + overlap))
+    start = 0
+    end = chunk_len
+    overlap_frames = overlap * sample_rate
+    fade = Fade(fade_in_len=0, fade_out_len=int(overlap_frames), fade_shape=chunk_fade_shape)
+    chunks = math.ceil((length - overlap_frames) / chunk_len)
+    # Progress bars
+    progress_bar_console = tqdm.tqdm(total=chunks)
+    if with_comfy:
+        comfy_progress_bar = comfy.utils.ProgressBar(chunks)
+
+    final = torch.zeros(batch, len(model.sources), channels, length, device=device)
+
+    while start < length - overlap_frames:
+        chunk = mix[:, :, start:end]
+        progress_bar_console.update(1)
+        if with_comfy:
+            comfy_progress_bar.update(1)
+        with torch.no_grad():
+            out = model.forward(chunk)
+        out = fade(out)
+        final[:, :, :, start:end] += out
+        if start == 0:
+            fade.fade_in_len = int(overlap_frames)
+            start += int(chunk_len - overlap_frames)
+        else:
+            start += chunk_len
+        end += chunk_len
+        if end >= length:
+            fade.fade_out_len = 0
+    return final
+
+
 class DemixerDemucs(DemixerGeneric):
     def __init__(self, d, device, models_dir):
         super().__init__(d, device, models_dir)
@@ -194,21 +253,38 @@ class DemixerDemucs(DemixerGeneric):
                     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,
-                )
+            if self.d.get('use_demucs_pt_process', False):
+                # This is for the model from torchaudio, using the Demucs code we get some strange noises
+                # Using their example they aren't produced
+                model.target_device = self.device
+                logger.debug("Using PyTorch Audio chunking for old model")
+                with model_to_target(model):
+                    separated_tensors = separate_sources(
+                        model,
+                        input_tensor_on_device,
+                        model.samplerate,
+                        segment=forced_segment or 16.0,
+                        overlap=overlap,
+                        device=self.device,
+                        chunk_fade_shape="half_sine"
+                    )
+            else:
+                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.
@@ -265,19 +341,23 @@ class DemixerDemucs(DemixerGeneric):
                     output_dict = {
                         "waveform": output_tensor,
                         "sample_rate": model.samplerate,
-                        "stem": stem_name.capitalize()
+                        "stem": stem_name.capitalize(),
+                        "generated": True,
                     }
                     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)
+                    if not input_was_batched:
+                        silent_waveform = silent_waveform.squeeze(0)
 
                     # Wrap the silent tensor in the ComfyUI AUDIO dictionary format
                     silent_dict = {
                         "waveform": silent_waveform,
                         "sample_rate": model.samplerate,
-                        "stem": stem_name.capitalize()
+                        "stem": stem_name.capitalize(),
+                        "generated": False,
                     }
                     final_outputs.append(silent_dict)
 

source/inference/demucs_log_helper.py

@@ -0,0 +1,322 @@
+from fractions import Fraction
+import logging
+from ..utils.misc import NODES_NAME
+
+rlogger = logging.getLogger(f"{NODES_NAME}.demucs_log")
+
+
+class DemucsModelInfo(object):
+    def __init__(self, index, klass_name: str, kwargs: dict, logger, weights, extra=False, sig=None):
+        super().__init__()
+        self.index = index
+        self.kwargs = kwargs
+        self.logger = logger
+        self.extra = extra
+        self.extra_indent = ""
+        self.weights = weights
+        self.signature = sig
+
+        # Always log these fundamental parameters
+        sr = kwargs.get('samplerate', 44100)
+        if sr != 44100:
+            rlogger.warning("Model not configured for 44.1 kHz sample rate")
+        a_ch = kwargs.get('audio_channels', 2)
+        if a_ch != 2:
+            rlogger.warning("Model not configured for stereo")
+
+        if klass_name == "HTDemucs":
+            self.htdemucs()
+        elif klass_name == "HDemucs":
+            self.hdemucs()
+        elif klass_name == "Demucs":
+            self.demucs()
+        else:
+            logger.warning(f"No specific logger for model class: {klass_name}. "
+                           "Displaying raw kwargs.")
+            for key, value in kwargs.items():
+                logger(f"     - {key}: {value}")
+
+    def get(self, key, default=None):
+        return self.kwargs.get(key, default)
+
+    def log_type(self, name):
+        start = "" if self.index < 0 else f"{self.index+1}. "
+        msg = f"  {start}Type: {name}"
+        if self.signature:
+            msg += f" [{self.signature}]"
+        self.logger(msg)
+
+    def _log_param(self, param_name, default, description="", unit="", indent="  ", can_skip=False):
+        """
+        Logs a parameter if its value is different from the default, or if it's a key parameter.
+
+        Args:
+            param_name (str): The name of the parameter to check.
+            default: The default value for this parameter.
+            description (str): A user-friendly description of the parameter.
+            unit (str): An optional unit to display after the value (e.g., 'Hz').
+            indent (str): The indentation string for the log message.
+        """
+        value = self.get(param_name, default)
+
+        # We log if the value is not the default, or if it's a fundamental parameter.
+        is_default = (value == default)
+        is_important = param_name in ['sources', 'segment']
+
+        if is_default and not is_important and not self.extra:
+            return None
+        if can_skip and is_default:
+            return None
+
+        if param_name == 'sources' and self.weights:
+            value = [s if w == 1.0 else ('' if not w else f'{w}*{s}') for s, w in zip(value, self.weights)]
+
+        if unit == '%':
+            value *= 100
+        unit_str = f" {unit}" if unit else ""
+        desc_str = description if description else param_name.capitalize().replace('_', ' ')
+        indent += self.extra_indent
+        value_str = f"{value.numerator}/{value.denominator}" if isinstance(value, Fraction) else str(value)
+        n = (39 - len(desc_str) - len(value_str) - len(unit_str))*" "
+        return f"     {indent}- {desc_str}: {value_str}{unit_str} {n}({param_name})"
+
+    def log_param(self, param_name, default, description="", unit="", indent="  ", can_skip=False):
+        res = self._log_param(param_name, default, description, unit, indent, can_skip)
+        if res is not None:
+            self.logger(res)
+
+    def add(self, param_name, default, description="", unit="", indent="  ", can_skip=False):
+        res = self._log_param(param_name, default, description, unit, indent, can_skip)
+        if res is not None:
+            self.params.append(res)
+
+    def reset(self):
+        self.params = []
+
+    def sub_section(self, name):
+        self.logger(f"       {name}:")
+
+    def section(self, name):
+        self.logger("       " + "-" * 40)
+        self.sub_section(name)
+
+    def flush(self, name, is_sub=False):
+        if self.params:
+            if is_sub:
+                self.sub_section(self.extra_indent + name)
+            else:
+                self.section(name)
+            for p in self.params:
+                self.logger(p)
+
+    def structure(self, ch=64, depth=6, with_lstm=False, with_ch_tm=False):
+        self.reset()
+        self.add('channels', ch, "Initial hidden channels")
+        self.add('depth', depth, "Number of U-Net layers")
+        self.add('growth', 2.0, "Channel growth factor per layer")
+        self.add('rewrite', True, "Use 1x1 convolutions in blocks")
+        if with_lstm:
+            self.add('lstm_layers', 0, "Number of main LSTM layers", can_skip=True)
+        if with_ch_tm:
+            self.add('channels_time', None, "Specific channels for time branch", can_skip=True)
+        self.flush("Structure")
+
+    def convolutions(self, advanced=False):
+        self.reset()
+        self.add('kernel_size', 8)
+        self.add('stride', 4)
+        if advanced:
+            self.add('time_stride', 2, "Final time layer stride")
+        self.add('context', 1, "Decoder context window size")
+        if advanced:
+            self.add('context_enc', 0, "Encoder context window size")
+        self.flush("Convolutions")
+
+    def normalization(self):
+        self.reset()
+        self.add('norm_starts', 4, "Start at layer")
+        self.add('norm_groups', 4, "Number of groups")
+        self.flush("Normalization")
+
+    def dconv(self, full=True):
+        if self.get('dconv_mode', 1) <= 0:
+            return
+        self.reset()
+        where = ['', 'In encoder', 'In decoder', 'In encoder and decoder'][self.get('dconv_mode', 1)]
+        self.add('dconv_mode', 1, where)
+        self.add('dconv_depth', 2, "Number of layers in DConv branch")
+        if full:
+            comp = 4
+            init = 1e-4
+        else:
+            comp = 8
+            init = 1e-3
+        self.add('dconv_comp', comp, "Channel compression factor")
+        self.add('dconv_init', init, "Initial scale")
+        if full:
+            self.add('dconv_attn', 4, "Layer to start attention in DConv")
+            self.add('dconv_lstm', 4, "Layer to start LSTM in DConv")
+        self.flush("DConv Residual Branch")
+
+    def stft(self):
+        self.reset()
+        self.add('nfft', 4096, "Frequency Bins")
+        # Decode the method
+        cac = self.get('cac')
+        niters = self.get('wiener_iters', 0)
+        if cac:
+            zout = "Complex as Channels (CaC)"
+        elif niters >= 0:
+            zout = "Wiener filtering"
+        else:
+            zout = "Naive iSTFT from masking"
+        self.add('___', zout, "Framework")
+        self.add('cac', True, "Use Complex as Channels")
+        if not self.get('cac', True):
+            self.add('wiener_iters', 0, "Wiener filter iterations")
+        self.flush("STFT")
+
+    def freq_branch(self):
+        self.reset()
+        def_ratio = None
+        if self.get('multi_freqs') == []:
+            def_ratio = []
+        self.add('multi_freqs', def_ratio, "Ratios for frequency band splitting")
+        if self.get('multi_freqs'):
+            self.add('multi_freqs_depth', 2, "Layers to apply frequency splitting")
+        self.add('freq_emb', 0.2, "Frequency embedding weight")
+        if self.get('freq_emb'):
+            indent = "    "
+            self.add('emb_scale', 10, "Scale", indent=indent)
+            self.add('emb_smooth', True, "Smooth", indent=indent)
+        self.flush("Frequency Branch")
+
+    def demucs(self):
+        """Logs the parameters for the original Demucs class."""
+        self.log_type("Classic Waveform Demucs (Demucs)")
+        self.log_param('sources', [], "Target source names")
+        self.log_param('segment', 40, "Segment size", unit="s")
+        # --- Structure & Channels ---
+        self.structure(ch=64, depth=6, with_lstm=True)
+        # --- Convolutions ---
+        self.convolutions()
+
+        self.reset()
+        self.add('gelu', True, "GeLU (not ReLU)")
+        if self.get('rewrite', True):
+            self.add('glu', True, "GLU in 1x1 rewrite (not ReLU)")
+        self.flush("Activations")
+        # --- Normalization ---
+        self.normalization()
+        # --- DConv Residual Branch ---
+        self.dconv()
+        # --- Pre/Post Processing ---
+        self.reset()
+        self.add('resample', True, "Use 2x resampling")
+        self.add('normalize', True, "Normalize audio on-the-fly")
+        self.flush("Processing")
+
+    def hdemucs(self):
+        """Logs the parameters for the HDemucs (Hybrid Spectrogram/Waveform) class."""
+        self.log_type("Hybrid Demucs (Spectrogram + Waveform) (HDemucs)")
+        self.log_param('sources', [], "Target source names")
+        self.log_param('segment', 40, "Segment size", unit="s")
+        # --- Structure & Channels ---
+        self.structure(ch=48, depth=6, with_ch_tm=True)
+        # --- STFT & Spectrogram ---
+        self.stft()
+        # --- Frequency Branch ---
+        self.freq_branch()
+        # --- Convolutions ---
+        self.convolutions(advanced=True)
+        # --- Normalization ---
+        self.normalization()
+        # --- DConv Residual Branch (defaults are different from Demucs) ---
+        self.dconv()
+
+    def htdemucs(self):
+        """Logs the parameters for the HTDemucs (Hybrid Transformer) class."""
+        self.log_type("Hybrid Transformer Demucs (HTDemucs)")
+        self.log_param('sources', [], "Target source names")
+        self.log_param('segment', 10, "Segment size", unit="s")
+        # --- Structure & Channels (defaults are different from HDemucs) ---
+        self.structure(ch=48, depth=4)
+        # --- STFT & Spectrogram ---
+        self.stft()
+        # --- Frequency Branch ---
+        self.freq_branch()
+        # --- Convolutions ---
+        self.convolutions(advanced=True)
+        # --- Normalization ---
+        self.normalization()
+        # --- DConv (defaults are different) ---
+        self.dconv(full=False)
+        # --- Transformer Block ---
+        if self.get('t_layers', 5) > 0:
+            self.extra_indent = "  "
+            # --- Main Transformer ---
+            self.reset()
+            if self.get('bottom_channels', 0):
+                self.add('bottom_channels', 0, "Channels forced to")
+            self.add('t_hidden_scale', 4.0, "Hidden scale")
+            self.add('t_layers', 5, "Number of transformer layers")
+            self.add('t_heads', 8, "Number of attention heads")
+            self.add('t_dropout', 0.0, "Dropout")
+            self.flush("Transformer")
+            # --- Positional Embeddings ---
+            self.reset()
+            self.add('t_emb', 'sin', "Type")
+            self.add('t_weight_pos_embed', 1.0, "Weight", can_skip=True)
+            t_emb = self.get('t_emb', 'sin')
+            if t_emb == 'scaled':
+                self.add('t_max_positions', 10000, "Max positions")
+            elif t_emb == 'sin':
+                self.add('t_max_period', 10000.0, "Max period")
+                self.add('t_sin_random_shift', 0, "Random shift", can_skip=True)
+            elif t_emb == 'cape':
+                self.add('t_cape_mean_normalize', True, "Cape normalize")
+                self.add('t_cape_glob_loc_scale', [5000.0, 1.0, 1.4], "Cape params")
+                if self.get('t_cape_augment', True):
+                    rlogger.warning("t_cape_augment is True in loaded model, should be False for inference.")
+            self.flush("Positional Embeddings", is_sub=True)
+            # --- Transformer Normalization ---
+            self.reset()
+            self.add('t_norm_first', True, "Before attention/FFN")
+            self.add('t_norm_in', True, "Before pos. embedding")
+            if self.get('t_norm_in', True):
+                self.add('t_norm_in_group', False, "On all timesteps")
+            self.add('t_group_norm', False, "Of encoder on all timesteps")
+            self.add('t_norm_out', True, "GroupNorm at end of layers")
+            self.flush("Normalization", is_sub=True)
+            # --- Transformer Misc ---
+            self.reset()
+            self.add('t_cross_first', False, "Cross-attention is the first layer")
+            self.add('t_layer_scale', True, "Layer scale")
+            self.add('t_gelu', True, "GeLU (not ReLU)")
+            self.flush("Various", is_sub=True)
+            # --- Sparsity ---
+            # Log sparsity details only if sparse attention is enabled
+            self.reset()
+            is_sparse = self.get('t_sparse_self_attn', False)
+            self.add('t_sparse_self_attn', False, "Use sparse self-attention")
+            if is_sparse:
+                self.add('t_sparse_cross_attn', False, "Sparse cross-attention")
+                self.add('t_auto_sparsity', False, "Automatic sparsity")
+                auto_sparsity = self.get('t_auto_sparsity', False)
+                if not auto_sparsity:
+                    self.add('t_mask_type', 'diag', "Masking pattern")
+                    self.add('t_mask_random_seed', 42, "Mask seed")
+                    mask_t = self.get('t_mask_type', 'diag')
+                    if 'diag' in mask_t:
+                        self.add('t_sparse_attn_window', 500, "Window size")
+                    if 'global' in mask_t:
+                        self.add('t_global_window', 100, "Window size")
+                    if 'random' in mask_t:
+                        self.add('t_sparsity', 0.95, "Sparsity for random mask", unit="%")
+            self.flush("Sparsity", is_sub=True)
+            # Training only
+            # self.add(logger, kwargs, 't_weight_decay', 0.0, "Weight decay", extra=False)
+            # self.add(logger, kwargs, 't_lr', None, "Learning rate", extra=False)
+            # self.add(logger, kwargs, 't_cape_augment', True, "Learning rate", extra=False)
+            # self.add(logger, kwargs, 'rescale', 0.1, "Rescale trick", extra=False)
+            self.extra_indent = ""

source/inference/get_model.py

@@ -9,14 +9,15 @@ import json
 import logging
 from safetensors import safe_open
 from .MDX_Net import MDX_Net
-from ..utils.misc import NODES_NAME, json_object_hook
+from ..utils.misc import NODES_NAME, json_object_hook, get_debug_level
 # Demucs class imports
 from .demucs_api import BagOfModels
+from .demucs_log_helper import DemucsModelInfo
 
 logger = logging.getLogger(f"{NODES_NAME}.get_model")
 
 
-def get_metadata(file_path, d):
+def get_metadata(file_path, d=None):
     """ Read the metadata from a safetensors file """
     logger.debug(f"Reading metadata from {file_path}")
     metadata = {}
@@ -25,6 +26,8 @@ def get_metadata(file_path, d):
     if not metadata:
         raise ValueError(f"Could not read metadata from safetensors file: {file_path}")
 
+    if d is None:
+        return metadata
     # Is this a child model?
     parent = d.get('parent')
     if parent:
@@ -57,7 +60,7 @@ def get_hyperparameter(metadata, parameter, as_type, default=None, warn_diff=Tru
 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'])
+    metadata = get_metadata(d['model_path'], d)
     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'])
@@ -102,8 +105,11 @@ def get_demucs_model(d):
         logger.debug(f"  - Redirecting class: {class_module} -> {local_class_module}")
         module = importlib.import_module(local_class_module)
         klass = getattr(module, class_name)
+        instance = klass(*args, **kwargs)
+        instance._signature = sig
+        instance._metadata = model_meta
 
-        sub_models.append({'sig': sig, 'model': klass(*args, **kwargs)})
+        sub_models.append({'sig': sig, 'model': instance})
 
     # Create a mapping from signature to model instance
     model_map = {m['sig']: m['model'] for m in sub_models}
@@ -113,17 +119,37 @@ def get_demucs_model(d):
     segment = float(metadata.get('segment', '0'))
 
     if not is_bag:
+        weights = None
         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
 
+    debug_level = get_debug_level(logger)
+    if debug_level >= 1:
+        # Show some information of the resulting model
+        logger.debug("Model information:")
+        logger.debug(f"- Total models {len(ordered_models)}")
+        if weights is not None and len(weights) != len(ordered_models):
+            raise ValueError(f"Invalid weights for {len(ordered_models)} models: {weights}")
+        for n, m in enumerate(ordered_models):
+            tp = m.__class__.__name__
+            w = weights[n] if weights else None
+            if n == 0:
+                ref_sources = m.sources
+            else:
+                if m.sources != ref_sources:
+                    logger.error("The sub-model outputs doesn't match")
+            if w and len(m.sources) != len(w):
+                raise ValueError(f"Invalid {w} weights for {m.sources} sources")
+            num = -1 if len(ordered_models) == 1 else n
+            DemucsModelInfo(num, tp, m._metadata['kwargs'], logger.debug, w, extra=debug_level > 1, sig=m._signature)
+
     return final_model
 
 

source/inference/resample_frac.py

@@ -1,220 +0,0 @@
-# File under the MIT license, see https://github.com/adefossez/julius/LICENSE for details.
-# Author: adefossez, 2020
-"""
-Differentiable, Pytorch based resampling.
-Implementation of Julius O. Smith algorithm for resampling.
-See https://ccrma.stanford.edu/~jos/resample/ for details.
-This implementation is specially optimized for when new_sr / old_sr is a fraction
-with a small numerator and denominator when removing the gcd (e.g. new_sr = 700, old_sr = 500).
-
-Very similar to [bmcfee/resampy](https://github.com/bmcfee/resampy) except this implementation
-is optimized for the case mentioned before, while resampy is slower but more general.
-
-"""
-import math
-from typing import Optional
-
-import torch
-from torch.nn import functional as F
-
-
-def sinc(x: torch.Tensor):
-    """
-    Implementation of sinc, i.e. sin(x) / x
-
-    __Warning__: the input is not multiplied by `pi`!
-    """
-    return torch.where(x == 0, torch.tensor(1., device=x.device, dtype=x.dtype), torch.sin(x) / x)
-
-
-class ResampleFrac(torch.nn.Module):
-    """
-    Resampling from the sample rate `old_sr` to `new_sr`.
-    """
-    def __init__(self, old_sr: int, new_sr: int, zeros: int = 24, rolloff: float = 0.945):
-        """
-        Args:
-            old_sr (int): sample rate of the input signal x.
-            new_sr (int): sample rate of the output.
-            zeros (int): number of zero crossing to keep in the sinc filter.
-            rolloff (float): use a lowpass filter that is `rolloff * new_sr / 2`,
-                to ensure sufficient margin due to the imperfection of the FIR filter used.
-                Lowering this value will reduce anti-aliasing, but will reduce some of the
-                highest frequencies.
-
-        Shape:
-
-            - Input: `[*, T]`
-            - Output: `[*, T']` with `T' = int(new_sr * T / old_sr)
-
-
-        .. caution::
-            After dividing `old_sr` and `new_sr` by their GCD, both should be small
-            for this implementation to be fast.
-
-        >>> import torch
-        >>> resample = ResampleFrac(4, 5)
-        >>> x = torch.randn(1000)
-        >>> print(len(resample(x)))
-        1250
-        """
-        super().__init__()
-        if not isinstance(old_sr, int) or not isinstance(new_sr, int):
-            raise ValueError("old_sr and new_sr should be integers")
-        gcd = math.gcd(old_sr, new_sr)
-        self.old_sr = old_sr // gcd
-        self.new_sr = new_sr // gcd
-        self.zeros = zeros
-        self.rolloff = rolloff
-
-        self._init_kernels()
-
-    def _init_kernels(self):
-        if self.old_sr == self.new_sr:
-            return
-
-        kernels = []
-        sr = min(self.new_sr, self.old_sr)
-        # rolloff will perform antialiasing filtering by removing the highest frequencies.
-        # At first I thought I only needed this when downsampling, but when upsampling
-        # you will get edge artifacts without this, the edge is equivalent to zero padding,
-        # which will add high freq artifacts.
-        sr *= self.rolloff
-
-        # The key idea of the algorithm is that x(t) can be exactly reconstructed from x[i] (tensor)
-        # using the sinc interpolation formula:
-        #   x(t) = sum_i x[i] sinc(pi * old_sr * (i / old_sr - t))
-        # We can then sample the function x(t) with a different sample rate:
-        #    y[j] = x(j / new_sr)
-        # or,
-        #    y[j] = sum_i x[i] sinc(pi * old_sr * (i / old_sr - j / new_sr))
-
-        # We see here that y[j] is the convolution of x[i] with a specific filter, for which
-        # we take an FIR approximation, stopping when we see at least `zeros` zeros crossing.
-        # But y[j+1] is going to have a different set of weights and so on, until y[j + new_sr].
-        # Indeed:
-        # y[j + new_sr] = sum_i x[i] sinc(pi * old_sr * ((i / old_sr - (j + new_sr) / new_sr))
-        #               = sum_i x[i] sinc(pi * old_sr * ((i - old_sr) / old_sr - j / new_sr))
-        #               = sum_i x[i + old_sr] sinc(pi * old_sr * (i / old_sr - j / new_sr))
-        # so y[j+new_sr] uses the same filter as y[j], but on a shifted version of x by `old_sr`.
-        # This will explain the F.conv1d after, with a stride of old_sr.
-        self._width = math.ceil(self.zeros * self.old_sr / sr)
-        # If old_sr is still big after GCD reduction, most filters will be very unbalanced, i.e.,
-        # they will have a lot of almost zero values to the left or to the right...
-        # There is probably a way to evaluate those filters more efficiently, but this is kept for
-        # future work.
-        idx = torch.arange(-self._width, self._width + self.old_sr).float()
-        for i in range(self.new_sr):
-            t = (-i/self.new_sr + idx/self.old_sr) * sr
-            t = t.clamp_(-self.zeros, self.zeros)
-            t *= math.pi
-            window = torch.cos(t/self.zeros/2)**2
-            kernel = sinc(t) * window
-            # Renormalize kernel to ensure a constant signal is preserved.
-            kernel.div_(kernel.sum())
-            kernels.append(kernel)
-
-        self.register_buffer("kernel", torch.stack(kernels).view(self.new_sr, 1, -1))
-
-    def forward(self, x: torch.Tensor, output_length: Optional[int] = None, full: bool = False):
-        """
-        Resample x.
-        Args:
-            x (Tensor): signal to resample, time should be the last dimension
-            output_length (None or int): This can be set to the desired output length
-                (last dimension). Allowed values are between 0 and
-                ceil(length * new_sr / old_sr). When None (default) is specified, the
-                floored output length will be used. In order to select the largest possible
-                size, use the `full` argument.
-            full (bool): return the longest possible output from the input. This can be useful
-                if you chain resampling operations, and want to give the `output_length` only
-                for the last one, while passing `full=True` to all the other ones.
-        """
-        if self.old_sr == self.new_sr:
-            return x
-        shape = x.shape
-        length = x.shape[-1]
-        x = x.reshape(-1, length)
-        x = F.pad(x[:, None], (self._width, self._width + self.old_sr), mode='replicate')
-        ys = F.conv1d(x, self.kernel, stride=self.old_sr)  # type: ignore
-        y = ys.transpose(1, 2).reshape(list(shape[:-1]) + [-1])
-
-        float_output_length = torch.as_tensor(self.new_sr * length / self.old_sr)
-        max_output_length = torch.ceil(float_output_length).long()
-        default_output_length = torch.floor(float_output_length).long()
-
-        if output_length is None:
-            applied_output_length = max_output_length if full else default_output_length
-        elif output_length < 0 or output_length > max_output_length:
-            raise ValueError(f"output_length must be between 0 and {max_output_length}")
-        else:
-            applied_output_length = torch.tensor(output_length)
-            if full:
-                raise ValueError("You cannot pass both full=True and output_length")
-        return y[..., :applied_output_length]  # type: ignore
-
-
-def resample_frac(x: torch.Tensor, old_sr: int, new_sr: int,
-                  zeros: int = 24, rolloff: float = 0.945,
-                  output_length: Optional[int] = None, full: bool = False):
-    """
-    Functional version of `ResampleFrac`, refer to its documentation for more information.
-
-    ..warning::
-        If you call repeatidly this functions with the same sample rates, then the
-        resampling kernel will be recomputed every time. For best performance, you should use
-        and cache an instance of `ResampleFrac`.
-    """
-    return ResampleFrac(old_sr, new_sr, zeros, rolloff).to(x)(x, output_length, full)
-
-
-# Easier implementations for downsampling and upsampling by a factor of 2
-# Kept for testing and reference
-
-def _kernel_upsample2_downsample2(zeros):
-    # Kernel for upsampling and downsampling by a factor of 2. Interestingly,
-    # it is the same kernel used for both.
-    win = torch.hann_window(4 * zeros + 1, periodic=False)
-    winodd = win[1::2]
-    t = torch.linspace(-zeros + 0.5, zeros - 0.5, 2 * zeros)
-    t *= math.pi
-    kernel = (sinc(t) * winodd).view(1, 1, -1)
-    return kernel
-
-
-def _upsample2(x, zeros=24):
-    """
-    Upsample x by a factor of two. The output will be exactly twice as long as the input.
-    Args:
-        x (Tensor): signal to upsample, time should be the last dimension
-        zeros (int): number of zero crossing to keep in the sinc filter.
-
-    This function is kept only for reference, you should use the more generic `resample_frac`
-    one. This function does not perform anti-aliasing filtering.
-    """
-    *other, time = x.shape
-    kernel = _kernel_upsample2_downsample2(zeros).to(x)
-    out = F.conv1d(x.view(-1, 1, time), kernel, padding=zeros)[..., 1:].view(*other, time)
-    y = torch.stack([x, out], dim=-1)
-    return y.view(*other, -1)
-
-
-def _downsample2(x, zeros=24):
-    """
-    Downsample x by a factor of two. The output length is half of the input, ceiled.
-    Args:
-        x (Tensor): signal to downsample, time should be the last dimension
-        zeros (int): number of zero crossing to keep in the sinc filter.
-
-    This function is kept only for reference, you should use the more generic `resample_frac`
-    one. This function does not perform anti-aliasing filtering.
-    """
-    if x.shape[-1] % 2 != 0:
-        x = F.pad(x, (0, 1))
-    xeven = x[..., ::2]
-    xodd = x[..., 1::2]
-    *other, time = xodd.shape
-    kernel = _kernel_upsample2_downsample2(zeros).to(x)
-    out = xeven + F.conv1d(xodd.view(-1, 1, time), kernel, padding=zeros)[..., :-1].view(
-        *other, time)
-    return out.view(*other, -1).mul(0.5)

source/utils/misc.py

@@ -10,8 +10,12 @@ NODES_NAME = "AudioSeparation"
 NODES_DEBUG_VAR = NODES_NAME.upper() + "_NODES_DEBUG"
 
 
+def get_debug_level(logger):
+    return logging.DEBUG - logger.getEffectiveLevel() + 1
+
+
 def debugl(logger, level, msg):
-    if logger.getEffectiveLevel() <= logging.DEBUG - (level - 1):
+    if get_debug_level(logger) >= level:
         logger.debug(msg)
 
 

tool/batch_convert.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3

tool/demix.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3
@@ -47,7 +48,7 @@ def demix(d, args):
 
     if args.save_complement:
         for wav in wavs[1:]:
-            if wav is None:
+            if wav is None or not wav['generated']:
                 continue
             out_path = f"{args.out_base or base}_{wav['stem']}{out_ext}"
             save_audio(wav['waveform'], wav['sample_rate'], out_path, out_format)

tool/demucs2safetensors.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3
@@ -25,7 +26,7 @@ try:
 except Exception:
     with_demuc_lib = False
 import bootstrap  # noqa: F401
-from source.utils.misc import cli_add_verbose, FractionEncoder
+from source.utils.misc import cli_add_verbose, FractionEncoder, debugl
 from source.utils.logger import main_logger, logger_set_standalone
 from source.db.models_db import cli_add_db, get_download_url, ModelsDB
 from source.db.hash import get_hash
@@ -36,6 +37,10 @@ MODULES_MAP = {'demucs': local_demucs_module,
                'demucs.demucs': local_demucs_module,
                'demucs.hdemucs': local_hdemucs_module,
                'demucs.htdemucs': local_htdemucs_module}
+MAP = {'freq_encoder': 'encoder',
+       'freq_decoder': 'decoder',
+       'time_encoder': 'tencoder',
+       'time_decoder': 'tdecoder'}
 logger = main_logger
 
 
@@ -61,6 +66,36 @@ def remap_module(modules_map):
                 del sys.modules[old_name]
 
 
+def solve_simple_pt(yaml_data, pkg):
+    """ PyTorch Audio lib has some raw models, we store the metadata in the YAML """
+    klass = yaml_data.get('klass')
+    if klass is None:
+        main_logger.error("No `klass` in YAML")
+        sys.exit(4)
+    if klass == 'Demucs':
+        klass = local_demucs_module.Demucs
+    elif klass == 'HDemucs':
+        klass = local_hdemucs_module.HDemucs
+    elif klass == 'HTDemucs':
+        klass = local_htdemucs_module.HTDemucs
+    else:
+        main_logger.error("Unknown model `klass` {klass}")
+        sys.exit(4)
+    args = yaml_data.get('args', {})
+    kwargs = yaml_data.get('kwargs', {})
+
+    # For PyTorch Audio model (very old code?)
+    new_dict = {}
+    for k, v in pkg.items():
+        parts = k.split('.')
+        gr = parts[0]
+        if gr in MAP:
+            k = MAP[gr] + '.' + '.'.join(parts[1:])
+        new_dict[k] = v
+
+    return klass, args, kwargs, new_dict
+
+
 def convert_demucs_model(yaml_path_str: str, model_paths: list[str], output_path_str: str, all_metadata, data: Dict):
     """
     Loads an original Demucs model bag, extracts weights and all necessary metadata
@@ -111,7 +146,11 @@ def convert_demucs_model(yaml_path_str: str, model_paths: list[str], output_path
             # This is the only "unsafe" part, loading the original pickle file
             pkg = torch.load(path, map_location='cpu', weights_only=False)
 
-        klass, args, kwargs, state = pkg["klass"], pkg["args"], pkg["kwargs"], pkg["state"]
+        debugl(logger, 2, f"PyTorch data type is {type(pkg)}")
+        if isinstance(pkg, dict) and 'klass' in pkg:
+            klass, args, kwargs, state = pkg["klass"], pkg["args"], pkg["kwargs"], pkg["state"]
+        else:
+            klass, args, kwargs, state = solve_simple_pt(yaml_data, pkg)
         main_logger.info(f"  - Model class: {klass.__module__}.{klass.__name__}")
 
         # Dequantize if necessary by letting the original code handle it
@@ -162,6 +201,7 @@ def convert_demucs_model(yaml_path_str: str, model_paths: list[str], output_path
 
 if __name__ == '__main__':
     parser = argparse.ArgumentParser(description="Convert Demucs .th models to a single .safetensors file.")
+
     parser.add_argument('--yaml', required=True, type=str, help="Path to the Demucs .yaml file.")
     parser.add_argument('--models', nargs='*', default=None, type=str,
                         help="Optional. Paths to .th files. If not provided, assumes they are "
@@ -205,7 +245,7 @@ if __name__ == '__main__':
 
         model_files = []
         for sig in set(signatures):  # Use set to avoid redundant searches
-            found = list(yaml_dir.glob(f'{sig}*.th'))
+            found = list(yaml_dir.glob(f'{sig}*.th')) + list(yaml_dir.glob(f'{sig}*.pt'))
             if not found:
                 raise FileNotFoundError(f"Could not automatically find a model file for signature '{sig}' in {yaml_dir}")
             if len(found) > 1:

tool/download_model.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3

tool/onnx2safetensors.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3

tool/show_class.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3

tool/show_db.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3

tool/show_metadata.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3
@@ -13,6 +14,7 @@ import bootstrap  # noqa: F401
 from source.utils.logger import main_logger, logger_set_standalone
 from source.utils.misc import cli_add_verbose, json_object_hook
 from source.inference.get_model import get_metadata
+from source.inference.demucs_log_helper import DemucsModelInfo
 
 
 def show_metadata(args):
@@ -23,11 +25,25 @@ def show_metadata(args):
     expanded = {k: json.loads(v, object_hook=json_object_hook) if v[0] in '{[' else v for k, v in d.items()}
     pprint(expanded)
 
+    sigs = expanded.get('signatures')
+    if sigs:
+        single = len(sigs) == 1
+        print("Demucs model")
+        if not single:
+            print(f"Composed by {len(sigs)} submodels")
+        weights = expanded.get('weights')
+        # Demucs model
+        for n, s in enumerate(sigs):
+            d = expanded[s]
+            num = -1 if single else n
+            DemucsModelInfo(num, d['class_name'], d['kwargs'], print, weights[n] if weights else None, extra=args.verbose,
+                            sig=s)
+
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser(description="Shows the safetensors metadata",
                                      formatter_class=argparse.RawTextHelpFormatter)
-    parser.add_argument('input_file', type=str, help="Path to the input ONNX model file.")
+    parser.add_argument('input_file', type=str, help="Path to the input safetensors model file.")
     cli_add_verbose(parser)
 
     args = parser.parse_args()

tool/show_onnx.py

@@ -1,3 +1,4 @@
+#!/usr/bin/env python3
 # Copyright (c) 2025 Salvador E. Tropea
 # Copyright (c) 2025 Instituto Nacional de Tecnología Industrial
 # License: GPLv3