node.py

"""
╔════════════════════════════════════════╗
║                                        ║
║   Developed by Abdullah OZMANTAR       ║
║   GitHub : github.com/abdozmantar      ║
║                                        ║
╚════════════════════════════════════════╝
"""

import sys

from dora.log import fatal
import torch as th
from pathlib import Path

import random
import string

from .audio import save_audio
from .helper import Separator
from .pretrained import SOURCES
from .repo import ModelLoadingError

class MultiInput(str):
    def __new__(cls, string, allowed_types="*"):
        res = super().__new__(cls, string)
        res.allowed_types=allowed_types
        return res
    def __ne__(self, other):
        if self.allowed_types == "*" or other == "*":
            return False
        return other not in self.allowed_types
    
audioOrNone = MultiInput("AUDIO", ["AUDIO", "*"])

class DeepExtractV2Node:
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "track": ("AUDIO", {"default": None, "multiline": True, "tooltip": "Input audio track to separate"}),
                "out": ("STRING", {"default": "separated", "tooltip": "Folder where to put extracted track. This folder will be created inside the main ComfyUI directory."}),
                "device": (["cuda", "cpu"], {"default": "cuda", "tooltip": "Device to use for inference"}),
                "overlap": ("FLOAT", {"default": 0.25, "min": 0.0, "max": 1.0, "step": 0.05, "tooltip": "Overlap between audio chunks"}),
                "shifts": ("INT", {"default": 1, "min": 1, "max": 20, "tooltip": "Number of random shifts for stabilization"}),
                "segment": ("INT", {"default": 0, "min": 0, "max": 600, "tooltip": "Split size of each chunk (0 = no split)"}),
                "bit_depth": (["default", "int24", "float32"], {"default": "default", "tooltip": "Audio output bit depth"}),
                "clip_mode": (["rescale", "clamp", "none"], {"default": "rescale", "tooltip": "Clipping handling strategy"}),
                "audio_format": (["wav", "flac", "mp3"], {"default": "wav", "tooltip": "Output audio file format"}),
                "mp3_bitrate": ("INT", {"default": 320, "min": 64, "max": 320, "tooltip": "Bitrate for MP3 output"}),
                "mp3_preset": ("INT", {"default": 2, "min": 2, "max": 7, "tooltip": "Encoder preset (2=best, 7=fast)"}),
                "jobs": ("INT", {"default": 0, "min": 0, "max": 32, "tooltip": "Number of parallel jobs"}),
                "split": ("BOOLEAN", {"default": True, "tooltip": "Split audio in chunks to save memory"}),
                "other_method": (["none", "add", "minus"], {"default": "add","tooltip": 'Decide how to get "no_{STEM}". "none" will not save "no_{STEM}". ''"add" will add all the other stems. "minus" will use the original track minus the selected stem.'}),
                "stem": ("STRING", { "default": None ,"tooltip": "Only separate audio into {STEM} and no_{STEM}."}),
                "filename": ("STRING", {"default": "{track}/{stem}.{ext}", "tooltip": "Set the name of output file. Use '{track}', '{stem}', '{ext}' to use variables of track name without extension, track extension,stem name and default output file extension. Default is {track}/{stem}.{ext}"}),
            }
        }
    
    OUTPUT_SOURCES = tuple(list(SOURCES) + ["derived_stem"])

    RETURN_TYPES = tuple([audioOrNone] * len(OUTPUT_SOURCES))
    RETURN_NAMES = tuple(OUTPUT_SOURCES)

    FUNCTION = "separate"
    CATEGORY = "DeepExtractV2"

   
    
    def separate(self,
                 track, 
                 out, 
                 device, 
                 overlap, 
                 shifts,
                 segment,
                 bit_depth,
                 clip_mode,
                 audio_format,
                 mp3_bitrate,
                 mp3_preset,
                 jobs,
                 split,
                 other_method,
                 stem,
                 filename):

        if(segment == 0):
            segment = None
            
        out = Path(out)
        
        if len(track) == 0:
                print("error: the following arguments are required: track", file=sys.stderr)
                sys.exit(1)

        try:
            separator = Separator(
                                model="htdemucs",
                                device=device,
                                shifts=shifts,
                                split=split,
                                overlap=overlap,
                                progress=True,
                                jobs=jobs,
                                segment=segment)
        except ModelLoadingError as error:
            fatal(error.args[0])

        if stem != "" and stem not in separator.model.sources:
            fatal(
                'error: stem "{stem}" is not in selected model. '
                "STEM must be one of {sources}.".format(
                    stem=stem, sources=", ".join(separator.model.sources)
                )
            )
        out = out
        out.mkdir(parents=True, exist_ok=True)
        print(f"Separated tracks will be stored in {out.resolve()}")

        def random_name(length=8):
            return ''.join(random.choices(string.ascii_lowercase + string.digits, k=length))

        track["name"] = random_name()

        origin, res = separator.separate_audio_file(track)

        if audio_format == "mp3":
            ext = "mp3"
        elif audio_format == "flac":
            ext = "flac"
        else:
            ext = "wav"
        kwargs = {
            "samplerate": separator.samplerate,
            "bitrate": mp3_bitrate,
            "preset": mp3_preset,
            "clip": clip_mode,
            "as_float": bit_depth == "float32",
            "bits_per_sample": 24 if  bit_depth == "int24" else 16,
        }

        outputs = {}

        if not stem or stem.strip() == "":
            for name, source in res.items():
                output_path = out / filename.format(
                    track=track["name"].rsplit(".", 1)[0],
                    stem=name,
                    ext=ext,
                )
                output_path.parent.mkdir(parents=True, exist_ok=True)
                outputs[name] = save_audio(source, str(output_path), **kwargs)

        else:
            selected_stem = stem  
            minus_stem_path = out / filename.format(
                track=track["name"].rsplit(".", 1)[0],
                stem=f"minus_{selected_stem}",
                ext=ext,
            )
            if other_method == "minus":
                minus_stem_path.parent.mkdir(parents=True, exist_ok=True)
                outputs["derived_stem"] = save_audio(origin - res[selected_stem], str(minus_stem_path), **kwargs)
            stem_path = out / filename.format(
                track=track["name"].rsplit(".", 1)[0],
                stem=selected_stem,
                ext=ext,
            )
            stem_path.parent.mkdir(parents=True, exist_ok=True)
            outputs[selected_stem] = save_audio(res.pop(selected_stem), str(stem_path), **kwargs)
            if other_method == "add":
                other_stem = th.zeros_like(next(iter(res.values())))
                for src in res.values():
                    other_stem += src
                no_stem_path = out / filename.format(
                    track=track["name"].rsplit(".", 1)[0],
                    stem=f"no_{selected_stem}",
                    ext=ext,
                )
                no_stem_path.parent.mkdir(parents=True, exist_ok=True)
                outputs["derived_stem"] = save_audio(other_stem, str(no_stem_path), **kwargs)

        for source in self.OUTPUT_SOURCES:
            if source not in outputs:
                outputs[source] = {
                    "waveform": th.zeros(2, 1).unsqueeze(0),
                    "sample_rate": 44100
                }

        ordered_outputs = [outputs[source] for source in self.OUTPUT_SOURCES]
        
        print(outputs)
        drums, bass, others, vocals, derived_stem = ordered_outputs

        return (drums, bass, others, vocals, derived_stem)