diff --git a/README.md b/README.md index 3f4ba4f..2b5ab4a 100644 --- a/README.md +++ b/README.md @@ -19,11 +19,12 @@ Might also double as a poor-man's cryptanalysis and steganography lab. - Bytes-like datatype and encoders/decoders for many nodes in the node pack. - Error correction codes for more robust operations - currently powered by [`reedsolo`](https://github.com/tomerfiliba-org/reedsolomon) - Multi-language Morse code encoder and decoder. +- Spectrogram to audio and vice versa. - (WIP) Modern cryptographic encryption tools and functions - currently powered by [`cryptography`](https://github.com/pyca/cryptography). - (Planned) Historical encoding and encrypting formats (Kansas City System, SNOW, SSTV,...) made user-friendly and available through a node interface. - (Planned) Recreations of historical cipher machines (Enigma, Lorenz, M-209,...) within ComfyUI. - (Planned) Audio spectrogram analysis tools. -- (Planned) Generative steganography models +- (Planned) ML-integrated steganography models. - (Far Future plans) Post-quantum cryptography tools. ## Testing methodology diff --git a/pyproject.toml b/pyproject.toml index 7592af6..2626aa5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,46 +1,44 @@ -[build-system] -requires = ["setuptools>=70.0"] -build-backend = "setuptools.build_meta" - [project] name = "arg-toolkit" version = "2.2.1" description = "A collection of ComfyUI nodes for cryptography, steganography, encryption, decryption, and tools needed to both make and solve ARGs (Alternative Reality Games) and secret messages." -authors = [ - {name = "AzelusLightvale", email = "NexusLightvale@protonmail.com"} -] readme = "README.md" -license = "GPL-3.0-only" requires-python = ">=3.10" -classifiers = ["Operating System :: OS Independent",] -dependencies = [ - "secretpy", - "cryptography", - "stegano", - "invisible-watermark", - "reedsolo", - "numpy", - "torch", - "pillow", +license = "GPL-3.0-only" +authors = [{ name = "AzelusLightvale", email = "NexusLightvale@protonmail.com" }] +classifiers = [ + "Operating System :: OS Independent", ] - - - -[project.optional-dependencies] -dev = [ - "bump-my-version", - "coverage", # testing - "mypy", # linting - "pre-commit", # runs linting on commit - "pytest", # testing - "ruff", # lintingclass +dependencies = [ + "cryptography", + "invisible-watermark", + "numpy", + "pillow", + "reedsolo", + "secretpy", + "stegano", + "torch", + "torchaudio", ] [project.urls] -Repository = "https://github.com/AzelusLightvale/ComfyUI-ARG-Toolkit" BugTracker = "https://github.com/AzelusLightvale/ComfyUI-ARG-Toolkit/issues" Documentation = "https://github.com/AzelusLightvale/ComfyUI-ARG-Toolkit/wiki" +Repository = "https://github.com/AzelusLightvale/ComfyUI-ARG-Toolkit" +[project.optional-dependencies] +dev = [ + "bump-my-version", + "coverage", # testing + "mypy", # linting + "pre-commit", # runs linting on commit + "pytest", # testing + "ruff", # lintingclass +] + +[build-system] +requires = ["setuptools>=70.0"] +build-backend = "setuptools.build_meta" [tool.comfy] PublisherId = "azeluslightvale" @@ -51,15 +49,6 @@ Repository = "https://github.com/AzelusLightvale/ComfyUI-ARG-Toolkit" includes = [] -[tool.setuptools.package-data] -"*" = ["*.*"] - -[tool.pytest.ini_options] -minversion = "8.0" -testpaths = [ - "tests", -] - [tool.mypy] files = "." @@ -68,6 +57,12 @@ strict = true warn_unreachable = true warn_no_return = true +[tool.pytest.ini_options] +minversion = "8.0" +testpaths = [ + "tests", +] + [tool.ruff] # extend-exclude = ["static", "ci/templates"] line-length = 140 @@ -77,12 +72,15 @@ target-version = "py39" # Add rules to ban exec/eval [tool.ruff.lint] select = [ - "S102", # exec-builtin - "S307", # eval-used - "W293", - "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. - # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f + "S102", # exec-builtin + "S307", # eval-used + "W293", + "F", # The "F" series in Ruff stands for "Pyflakes" rules, which catch various Python syntax errors and undefined names. + # See all rules here: https://docs.astral.sh/ruff/rules/#pyflakes-f ] [tool.ruff.lint.flake8-quotes] inline-quotes = "double" + +[tool.setuptools.package-data] +"*" = ["*.*"] diff --git a/requirements.txt b/requirements.txt index 67e49e9..1c60e12 100644 --- a/requirements.txt +++ b/requirements.txt @@ -57,5 +57,7 @@ torch==2.13.0 # via # arg-toolkit (pyproject.toml) # invisible-watermark +torchaudio==2.11.0 + # via arg-toolkit (pyproject.toml) typing-extensions==4.16.0 # via torch diff --git a/src/spectrogram.py b/src/spectrogram.py index 7e0fd95..00f5cae 100644 --- a/src/spectrogram.py +++ b/src/spectrogram.py @@ -1,4 +1,5 @@ import torch +import torchaudio.transforms as torch_trans import numpy as np @@ -20,9 +21,9 @@ class SpectrogramEncoder: "tooltip": "The audio file to generate the spectrogram from", }, ), - "window_size": ("INT", {"default": 25, "step": 1, "tooltip": "The size of each 'frame' in miliseconds."}), - "hop_size": ("INT", {"default": 15, "step": 1, "tooltip": "The size of the overlap between each 'frames' in miliseconds"}), - "windowing_shape": ( + "window_size": ("FLOAT", {"default": 25.000, "min": 0.001, "max": 100000.000, "step": 0.0011, "tooltip": "The size of each 'frame' in miliseconds."}), + "hop_size": ("FLOAT", {"default": 15.000, "min": 0.001, "max": 100000.000, "step": 0.001, "tooltip": "The size of the overlap between each 'frames' in miliseconds"}), + "windowing_func": ( ["Hanning", "Bartlett", "Blackman", "Hamming"], { "default": "Hanning", @@ -116,16 +117,16 @@ class SpectrogramEncoder: def execute( self, audio, - window_size, - hop_size, - windowing_shape, - db_norm, - colormap_low, - colormap_mid, - colormap_high, - cutoff_switch, - scaling_method, - active_db, + window_size: int, + hop_size: int, + windowing_func: str, + db_norm: float, + colormap_low: str, + colormap_mid: str, + colormap_high: str, + cutoff_switch: bool, + scaling_method: str, + active_db: float, ): # Defining the base variables for audio waveforms = audio["waveform"] @@ -136,15 +137,15 @@ class SpectrogramEncoder: # The framing and windowing process raw_frame_length = int(np.round((window_size / 1000.0) * sampling_rate)) hop_length = int(np.round((hop_size / 1000.0) * sampling_rate)) - frames = np.lib.stride_tricks.sliding_window_view(signal, window_shape=raw_frame_length)[::hop_length] - windowing_func = getattr(np, windowing_shape.lower().strip()) + frames = np.lib.stride_tricks.sliding_window_view(signal, window_shape=raw_frame_length, axis=-1)[::hop_length] + windowing_func = getattr(np, windowing_func.lower().strip()) windowed_frames = frames * windowing_func(raw_frame_length) # Zero-padding the frame to a power of two for FFT performance p2_fft = 2 ** int(np.ceil(np.log2(raw_frame_length))) # The FFT and spectrogram functions - complex_spectra = np.fft.rfft(windowed_frames, n=p2_fft, axis=-1) + complex_spectra = np.fft.rfft(windowed_frames, n=p2_fft) magnitude_spectra = np.abs(complex_spectra).T # Scaling function @@ -164,13 +165,28 @@ class SpectrogramEncoder: mel_num = int(np.round(mel_max / (2595.0 * np.log10(1.0 + hz_per_bin / 700.0)))) mel_num = int(np.clip(mel_num, 32, p2_fft // 2 + 1)) - mel_points = np.linspace(0, mel_max, mel_num) + mel_points = np.linspace(0, mel_max, mel_num +2) hz_points = 700.0 * (10.0 ** (mel_points / 2595.0) - 1.0) - freq_indices = np.clip((hz_points / nyquist * (freq_num - 1)).astype(int), 0, freq_num - 1) + freq_indices = (hz_points / nyquist) * (freq_num - 1) - db_spectrogram = magnitude_spectra[freq_indices, :] - db_spectrogram = 20 * np.log10(magnitude_spectra + 1e-10) + mel_matrix = np.zeros((mel_num, freq_num), dtype=np.float32) + for i in range(1, mel_num + 1): + left, center, right = freq_indices[i - 1], freq_indices[i], freq_indices[i + 1] + bins = np.arange(int(np.floor(left)), int(np.ceil(right)) + 1) + bins = bins[(bins >= 0) & (bins < freq_num)] + + up_mask = (bins >= left) & (bins <= center) + down_mask = (bins > center) & (bins <= right) + + if center > left: + mel_matrix[i - 1, bins[up_mask]] = (bins[up_mask] - left) / (center - left) + if right > center: + mel_matrix[i - 1, bins[down_mask]] = (right - bins[down_mask]) / (right - center) + + db_spectrogram = np.dot(mel_matrix, magnitude_spectra) + db_spectrogram = 20 * np.log10(db_spectrogram + 1e-10) db_spectrogram_norm = np.maximum(db_spectrogram, np.max(db_spectrogram) - db_norm) + # Colormap if cutoff_switch: colormap_db = self.crop_spectrogram_dimensions(db_spectrogram_norm, active_db) @@ -182,9 +198,173 @@ class SpectrogramEncoder: return (colormap_tensor,) +class SpectrogramDecoder: + CATEGORY = "ARG Toolkit/Steganography/Analysis" + + @classmethod + def __init__(self): + pass + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "image": ("IMAGE", {"forceInput": True,"tooltip": "The spectrogram image to decode into audio."}), + "image_type": ( + "BOOLEAN", + { + "label_on": "Spectrogram", + "label_off": "Image", + "default": True, + "tooltip": "The type of image fed into the model. `Spectrogram` flips the image to correct for orientation, `Image` does not.", + } + ), + "channel": ( + ["colormap", "channel_red", "channel_green", "channel_blue", "luminance"], + { + "default": "luminance", + "tooltip": "The color channel used to collapse and deconstruct the image. For spectrograms with a known colormap, use `colormap`, otherwise use `luminance` or any of the RGB channels.", + } + ), + "window_size": ("INT", {"default": 25, "step": 1, "tooltip": "The size of each 'frame' in miliseconds."}), + "hop_size": ("INT", {"default": 15, "step": 1, "tooltip": "The size of the overlap between each 'frames' in miliseconds"}), + "windowing_func": ( + ["Hanning", "Bartlett", "Blackman", "Hamming"], + { + "default": "Hanning", + "tooltip": "The windowing function to use. Unless you know what you're doing, leave this as default.", + }, + ), + "scaling_method": ( + ["Linear", "Quadratic", "Logarithmic", "Logarithmic Mel"], + { + "default": "Logarithmic", + "tooltip": "The scaling method used for the spectrogram.", + } + ), + "autogain": ( + "BOOLEAN", + { + "tooltip": "For Linear and Quadratic scaling method, dynamically scale the amplitude based on the peak. If disabled, use `max_amp` to define it instead. If used with Logarithmic Mel scaling, will enable area normalization if used with `mel_scale=slaney`.", + } + ), + "top_db": ( + "FLOAT", + { + "default": 80.0, + "min": 0.0, + "max": 255.0, + "tooltip": "The dynamic range (in negative dB) to scale against for both Logarithmic scaling methods.", + }, + ), + "sampling_rate": ("INT", {"max": 384000, "min":100, "default": 44100, "step": 1, "tooltip": "The target sample rate for the audio file."}), + "max_iter": ("INT", {"min": 1, "max": 1024, "default": 32, "tooltip": "Maximum iteration to run the Griffin-Lim algorithm for."}), + "gl_momentum": ("FLOAT", {"min": 0.00, "max": 2.00, "default": 0.99, "step": 0.01, "tooltip": "For use with Fast Griffin-Lim. 0 disables Fast Griffin-Lim."}) + }, + "optional": { + "colormap_low": ("STRING", {"default": "#000000", "tooltip": "The color of the lowest point of the spectrogram.",}), + "colormap_mid": ("STRING", {"default": "#B73779", "tooltip": "The color of the middle point of the spectrogram.",}), + "colormap_high": ("STRING", {"default": "#FCFDBF", "tooltip": "The color of the highest point of the spectrogram.",}), + "max_amp": ("FLOAT", {"step": 0.1, "min": 0.0, "max": 1.0, "default": 1.0, "tooltip": "If `autogain is disabled, this defines the maximum amplitude to scale with in Linear and Quadratic scaling.",}), + "mel_scale": ("BOOLEAN", {"label_off": "htk", "label_on": "slaney", "default": False, "tooltip": "Mel scaling method to use."}), + }, + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "execute" + + def colormap_collapse(self, image, colormap_low, colormap_mid, colormap_high): + c1 = torch.tensor([int(colormap_low[i : i + 2], 16) for i in (1, 3, 5)], dtype=torch.float32) / 255.0 + c2 = torch.tensor([int(colormap_mid[i : i + 2], 16) for i in (1, 3, 5)], dtype=torch.float32) / 255.0 + c3 = torch.tensor([int(colormap_high[i : i + 2], 16) for i in (1, 3, 5)], dtype=torch.float32) / 255.0 + + vec12, vec23 = c2 - c1, c3 - c2 + len12_sq, len23_sq = torch.sum(vec12**2), torch.sum(vec23**2) + + diff1 = image - c1 + t1 = torch.clamp(torch.sum(diff1 * vec12, dim=-1) / (len12_sq + 1e-8), 0.0, 1.0) * 0.5 + + diff2 = image - c2 + t2 = torch.clamp(torch.sum(diff2 * vec23, dim=-1) / (len23_sq + 1e-8), 0.0, 1.0) * 0.5 + 0.5 + + dist1 = torch.norm(image - (c1 + t1.unsqueeze(-1) / 0.5 * vec12), dim=-1) + dist2 = torch.norm(image - (c2 + (t2.unsqueeze(-1) - 0.5) / 0.5 * vec23), dim=-1) + + norm_val = torch.where(dist1 < dist2, t1, t2) + return norm_val + + def execute(self, image, image_type, window_size, hop_size, windowing_func, autogain, max_amp, max_iter, top_db, sampling_rate, gl_momentum, scaling_method, mel_scale, channel, colormap_low, colormap_mid, colormap_high): + # Rotate the image in case it's a spectrogram + if image_type: + image = torch.flip(image, dims=[1]) + elif not image_type: + image = image + + raw_frame_length = int(np.round((window_size / 1000.0) * sampling_rate)) + hop_length = int(np.round((hop_size / 1000.0) * sampling_rate)) + p2_fft = 2 ** int(np.ceil(np.log2(raw_frame_length))) + n_stft = (p2_fft//2)+1 + + + # Channel collapse + if channel == "colormap": + norm = self.colormap_collapse(image, colormap_low, colormap_mid, colormap_high) + elif channel == "luminance": + weights = torch.tensor([0.299, 0.587, 0.114]) # Luminance value based on ITU-R BT.601 perceptual weights. + norm = torch.sum(image * weights, dim=-1) + elif channel == "channel_red": + norm = image[..., 0] + elif channel == "channel_green": + norm = image[..., 1] + elif channel == "channel_blue": + norm = image[..., 2] + + # Scaling the array + if scaling_method in ["Linear", "Quadratic"]: + if scaling_method == "Linear": + image_peak = torch.max(norm) + elif scaling_method == "Quadratic": + image_peak = torch.max(norm**2) + + if autogain and image_peak>0: + magnitude = (norm/image_peak) * 1.0 + else: + magnitude = norm * max_amp + elif scaling_method == "Logarithmic": + log_mag = torch.pow(10.0, ((norm * top_db) - top_db)/20) + magnitude = torch.clamp(log_mag, min=0.0) + elif scaling_method == "Logarithmic Mel": + log_mag = torch.pow(10.0, ((norm * top_db) - top_db)/20) + lin_mag = torch.clamp(log_mag, min=0.0) + inverse_scale = torch_trans.InverseMelScale(n_stft=n_stft, n_mels=norm.shape[1], sample_rate=sampling_rate, f_min=0.0, f_max=sampling_rate/2, norm="slaney" if autogain==True else None, mel_scale="htk" if mel_scale==False else "slaney") + magnitude = inverse_scale(lin_mag) + + # Griffin-Lim + windowing = { + "Hanning": torch.hann_window, + "Hamming": torch.hamming_window, + "Blackman": torch.blackman_window, + "Bartlett": torch.bartlett_window, + } + windowing_func = windowing.get(windowing_func) + + griffin_lim = torch_trans.GriffinLim(n_fft=p2_fft, win_length=raw_frame_length, hop_length=hop_length, window_fn=windowing_func, n_iter=max_iter, momentum=gl_momentum, power=1.0) + waveform = griffin_lim(magnitude) + max_val = torch.max(torch.abs(waveform)) + if max_val > 1e-8 : + waveform = waveform / max_val + waveform = waveform.unsqueeze(0).cpu() + final_audio = {"waveform": waveform, "sample_rate": sampling_rate} + + return (final_audio,) + + NODE_CLASS_MAPPINGS = { "SpectrogramEncoder": SpectrogramEncoder, + "SpectrogramDecoder": SpectrogramDecoder, } NODE_DISPLAY_NAME_MAPPINGS = { "SpectrogramEncoder": "Audio to Spectrogram", + "SpectrogramDecoder": "Spectrogram to Audio", }