diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..d721463 --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..19837a1 --- /dev/null +++ b/nodes.py @@ -0,0 +1,402 @@ +import torch + +import os, sys, re, shutil, subprocess +import hashlib +import numpy as np +import math +from typing import Tuple, List, Dict +import torchaudio.functional as F +from scipy.signal import get_window + +sys.path.insert(0, os.path.join(os.path.dirname(os.path.realpath(__file__)), "comfy")) + +import torchaudio + +import comfy.model_management +import comfy.sample +import comfy.sampler_helpers +import folder_paths + +MAX_RESOLUTION=8192 + +def is_url(url): + return url.split("://")[0] in ["http", "https"] + +def is_safe_path(path): + if "VHS_STRICT_PATHS" not in os.environ: + return True + basedir = os.path.abspath('.') + try: + common_path = os.path.commonpath([basedir, path]) + except: + #Different drive on windows + return False + return common_path == basedir + +def strip_path(path): + path = path.strip() + if path.startswith("\""): + path = path[1:] + if path.endswith("\""): + path = path[:-1] + return path + +def validate_path(path, allow_none=False, allow_url=True): + if path is None: + return allow_none + if is_url(path): + #Probably not feasible to check if url resolves here + if not allow_url: + return "URLs are unsupported for this path" + return is_safe_path(path) + if not os.path.isfile(strip_path(path)): + return "Invalid file path: {}".format(path) + return is_safe_path(path) +def calculate_file_hash(filename: str, hash_every_n: int = 1): + #Larger video files were taking >.5 seconds to hash even when cached, + #so instead the modified time from the filesystem is used as a hash + h = hashlib.sha256() + h.update(filename.encode()) + h.update(str(os.path.getmtime(filename)).encode()) + return h.hexdigest() +def hash_path(path): + if path is None: + return "input" + if is_url(path): + return "url" + return calculate_file_hash(strip_path(path)) + +ytdl_path = os.environ.get("VHS_YTDL", None) or shutil.which('yt-dlp') \ + or shutil.which('youtube-dl') + +download_history = {} +def try_download_video(url): + if ytdl_path is None: + return None + if url in download_history: + return download_history[url] + os.makedirs(folder_paths.get_temp_directory(), exist_ok=True) + #Format information could be added to only download audio for Load Audio, + #but this gets hairy if same url is also used for video. + #Best to just always keep defaults + #dl_format = ['-f', 'ba'] if is_audio else [] + try: + res = subprocess.run([ytdl_path, "--print", "after_move:filepath", + "-P", folder_paths.get_temp_directory(), url], + capture_output=True, check=True) + #strip newline + file = res.stdout.decode('utf-8')[:-1] + except subprocess.CalledProcessError as e: + raise Exception("An error occurred in the yt-dl process:\n" \ + + e.stderr.decode("utf-8")) + file = None + download_history[url] = file + return file + + + +class SignalProcessingLoadAudio(): + @classmethod + def INPUT_TYPES(s): + #Hide ffmpeg formats if ffmpeg isn't available + return { + "required": { + "audio_file": ("STRING", {"default": "input/", "signal_processins": ['wav','mp3','ogg','m4a','flac']}), + }, + "optional" : {"seek_seconds": ("FLOAT", {"default": 0, "min": 0})} + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + CATEGORY = "Signal Processing" + FUNCTION = "process" + + def process(self, audio_file, seek_seconds): + audio_file = strip_path(audio_file) + if audio_file is None or validate_path(audio_file) != True: + raise Exception("audio_file is not a valid path: " + audio_file) + if is_url(audio_file): + audio_file = try_download_video(audio_file) or audio_file + + print("Supported audio backends:", torchaudio.list_audio_backends()) + metadata = torchaudio.info(audio_file) + num_channels = metadata.num_channels + waveform, sample_rate = torchaudio.load(audio_file) + print('sample_rate',sample_rate) + # flatten to [sample_count] + waveform = waveform.unsqueeze(0).contiguous() + #waveform = waveform.transpose(0, 1).reshape(-1) + # flatten to [1, 2, sample_count] + #waveform = waveform.reshape((-1,num_channels)).transpose(0,1).unsqueeze(0) + #audio = waveform.reshape((-1,2)).transpose(0,1).unsqueeze(0) + return ({'waveform': waveform, 'sample_rate': sample_rate},) + #return (get_audio(audio_file, start_time=seek_seconds),) + + @classmethod + def IS_CHANGED(s, audio_file, seek_seconds): + return hash_path(audio_file) + + @classmethod + def VALIDATE_INPUTS(s, audio_file, **kwargs): + return validate_path(audio_file, allow_none=True) + +class SignalProcessingFilter: + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO", {"forceInput": True}), + "cutoff": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "filter_type": (["lowpass", "highpass", "bandpass", "bandstop"], {"default": "lowpass"}), + "q_factor": ("FLOAT", {"default": 0.707, "min": 0.1, "max": 5.0, "step": 0.01}), # For resonance/bandwidth + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + CATEGORY = "Signal Processing" + FUNCTION = "apply_filter" + + def apply_filter(self, audio: Dict[str, torch.Tensor], cutoff: float, filter_type: str, q_factor: float): + """ + Apply a specified filter to the input audio. + + Parameters: + audio (Dict[str, torch.Tensor]): Input audio with 'waveform' and 'sample_rate'. + cutoff (float): Normalized cutoff frequency (0.0 to 1.0). + filter_type (str): Type of filter ('lowpass', 'highpass', 'bandpass', 'bandstop'). + q_factor (float): Quality factor determining the filter's bandwidth. + + Returns: + Tuple[Dict[str, torch.Tensor]]: Filtered audio. + """ + waveform = audio['waveform'].squeeze(0) + sample_rate = audio['sample_rate'] + nyquist = sample_rate / 2.0 + + # Define minimum and maximum frequencies for mapping + log_min = 20.0 # 20 Hz, typical lower bound of human hearing + log_max = nyquist - 100.0 # Slightly below Nyquist to prevent instability + + # Avoid log(0) by ensuring cutoff is within (0,1) + cutoff = min(max(cutoff, 1e-6), 1.0 - 1e-6) + + # Logarithmic mapping + log_min = torch.log(torch.tensor(log_min)) + log_max = torch.log(torch.tensor(log_max)) + log_cutoff = log_min + cutoff * (log_max - log_min) + cutoff_freq = torch.exp(log_cutoff).item() + + # Debug: Print mapped cutoff frequency + print(f"Normalized cutoff: {cutoff} mapped to frequency: {cutoff_freq:.2f} Hz") + + # Choose filter type + if filter_type == "lowpass": + filtered_waveform = torchaudio.functional.lowpass_biquad( + waveform, sample_rate, cutoff_freq, Q=q_factor + ) + elif filter_type == "highpass": + filtered_waveform = torchaudio.functional.highpass_biquad( + waveform, sample_rate, cutoff_freq, Q=q_factor + ) + elif filter_type in ["bandpass", "bandstop"]: + center_freq = cutoff_freq + # Ensure that the bandwidth does not exceed the Nyquist frequency + bandwidth = center_freq / q_factor + lower_freq = max(center_freq - bandwidth / 2.0, 20.0) # Prevent dropping below 20 Hz + upper_freq = min(center_freq + bandwidth / 2.0, nyquist - 100.0) # Prevent exceeding Nyquist + + # Debug: Print bandwidth details + print(f"Band filter with center frequency: {center_freq:.2f} Hz, " + f"lower: {lower_freq:.2f} Hz, upper: {upper_freq:.2f} Hz") + + if filter_type == "bandpass": + filtered_waveform = torchaudio.functional.bandpass_biquad( + waveform, sample_rate, center_freq, Q=q_factor + ) + else: # bandstop + filtered_waveform = torchaudio.functional.band_biquad( + waveform, sample_rate, center_freq, Q=q_factor + ) + else: + raise ValueError(f"Unsupported filter type: {filter_type}") + + return ({'waveform': filtered_waveform.unsqueeze(0), 'sample_rate': sample_rate},) + +#https://github.com/paulnasca/paulstretch_python +#https://github.com/paulnasca/paulstretch_python/blob/master/paulstretch_stereo.py +class SignalProcessingPaulStretch: + @classmethod + def INPUT_TYPES(cls) -> Dict: + return { + "required": { + "audio": ("AUDIO", {"forceInput": True}), + "stretch_factor": ("FLOAT", {"default": 8.0, "min": 1.0, "max": 100.0, "step": 0.1}), + "window_size_seconds": ("FLOAT", {"default": 0.25, "min": 0.05, "max": 10.0, "step": 0.05}), + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + CATEGORY = "Signal Processing" + FUNCTION = "process" + + def process( + self, + audio: Dict[str, torch.Tensor], + stretch_factor: float, + window_size_seconds: float + ) -> Tuple[Dict[str, torch.Tensor]]: + """ + Apply the PaulStretch algorithm to time-stretch the input audio without altering its pitch. + + Parameters: + audio (Dict[str, torch.Tensor]): Input audio with 'waveform' and 'sample_rate'. + stretch_factor (float): Factor by which to stretch the audio (1.0 = no stretch). + window_size_seconds (float): Window size in seconds. + + Returns: + Tuple[Dict[str, torch.Tensor]]: PaulStretched audio with original sample rate. + """ + # Initialize logging + + # Conditional processing: If stretch_factor is 1.0, return original audio + if stretch_factor == 1.0: + return ({'waveform': audio['waveform'], 'sample_rate': audio['sample_rate']},) + + # Squeeze the first dimension (e.g., batch dimension) to get shape (C, N) or (N,) + waveform = audio['waveform'].squeeze(0).cpu().numpy() # Shape: (C, N) or (N,) for mono + sample_rate = audio['sample_rate'] + + # Handle mono audio by converting to stereo by duplicating the channel + if waveform.ndim == 1: + waveform = np.tile(waveform, (2, 1)) # Shape: (2, N) + + nchannels, nsamples = waveform.shape + + # Optimize window size to be divisible by 2, 3, and 5 + window_size = int(window_size_seconds * sample_rate) + if window_size < 16: + window_size = 16 + window_size = self.optimize_windowsize(window_size) + window_size = int(window_size / 2) * 2 # Ensure even window size + half_window_size = int(window_size / 2) + + # Correct the end of the waveform by applying a fade-out + end_size = int(sample_rate * 0.05) + if end_size < 16: + end_size = 16 + waveform[:, -end_size:] *= np.linspace(1, 0, end_size) + + # Compute displacement inside the input file + start_pos = 0.0 + displace_pos = (window_size * 0.5) / stretch_factor + + # Create custom window function as in original code + window = np.power(1.0 - np.power(np.linspace(-1.0, 1.0, window_size), 2.0), 1.25) + + # Initialize old windowed buffer + old_windowed_buf = np.zeros((nchannels, window_size)) + + # Initialize list to store output frames + output_frames = [] + + # Processing loop + frame_count = 0 + while True: + # Get the windowed buffer + istart_pos = int(math.floor(start_pos)) + buf = waveform[:, istart_pos:istart_pos + window_size] + if buf.shape[1] < window_size: + buf = np.pad(buf, ((0, 0), (0, window_size - buf.shape[1])), 'constant') + buf = buf * window + + # Get the amplitudes of the frequency components and discard the phases + freqs = np.abs(np.fft.rfft(buf, axis=1)) + + # Randomize the phases by multiplication with a random complex number with modulus=1 + ph = np.random.uniform(0, 2 * np.pi, (nchannels, freqs.shape[1])) * 1j + freqs = freqs * np.exp(ph) + + # Do the inverse FFT + buf_ifft = np.fft.irfft(freqs, n=window_size, axis=1) + + # Window again the output buffer + buf_ifft *= window + + # Overlap-add the output + output = buf_ifft[:, 0:half_window_size] + old_windowed_buf[:, half_window_size:window_size] + old_windowed_buf = buf_ifft + + # Clamp the values to -1..1 + output = np.clip(output, -1.0, 1.0) + + # Append to output_frames + output_frames.append(output) + + # Increment start_pos + start_pos += displace_pos + frame_count += 1 + + # Check if we have reached the end of the input + if start_pos >= nsamples: + break + + # Optional: Log progress every 10% + num_frames = max(int(math.ceil((nsamples - window_size) / (window_size * 0.5)) + 1), 1) + if frame_count % max(int(num_frames / 10), 1) == 0: + progress = (start_pos / nsamples) * 100 + + # Concatenate all output frames horizontally + output_array = np.hstack(output_frames) + + # Clamp final output to ensure no clipping + output_array = np.clip(output_array, -1.0, 1.0) + + # Convert back to torch tensor + output_tensor = torch.from_numpy(output_array).float() + + # Unsqueeze to add the first dimension back (e.g., batch dimension) + output_tensor = output_tensor.unsqueeze(0) # Shape: (1, C, N_stretched) + + + return ({'waveform': output_tensor, 'sample_rate': sample_rate},) + + @staticmethod + def optimize_windowsize(n: int) -> int: + """ + Optimize the window size to be divisible by 2, 3, and 5. + + Parameters: + n (int): Initial window size. + + Returns: + int: Optimized window size. + """ + orig_n = n + while True: + n = orig_n + while (n % 2) == 0: + n /= 2 + while (n % 3) == 0: + n /= 3 + while (n % 5) == 0: + n /= 5 + + if n < 2: + break + orig_n += 1 + return orig_n + + +NODE_CLASS_MAPPINGS = { + "SignalProcessingLoadAudio": SignalProcessingLoadAudio, + "SignalProcessingFilter": SignalProcessingFilter, + "SignalProcessingPaulStretch": SignalProcessingPaulStretch +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "SignalProcessingLoadAudio": "(SP) Load Audio", + "SignalProcessingFilter": "(SP) Filter", + "SignalProcessingPaulStretch" : "(SP) PaulStretch" +}