From c094cbba281c605f95f42f5ccb277210d5e97851 Mon Sep 17 00:00:00 2001 From: Willie G <111990299+GeekyGhost@users.noreply.github.com> Date: Sat, 15 Mar 2025 02:35:49 -0400 Subject: [PATCH] Update nodes.py Added batch node, audio out on main node, fixed some things, improved a few scripts, improved functionality and quality a bit. --- nodes.py | 465 ++++++++++++++++++++++++++++++++++++++++--------------- 1 file changed, 338 insertions(+), 127 deletions(-) diff --git a/nodes.py b/nodes.py index 1515bff..8ca9054 100644 --- a/nodes.py +++ b/nodes.py @@ -9,7 +9,6 @@ import torch import torchaudio from pedalboard.io import AudioFile import glob -import wave import platform MAX_BATCH_SIZE = 8 @@ -23,168 +22,380 @@ class RiffusionNode: return { "required": { 'spectrogram': ("IMAGE",), - 'filename_prefix': ("STRING", {"default": "audio\Riffusion"}), - }, + "optional": { + 'filename_prefix': ("STRING", {"default": "Riffusion_"}), + 'save_to_file': ("BOOLEAN", {"default": False}), + 'max_volume': ("FLOAT", {"default": 50.0, "min": 10.0, "max": 100.0, "step": 5.0}), + 'quality_level': ("INT", {"default": 2, "min": 0, "max": 3, "step": 1}), + 'apply_filter': ("BOOLEAN", {"default": True}), + } } - RETURN_TYPES = () + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) FUNCTION = "Process_Riffusion" CATEGORY = "audio" - OUTPUT_NODE = True def spectrogram_from_image(self, image, max_volume=50, power_for_image=0.25) -> np.ndarray: - # Convert the image to a numpy array + """ + Compute a spectrogram magnitude array from a spectrogram image. + """ + # Convert to a numpy array of floats data = np.array(image).astype(np.float32) - # Reverse the image vertically - data = data[::-1,:,0] - # Invert the image - data = 255-data - # Normalize the image - data = data*max_volume/255 - # Invert power curve - data = np.power(data, 1/power_for_image) + + # Flip Y take a single channel + data = data[::-1, :, 0] + + # Invert + data = 255 - data + + # Rescale to max volume + data = data * max_volume / 255 + + # Reverse the power curve + data = np.power(data, 1 / power_for_image) + return data - def waveform_from_spectrogram(self, Sxx: np.ndarray, n_fft: int, hop_length: int, win_length: int, num_samples: int, - sample_rate: int, mel_scale: bool = True, nmels: int = 512, max_mel_iters: int = 200, num_griffin_lim_iters: int = 32, - device: str = platform.system() == "Darwin" and "cpu" or "cuda" + def waveform_from_spectrogram( + self, + Sxx: np.ndarray, + n_fft: int, + hop_length: int, + win_length: int, + num_samples: int, + sample_rate: int, + mel_scale: bool = True, + n_mels: int = 512, + num_griffin_lim_iters: int = 32, + device: str = None, + quality_level: int = 2 ) -> np.ndarray: + """ + Reconstruct a waveform from a spectrogram using the Griffin-Lim algorithm. + Quality level controls the number of iterations: + 0 = fast/low quality (16 iterations) + 1 = balanced (32 iterations) + 2 = good (48 iterations) + 3 = best (64 iterations) + """ + # Adjust iterations based on quality level + iterations_map = {0: 16, 1: 32, 2: 48, 3: 64} + num_griffin_lim_iters = iterations_map.get(quality_level, 32) + + # Select appropriate device + if device is None: + device = "cuda" if torch.cuda.is_available() else "cpu" + if platform.system() == "Darwin": # macOS + device = "cpu" # Force CPU on macOS + + # Convert to torch tensor and move to device Sxx_torch = torch.from_numpy(Sxx).to(device) if mel_scale: - mel_inv_scaler = torchaudio.transforms.InverseMelScale(n_mels=nmels, sample_rate=sample_rate, f_min=0, f_max=10000, - n_stft=n_fft // 2 + 1, norm=None, mel_scale="htk").to(device) - Sxx_torch = mel_inv_scaler(Sxx_torch) + try: + # Fix dimensions for mel inverse scaling + if Sxx_torch.ndim == 2: + Sxx_torch = Sxx_torch.unsqueeze(0) # Add batch dimension + + mel_inv_scaler = torchaudio.transforms.InverseMelScale( + n_mels=n_mels, + sample_rate=sample_rate, + f_min=0, + f_max=10000, + n_stft=n_fft // 2 + 1, + norm=None, + mel_scale="htk" + ).to(device) + + # Ensure we have the right shape for InverseMelScale + if Sxx_torch.shape[-2] != n_mels: + print(f"Warning: Reshaping Sxx_torch from {Sxx_torch.shape} to fit n_mels={n_mels}") + # If the shape doesn't match, try a simpler approach + mel_scale = False + else: + Sxx_torch = mel_inv_scaler(Sxx_torch) + except Exception as e: + print(f"Error in mel scale inverse transformation: {e}") + mel_scale = False + print("Continuing without mel scale transformation") - griffin_lim = torchaudio.transforms.GriffinLim(n_fft=n_fft, hop_length=hop_length, win_length=win_length, power=1.0, - n_iter=num_griffin_lim_iters).to(device) + try: + # If not using mel scale, ensure dimensions match n_fft + if not mel_scale: + # Need to resize to match expected STFT dimensions + if Sxx_torch.ndim == 2: + Sxx_torch = Sxx_torch.unsqueeze(0) # Add batch dimension + + # Get the expected frequency dimension size + freq_dim = n_fft // 2 + 1 + + if Sxx_torch.shape[-2] != freq_dim: + print(f"Reshaping frequency dimension from {Sxx_torch.shape[-2]} to {freq_dim}") + # Use interpolation to resize + if hasattr(torch.nn.functional, 'interpolate'): + # For newer PyTorch versions + Sxx_resized = torch.nn.functional.interpolate( + Sxx_torch.permute(0, 2, 1), # [batch, time, freq] -> [batch, freq, time] + size=freq_dim, + mode='linear', + align_corners=False + ).permute(0, 2, 1) # Back to [batch, time, freq] + Sxx_torch = Sxx_resized + else: + # Manual resize using numpy (fallback) + Sxx_np = Sxx_torch.cpu().numpy() + Sxx_resized = np.zeros((Sxx_np.shape[0], freq_dim, Sxx_np.shape[2]), dtype=Sxx_np.dtype) + for b in range(Sxx_np.shape[0]): + for t in range(Sxx_np.shape[2]): + Sxx_resized[b, :, t] = np.interp( + np.linspace(0, 1, freq_dim), + np.linspace(0, 1, Sxx_np.shape[1]), + Sxx_np[b, :, t] + ) + Sxx_torch = torch.from_numpy(Sxx_resized).to(device) + + print(f"Griffin-Lim input shape: {Sxx_torch.shape}, n_fft: {n_fft}, expected: {n_fft//2+1}") + print(f"Using {num_griffin_lim_iters} Griffin-Lim iterations (quality level {quality_level})") + + griffin_lim = torchaudio.transforms.GriffinLim( + n_fft=n_fft, + win_length=win_length, + hop_length=hop_length, + power=1.0, + n_iter=num_griffin_lim_iters, + momentum=0.99, # Higher momentum helps convergence quality + rand_init=False # Deterministic initialization + ).to(device) + + waveform = griffin_lim(Sxx_torch).cpu().numpy() + + # If we got multiple waveforms (batch dimension > 1), just take the first + if waveform.ndim > 1: + waveform = waveform[0] + + return waveform + except Exception as e: + print(f"Error in Griffin-Lim algorithm: {e}") + # Return zeros if Griffin-Lim fails + print("Returning empty waveform") + return np.zeros(num_samples, dtype=np.float32) - waveform = griffin_lim(Sxx_torch).cpu().numpy() + def apply_audio_filter(self, samples, sample_rate): + """Apply some basic audio filtering to improve audio quality""" + try: + from pedalboard import Pedalboard, LowpassFilter, HighpassFilter, Compressor + + # Convert to the format pedalboard expects + if samples.ndim == 1: + samples = samples.reshape(1, -1) # mono -> [1, samples] + + # Create a pedalboard with some basic filters + board = Pedalboard([ + HighpassFilter(cutoff_frequency_hz=60), # Remove rumble + LowpassFilter(cutoff_frequency_hz=16000), # Remove high-frequency noise + Compressor(threshold_db=-20, ratio=3.0) # Smooth out dynamics + ]) + + # Apply the pedalboard + processed_samples = board(samples, sample_rate) + + # Convert back to the original shape + if samples.ndim == 2 and samples.shape[0] == 1: + processed_samples = processed_samples.reshape(-1) + + return processed_samples + except Exception as e: + print(f"Error applying audio filter: {e}") + return samples - return waveform - - - def get_wave_bytes_from_spectrogram(self, spec) -> T.Tuple[io.BytesIO, float]: - max_volume = 50 + def get_audio_from_spectrogram(self, spec, max_volume=50, quality_level=2, apply_filter=True): + """Convert spectrogram to audio waveform and sample rate""" power_for_image = 0.25 Sxx = self.spectrogram_from_image(spec, max_volume, power_for_image) - sample_rate = 44100 - clip_duration_ms = 5000 - + sample_rate = 44100 # [Hz] + clip_duration_ms = 5000 # [ms] + bins_per_image = spec.height n_mels = spec.height - - #FFT parameters - window_duration_ms = 100 - padded_duration_ms = 400 - step_size_ms = 10 - - #Derived Parameters - num_samples = ( - int(spec.width/float(bins_per_image) * clip_duration_ms) * sample_rate + + # FFT parameters - adjusted for better quality + window_duration_ms = 100 # [ms] + padded_duration_ms = 400 # [ms] + step_size_ms = 10 # [ms] + + # Derived parameters + num_samples = int(spec.width / float(bins_per_image) * clip_duration_ms * sample_rate / 1000) + n_fft = int(padded_duration_ms / 1000.0 * sample_rate) + hop_length = int(step_size_ms / 1000.0 * sample_rate) + win_length = int(window_duration_ms / 1000.0 * sample_rate) + + # Use CUDA if available + device = "cuda" if torch.cuda.is_available() else "cpu" + if platform.system() == "Darwin": # macOS + device = "cpu" # Force CPU on macOS + + print(f"Generating audio with max_volume={max_volume}, quality_level={quality_level}") + print(f"Using device: {device}") + print(f"Shapes - n_fft: {n_fft}, hop_length: {hop_length}, win_length: {win_length}, Sxx: {Sxx.shape}") + + # Generate the waveform + samples = self.waveform_from_spectrogram( + Sxx=Sxx, + n_fft=n_fft, + hop_length=hop_length, + win_length=win_length, + num_samples=num_samples, + sample_rate=sample_rate, + mel_scale=True, + n_mels=n_mels, + quality_level=quality_level, + device=device ) - n_fft = int(padded_duration_ms / 1000 * sample_rate) - hop_length = int(step_size_ms / 1000 * sample_rate) - win_length = int(window_duration_ms / 1000 * sample_rate) - - #Generate the waveform - samples = self.waveform_from_spectrogram(Sxx=Sxx, n_fft=n_fft, hop_length=hop_length, win_length=win_length, - num_samples=num_samples, sample_rate=sample_rate, mel_scale=True, nmels=n_mels, max_mel_iters=200, - num_griffin_lim_iters=32) - - #Encode samples and return to calling function - wav_bytes = io.BytesIO() - wavfile.write(wav_bytes, sample_rate, samples.astype(np.int16)) - wav_bytes.seek(0) - - duration_s = float(len(samples))/sample_rate - - return wav_bytes, duration_s - - def process_wav(self, wav_file): - with AudioFile(wav_file) as f: - audio = f.read(f.frames) - sample_rate = f.samplerate - - filename = wav_file.replace(".wav", ".mp3") - - with AudioFile(filename, "w", sample_rate, audio.shape[0]) as f: - f.write(audio) - - def Process_Riffusion(self, spectrogram, filename_prefix): - # Get the current file's directory and go up two levels - current_dir = Path(__file__).parent - output_dir = current_dir.parent.parent / 'output' - results = list() + # Apply audio filtering if requested + if apply_filter: + samples = self.apply_audio_filter(samples, sample_rate) - # Create the output directory if it doesn't exist - output_dir.mkdir(parents=True, exist_ok=True) - - # Search output_dir for files with the filename prefix - list_of_files = glob.glob(str(output_dir / f"{filename_prefix}*")) + # Apply light normalization to prevent clipping + if np.max(np.abs(samples)) > 0: + normalization_factor = 0.95 / np.max(np.abs(samples)) + samples = samples * normalization_factor + + return samples, sample_rate - # Get the highest number in the list of files - if not list_of_files: - new_number = 0 - else: - max_number = 0 - for file in list_of_files: - number = int(file.split("_")[-2]) - max_number = max(max_number, number) - # Increment the number and add it to the filename and add the extension - new_number = max_number + 1 - - new_filename = f"{filename_prefix}_{new_number:05}_.wav" - - # now load the spectrogram and convert it to a audio file - # TODO: Add support for multiple images (as input) to generate a single audio file for longer audio lengths - # TODO: Add support for multiple images to generate multiple audio files - MAX_FILES = 1 + def Process_Riffusion(self, spectrogram, filename_prefix="Riffusion_", save_to_file=False, + max_volume=50.0, quality_level=2, apply_filter=True): + """Convert a spectrogram image to audio""" + # Process first spectrogram in batch spec = None for (batch_number, image) in enumerate(spectrogram): - print("Batch Number ", batch_number) - if batch_number >= MAX_FILES: + if batch_number >= 1: # Process only the first image break i = 255. * image.cpu().numpy() spec = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8)) - audio, duration = self.get_wave_bytes_from_spectrogram(spec) - print(f"Duration of {new_filename}: {duration}") - - # get full output path - full_output_path = os.path.join(output_dir, new_filename) - - # save the audio - with open(full_output_path, 'wb') as f: - f.write(audio.getbuffer()) - - # TODO: add support for other audio types (wav, flac, mp3, ogg, aac) - # deprecated until I figure out how to save as other audio types - # Audio file formats to use: - # .wav - # .flac - # .mp3 - # .ogg - # .aac - #self.process_wav(full_output_path) - - results.append({ - "filename": new_filename, - "subfolder": str(output_dir), - "type": "output" - }) - - return { "ui": { "audio": results } } + if spec is None: + print("Warning: No valid spectrogram found in batch") + # Return empty audio + empty_audio = { + "waveform": torch.zeros((1, 1, 1000), dtype=torch.float32), + "sample_rate": 44100 + } + return (empty_audio,) + # Convert spectrogram to audio with parameters + samples, sample_rate = self.get_audio_from_spectrogram( + spec, + max_volume=float(max_volume), + quality_level=int(quality_level), + apply_filter=apply_filter + ) + + # Format waveform for ComfyUI's AUDIO type + samples = samples.astype(np.float32) + + # Create waveform tensor with shape [1, 1, samples] + waveform = torch.from_numpy(samples).float().unsqueeze(0).unsqueeze(0) + + # ComfyUI expects a dictionary with waveform and sample_rate + audio_output = { + "waveform": waveform, + "sample_rate": sample_rate + } + + # Save to disk if requested + if save_to_file: + try: + # Get the current file's directory and go up two levels to output directory + current_dir = Path(__file__).parent + output_dir = current_dir.parent.parent / 'output' + + # Create the output directory if it doesn't exist + output_dir.mkdir(parents=True, exist_ok=True) + + # Search output_dir for files with the filename prefix + list_of_files = glob.glob(str(output_dir / f"{filename_prefix}*")) + + # Get the highest number in the list of files + max_number = 0 + for file in list_of_files: + try: + number = int(file.split("_")[-1].split(".")[0]) + max_number = max(max_number, number) + except (ValueError, IndexError): + continue + + # Create new filename with incremented number + new_number = max_number + 1 + new_filename = output_dir / f"{filename_prefix}{new_number}.wav" + + # Save the audio file + wavfile.write(str(new_filename), sample_rate, (samples * 32767).astype(np.int16)) + print(f"Saved audio to {new_filename}") + + # Optionally convert to MP3 + try: + mp3_filename = str(new_filename).replace(".wav", ".mp3") + + with AudioFile(new_filename) as f: + audio_data = f.read(f.frames) + sr = f.samplerate + + with AudioFile(mp3_filename, 'w', sr, audio_data.shape[0]) as f: + f.write(audio_data) + + print(f"Converted to MP3: {mp3_filename}") + + except Exception as e: + print(f"Error converting to MP3: {e}") + + except Exception as e: + print(f"Error saving audio file: {e}") + + return (audio_output,) + + +class RiffusionToBatchNode: + """Node to convert a single audio to a batch for further processing""" + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO",), + "batch_size": ("INT", {"default": 1, "min": 1, "max": 64, "step": 1}), + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio_batch",) + FUNCTION = "create_batch" + CATEGORY = "audio" + + def create_batch(self, audio, batch_size=1): + """Create a batch of copies of the input audio""" + waveform = audio["waveform"] + sample_rate = audio["sample_rate"] + + # If batch_size is 1, just return the original + if batch_size <= 1: + return (audio,) + + # Create batch of identical audio + batch = waveform.repeat(batch_size, 1, 1) + + return ({"waveform": batch, "sample_rate": sample_rate},) NODE_CLASS_MAPPINGS = { - "RiffusionNode": RiffusionNode + "RiffusionNode": RiffusionNode, + "RiffusionToBatchNode": RiffusionToBatchNode } NODE_DISPLAY_NAME_MAPPINGS = { - "RiffusionNode": "Riffusion" -} \ No newline at end of file + "RiffusionNode": "Riffusion Audio Generator", + "RiffusionToBatchNode": "Riffusion Audio Batch" +}