diff --git a/README.md b/README.md index 601bc12..15b8471 100644 --- a/README.md +++ b/README.md @@ -1,297 +1,73 @@ # Network Bending for ComfyUI -A custom node pack for ComfyUI that enables creative manipulation and "bending" of neural network models. Perform various operations on loaded model checkpoints to create unique and experimental effects. +A custom node pack for ComfyUI that enables the creative manipulation of generative models. Allows you to corrupt, modify, and blend neural networks. ## Features -### 🎛️ Network Bending Node -The main node for applying various network bending operations: - - **Add Noise**: Inject Gaussian noise into model weights -- **Scale Weights**: Scale model weights up or down -- **Prune Weights**: Remove small weights (sparsification) +- **Scale Weights**: Scale model weights up or down to amplify or dampen features +- **Prune Weights**: Remove small weights for sparsification - **Randomize Weights**: Replace portions of weights with random values - **Smooth Weights**: Apply spatial smoothing to weight matrices - **Quantize Weights**: Reduce weight precision to discrete levels - -### 🔧 Network Bending Advanced Node -Advanced operations for more complex manipulations: - - **Layer Swap**: Swap weights between similar layers - **Activation Replace**: Replace activation functions - **Weight Transpose**: Transpose weight matrices -- **Channel Shuffle**: Randomly shuffle channels in conv layers +- **Channel Shuffle**: Randomly shuffle channels in convolutional layers - **Frequency Filter**: Apply frequency domain filtering - **Weight Clustering**: Group similar weights together -### 🎨 Model Mixer Node Mix two models together with various blending modes: - - **Linear Interpolation**: Simple weighted average of weights - **Weighted Sum**: Weighted sum with normalization - **Layer-wise Mix**: Different mix ratios for different layers - **Frequency Blend**: Mix models in frequency domain - **Random Mix**: Randomly select weights from either model -### 🔄 Latent Format Converter Node -Convert between audio and image latent formats to enable cross-modal workflows: +### Audio Processing Nodes -- **Audio to Image**: Convert 3D audio latents [B,C,L] to 4D image latents [B,C,H,W] -- **Image to Audio**: Convert 4D image latents [B,C,H,W] to 3D audio latents [B,C,L] -- **Auto Detect**: Automatically detect format and convert appropriately -- **Channel Conversion**: Convert between different channel counts (e.g., 4→64 for SD→StableAudio) -- **Multiple Methods**: Choose from reshape, interpolate, fold, or tile methods for both spatial and channel dims +The package also includes comprehensive audio processing capabilities for Stable Audio models: -### 🌀 VAE Network Bending Node -Apply network bending operations specifically to VAE models: +#### Audio VAE Nodes +- **Audio VAE Encode**: Encode audio waveforms into latent space +- **Audio VAE Decode**: Decode latent representations back to audio -- **Add Noise**: Inject noise into encoder/decoder weights -- **Scale Weights**: Scale VAE weights for more/less aggressive encoding -- **Asymmetric Noise**: Different noise levels for encoder vs decoder -- **Channel Corruption**: Corrupt specific channels in the VAE -- **Progressive Corruption**: Gradually increase corruption through layers -- **Target Components**: Choose encoder, decoder, or both +#### Audio Latent Manipulation +- **Audio Latent Interpolate**: Smoothly interpolate between audio latents +- **Audio Latent Blend**: Blend multiple audio latents with various modes +- **Audio Latent Manipulator**: Direct manipulation of audio latent features +- **Audio Feature Extractor**: Extract specific features from audio latents -### 🎯 VAE Mixer Node -Mix two VAE models with specialized blending strategies: - -- **Linear Blend**: Blend both encoder and decoder linearly -- **Encoder Swap**: Use encoder from one VAE with decoder from another -- **Decoder Swap**: Use decoder from one VAE with encoder from another -- **Cross Architecture**: Mix encoder/decoder with different ratios -- **Layer Shuffle**: Randomly select layers from each VAE - -### 💫 VAE Latent Bending Node -Perform operations directly on VAE latent streams: - -- **Add Noise**: Add noise to latent values for texture variations -- **Channel Swap**: Swap latent channels for color shifts -- **Frequency Filter**: Apply frequency domain filtering -- **Spatial Corruption**: Corrupt specific spatial regions -- **Value Quantization**: Create posterization effects -- **Dimension Warp**: Warp spatial dimensions for distortions - -### 🔮 VAE Channel Manipulator Node -Advanced channel-level operations on VAE latents: - -- **Channel Attention**: Apply self-attention across channels -- **Cross Channel Mixing**: Mix information between channels -- **Channel Dropout**: Randomly drop channels -- **Channel Amplification**: Amplify important channels -- **Channel Rotation**: Rotate channels in feature space -- **Channel Statistics Swap**: Swap statistics with reference latent +#### Audio Style Transfer +- **Audio Style Transfer**: Transfer style from one audio to another +- **Audio Latent Guidance**: Guide generation with reference audio +- **Audio Reference Encoder**: Encode reference audio for style transfer ## Installation +### Basic Installation + 1. Clone this repository into your ComfyUI `custom_nodes` directory: ```bash cd ComfyUI/custom_nodes -git clone https://github.com/yourusername/network_bending.git +git clone https://github.com/DavidPiazza/network_bending.git ``` -2. Install dependencies (if any): -```bash -cd network_bending -pip install -r requirements.txt # if requirements.txt exists -``` +2. Restart ComfyUI -3. Restart ComfyUI +### Audio Network Bending -## Usage +1. **Load an audio model** (e.g., Stable Audio VAE) +2. **Add audio processing nodes** from the network_bending/audio category +3. **Connect audio inputs** either from file loaders or generated audio +4. **Apply bending operations** to the audio latents +5. **Decode back to audio** using Audio VAE Decode -### Basic Network Bending - -1. Load a model checkpoint using the standard ComfyUI loader nodes -2. Add a "Network Bending" node to your workflow -3. Connect the MODEL output to the Network Bending input -4. Configure parameters: - - **Operation**: Select the type of bending operation - - **Intensity**: Control the strength (0.0 = no effect, 1.0 = maximum) - - **Target Layers**: Specify which layers to modify (e.g., "conv", "attention", or "all") - - **Seed**: Set for reproducible results (-1 for random) - -### Layer Targeting - -You can target specific layers using patterns: -- `all` - Modify all layers -- `conv` - Target convolutional layers -- `attention` - Target attention layers -- `linear` - Target linear/dense layers -- `norm` - Target normalization layers -- `embedding` - Target embedding layers -- Combine patterns with commas: `conv,attention` - -### Model Mixing - -1. Load two model checkpoints -2. Add a "Model Mixer" node -3. Connect both models to the mixer inputs -4. Set the mix mode and ratio -5. Connect the output to your generation pipeline - -### Latent Format Conversion - -1. When you encounter dimension mismatch errors between audio and image models -2. Add a "Latent Format Converter" node between the latent source and destination -3. Set conversion mode (or use auto_detect) -4. Choose appropriate reshape method (fold works well for audio→image) -5. Adjust target dimensions if converting to image format -6. Set target_channels to match your model (4 for SD/SDXL, 64 for StableAudio) -7. Use "project" channel mode for best results when converting channels - -### VAE Network Bending - -1. Load a VAE or get one from a checkpoint loader -2. Add a "VAE Network Bending" node -3. Connect the VAE to the input -4. Choose operation and target component (encoder/decoder/both) -5. Set intensity and preserve_mean option -6. Use the bent VAE for encoding/decoding - -### VAE Latent Manipulation - -1. Encode an image to latent space using VAE -2. Add "VAE Latent Bending" or "VAE Channel Manipulator" nodes -3. Connect latents through the manipulation nodes -4. Experiment with different operations and intensities -5. Decode back to image to see the effects - -## Parameters - -### Network Bending Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| operation | Select | Type of bending operation | See operations list | -| intensity | Float | Strength of the operation | 0.0 - 1.0 | -| target_layers | String | Layer patterns to target | Comma-separated patterns | -| seed | Int | Random seed | -1 to 2^63-1 | - -### Model Mixer Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| mix_mode | Select | Blending mode | See mix modes list | -| mix_ratio | Float | Blend ratio | 0.0 - 1.0 | - -### Latent Format Converter Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| conversion_mode | Select | Conversion direction | audio_to_image, image_to_audio, auto_detect | -| target_height | Int | Height for image format | 8 - 512 (steps of 8) | -| target_width | Int | Width for image format | 8 - 512 (steps of 8) | -| reshape_method | Select | Method for reshaping | reshape, interpolate, fold, tile | -| target_channels | Int | Target channel count | -1 to 256 (-1 keeps original) | -| channel_mode | Select | Channel conversion method | project, pad, interpolate, tile | - -### VAE Network Bending Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| operation | Select | VAE bending operation | See VAE operations list | -| intensity | Float | Strength of the operation | 0.0 - 1.0 | -| target_component | Select | Part of VAE to modify | both, encoder, decoder, latent_layers | -| preserve_mean | Boolean | Preserve weight means | True/False | -| seed | Int | Random seed | -1 to 2^63-1 | - -### VAE Mixer Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| mix_mode | Select | VAE blending mode | See VAE mix modes list | -| mix_ratio | Float | Blend ratio | 0.0 - 1.0 | - -### VAE Latent Bending Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| operation | Select | Latent operation | See latent operations list | -| intensity | Float | Operation strength | 0.0 - 1.0 | -| channel_specific | Boolean | Target specific channels | True/False | -| target_channels | String | Channel indices | Comma-separated (e.g., "0,1,2,3") | - -### VAE Channel Manipulator Node - -| Parameter | Type | Description | Range | -|-----------|------|-------------|-------| -| operation | Select | Channel operation | See channel operations list | -| intensity | Float | Operation strength | 0.0 - 1.0 | -| preserve_energy | Boolean | Preserve signal energy | True/False | -| reference_latent | Latent | Reference for some ops | Optional latent input | - -## UI Features - -- **Visual Feedback**: Operations display information about modified layers -- **Help Buttons**: Click "Layer Patterns Help" or "Operation Info" for detailed information -- **Color Coding**: Different node types have distinct colors for easy identification -- **Real-time Updates**: See which layers were modified after each operation - -## Examples - -### Subtle Model Variation -- Operation: `add_noise` -- Intensity: `0.05` -- Target Layers: `attention` -- Creates subtle variations while preserving model behavior - -### Aggressive Glitch Art -- Operation: `randomize_weights` -- Intensity: `0.3` -- Target Layers: `conv` -- Creates glitchy, unpredictable outputs - -### Model Hybridization -- Use Model Mixer with `linear_interpolation` -- Mix Ratio: `0.5` -- Creates a balanced hybrid of two models - -### Cross-Modal Generation -- Use Latent Format Converter with `auto_detect` -- Reshape Method: `fold` (for audio→image) -- Enables using audio latents with image models - -### VAE Glitch Effects -- VAE Network Bending: `channel_corruption` -- Target Component: `decoder` -- Intensity: `0.2` -- Creates colorful decoding artifacts - -### Dreamy VAE Processing -- VAE Latent Bending: `frequency_filter` -- Intensity: `0.4` -- Creates soft, dream-like images - -### Style Transfer via VAE -- VAE Channel Manipulator: `channel_statistics_swap` -- Intensity: `0.7` -- Use reference latent from style image -- Transfers color characteristics between images - -## Tips - -1. **Start with low intensity values** (0.01-0.1) and increase gradually -2. **Target specific layers** for more controlled effects -3. **Use seeds** for reproducible results -4. **Combine operations** by chaining multiple bending nodes -5. **Save interesting results** as new model checkpoints - -## Troubleshooting - -- **No effect visible**: Increase intensity or check target_layers pattern -- **Model breaks completely**: Reduce intensity or target fewer layers -- **Out of memory**: Operations clone the model, ensure sufficient VRAM - -## Contributing - -Contributions are welcome! Please feel free to submit pull requests or open issues for bugs and feature requests. +#### Cross-Modal Bending +Apply image model bending techniques to audio models or vice versa for experimental results. ## License -This project is licensed under the MIT License - see the LICENSE file for details. - -## Acknowledgments - -- ComfyUI team for the excellent framework -- Inspired by circuit bending and glitch art techniques +This project is licensed under the GNU General Public License v3 (GPL-3.0) - see the LICENSE file for details. +## diff --git a/src/network_bending/audio_nodes/README.md b/src/network_bending/audio_nodes/README.md new file mode 100644 index 0000000..7bddf0a --- /dev/null +++ b/src/network_bending/audio_nodes/README.md @@ -0,0 +1,200 @@ +# Audio Latent Space Manipulation Nodes for Stable Audio + +This collection of nodes enables sophisticated audio manipulation and style transfer by working directly in the latent space of Stable Audio's VAE. This approach allows for fine-grained control over audio generation without relying solely on text prompts. + +## Overview + +The latent space approach offers several advantages: +- **Direct audio-to-audio transformations** without text bottlenecks +- **Style transfer** between different audio sources +- **Hybrid sound generation** by blending multiple audio characteristics +- **Feature-based manipulation** using audio analysis +- **Real-time guidance** during the diffusion process + +## Node Categories + +### 1. Basic Encoding/Decoding + +#### AudioVAEEncode +Encodes audio waveforms into Stable Audio's latent space. +- **Inputs**: Audio waveform, VAE model +- **Parameters**: Normalize (bool), target sample rate +- **Outputs**: Audio latent, latent info + +#### AudioVAEDecode +Decodes audio latents back to waveforms. +- **Inputs**: Audio latent, VAE model +- **Parameters**: Denormalize (bool) +- **Outputs**: Audio waveform + +### 2. Latent Manipulation + +#### AudioLatentInterpolate +Smoothly interpolate between two audio latents. +- **Modes**: Linear, spherical (SLERP), cubic, sine +- **Use cases**: Morphing between sounds, creating transitions + +#### AudioLatentBlend +Blend multiple audio latents with various blend modes. +- **Modes**: Add, multiply, screen, overlay, soft light, difference, exclusion +- **Supports**: Up to 4 audio inputs with individual weights + +#### AudioLatentManipulator +Apply various transformations to audio latents. +- **Operations**: + - Frequency shift: Shift frequency content + - Temporal stretch: Time-stretch without pitch change + - Harmonic emphasis: Enhance harmonic content + - Noise injection: Add controlled noise + - Dynamic range: Compress/expand dynamics + - Spatial transform: Rotate in latent space + - Resonance filter: Apply resonant filtering + +### 3. Audio Analysis + +#### AudioFeatureExtractor +Extract audio features for guided manipulation. +- **Features**: Spectral centroid, rolloff, zero-crossing rate, MFCC, bandwidth, RMS energy, tempo, onset strength +- **Use**: Guide latent manipulations based on audio characteristics + +### 4. Style Transfer + +#### AudioStyleTransfer +Transfer style characteristics between audio sources. +- **Modes**: + - Global: Statistics matching + - Frequency bands: Band-wise style transfer + - Temporal segments: Time-based style transfer + - Adaptive: Correlation-based transfer + - Neural style: Deep feature matching + +#### AudioReferenceEncoder +Encode multiple audio references into a combined representation. +- **Combination modes**: Average, weighted, PCA, attention +- **Features**: Segment extraction, weighted blending + +### 5. Generation Guidance + +#### AudioLatentGuidance +Guide the diffusion process with reference latents. +- **Modes**: Additive, multiplicative, attention, gradient +- **Control**: Start/end steps, guidance strength + +## Workflow Examples + +### Basic Style Transfer +``` +1. Load content and style audio +2. Encode both with AudioVAEEncode +3. Apply AudioStyleTransfer +4. Decode result with AudioVAEDecode +``` + +### Hybrid Sound Creation +``` +1. Load multiple audio sources +2. Encode all with AudioVAEEncode +3. Use AudioLatentBlend to combine +4. Apply AudioLatentManipulator for fine-tuning +5. Decode final result +``` + +### Feature-Guided Manipulation +``` +1. Load reference audio +2. Extract features with AudioFeatureExtractor +3. Encode target audio +4. Use AudioLatentManipulator with features +5. Decode result +``` + +### Guided Generation +``` +1. Encode reference audio +2. Apply AudioLatentGuidance to model +3. Generate new audio with guidance +4. Fine-tune with style transfer if needed +``` + +## Technical Details + +### Latent Space Properties +- Stable Audio's VAE compresses audio at ~21.5Hz latent rate +- Latents preserve both spectral and temporal information +- Manipulation in latent space is computationally efficient + +### Best Practices +1. **Normalize audio** before encoding for consistent results +2. **Match sample rates** between audio sources +3. **Use appropriate blend modes** for your use case +4. **Experiment with strength parameters** - start low and increase +5. **Combine multiple techniques** for complex transformations + +### Performance Considerations +- Encoding/decoding is GPU-accelerated +- Latent manipulations are memory-efficient +- Real-time processing possible with optimization +- Batch processing supported for multiple audio files + +## Advanced Techniques + +### Creating Instrument Hybrids +Blend characteristics of different instruments: +1. Encode each instrument separately +2. Use frequency-band style transfer +3. Apply harmonic emphasis +4. Fine-tune with latent interpolation + +### Environmental Audio Design +Create complex soundscapes: +1. Encode multiple environmental sounds +2. Use temporal segment transfer +3. Apply spatial transformations +4. Blend with attention-based combination + +### Rhythmic Pattern Transfer +Transfer rhythm while preserving timbre: +1. Extract tempo/onset features +2. Use temporal stretch with feature guidance +3. Apply style transfer in temporal segments + +## Troubleshooting + +### Common Issues +- **Artifacts in output**: Reduce manipulation strength +- **Loss of clarity**: Check normalization settings +- **Mismatched lengths**: Ensure proper padding/cropping +- **Memory issues**: Process in smaller segments + +### Tips for Quality +- Use high-quality source audio (44.1kHz recommended) +- Avoid extreme parameter values +- Test different interpolation modes +- Monitor latent statistics for anomalies + +## Integration with Stable Audio + +These nodes are designed to work seamlessly with Stable Audio's generation pipeline: +- Compatible with text conditioning +- Can be used before/after generation +- Supports timing embeddings +- Works with all Stable Audio checkpoints + +## Future Enhancements + +Planned features: +- Multi-scale latent manipulation +- Learned projection heads for other modalities +- Advanced spectral processing +- Real-time preview capabilities +- Extended feature extraction options + +## Credits + +This implementation is based on research in: +- Audio VAE architectures +- Style transfer techniques +- Latent space manipulation +- Audio feature analysis + +For more information and updates, visit the project repository. \ No newline at end of file diff --git a/src/network_bending/audio_nodes/__init__.py b/src/network_bending/audio_nodes/__init__.py new file mode 100644 index 0000000..196f4a8 --- /dev/null +++ b/src/network_bending/audio_nodes/__init__.py @@ -0,0 +1,44 @@ +""" +Audio conditioning nodes for Stable Audio in ComfyUI +""" + +from .audio_latent_nodes import ( + AudioVAEEncode, + AudioVAEDecode, + AudioLatentInterpolate, + AudioLatentBlend, + AudioFeatureExtractor, + AudioLatentManipulator, +) + +from .audio_style_transfer import ( + AudioStyleTransfer, + AudioLatentGuidance, + AudioReferenceEncoder +) + +NODE_CLASS_MAPPINGS = { + "AudioVAEEncode": AudioVAEEncode, + "AudioVAEDecode": AudioVAEDecode, + "AudioLatentInterpolate": AudioLatentInterpolate, + "AudioLatentBlend": AudioLatentBlend, + "AudioFeatureExtractor": AudioFeatureExtractor, + "AudioLatentManipulator": AudioLatentManipulator, + "AudioStyleTransfer": AudioStyleTransfer, + "AudioLatentGuidance": AudioLatentGuidance, + "AudioReferenceEncoder": AudioReferenceEncoder, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "AudioVAEEncode": "Audio VAE Encode", + "AudioVAEDecode": "Audio VAE Decode", + "AudioLatentInterpolate": "Audio Latent Interpolate", + "AudioLatentBlend": "Audio Latent Blend", + "AudioFeatureExtractor": "Audio Feature Extractor", + "AudioLatentManipulator": "Audio Latent Manipulator", + "AudioStyleTransfer": "Audio Style Transfer", + "AudioLatentGuidance": "Audio Latent Guidance", + "AudioReferenceEncoder": "Audio Reference Encoder", +} + +__all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS'] \ No newline at end of file diff --git a/src/network_bending/audio_nodes/audio_latent_nodes.py b/src/network_bending/audio_nodes/audio_latent_nodes.py new file mode 100644 index 0000000..e4343d8 --- /dev/null +++ b/src/network_bending/audio_nodes/audio_latent_nodes.py @@ -0,0 +1,638 @@ +""" +Audio Latent Space Manipulation Nodes for Stable Audio +These nodes enable working directly with audio latents for style transfer and hybrid generation +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from typing import Dict, List, Tuple, Any, Optional +import torchaudio +import librosa + +# Import the additional nodes from the style transfer module +from .audio_style_transfer import ( + AudioStyleTransfer, + AudioLatentGuidance, + AudioReferenceEncoder +) + +class AudioVAEEncode: + """ + Encode audio waveforms into Stable Audio latent space + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO", {"tooltip": "Audio waveform to encode"}), + "vae": ("VAE", {"tooltip": "Stable Audio VAE model"}), + }, + "optional": { + "normalize": ("BOOLEAN", { + "default": True, + "tooltip": "Normalize audio before encoding" + }), + "target_sample_rate": ("INT", { + "default": 44100, + "min": 16000, + "max": 48000, + "tooltip": "Target sample rate for encoding" + }), + } + } + + RETURN_TYPES = ("AUDIO_LATENT", "LATENT_INFO") + RETURN_NAMES = ("latent", "info") + FUNCTION = "encode_audio" + CATEGORY = "audio/latent" + + def encode_audio(self, audio, vae, normalize=True, target_sample_rate=44100): + # Extract audio tensor and sample rate + waveform, sample_rate = audio + + # Resample if necessary + if sample_rate != target_sample_rate: + resampler = torchaudio.transforms.Resample(sample_rate, target_sample_rate) + waveform = resampler(waveform) + + # Normalize audio + if normalize: + waveform = waveform / torch.max(torch.abs(waveform)) + + # Ensure correct shape for VAE (batch, channels, samples) + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) + + # Encode with VAE + with torch.no_grad(): + # Move to same device as VAE + device = next(vae.parameters()).device + waveform = waveform.to(device) + + # Encode + latent = vae.encode(waveform) + + # Get latent statistics + latent_mean = latent.mean() + latent_std = latent.std() + + # Create info dictionary + info = { + "original_shape": waveform.shape, + "latent_shape": latent.shape, + "sample_rate": target_sample_rate, + "latent_mean": latent_mean.item(), + "latent_std": latent_std.item(), + } + + return (latent, info) + +class AudioVAEDecode: + """ + Decode audio latents back to waveforms + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent": ("AUDIO_LATENT", {"tooltip": "Audio latent to decode"}), + "vae": ("VAE", {"tooltip": "Stable Audio VAE model"}), + }, + "optional": { + "denormalize": ("BOOLEAN", { + "default": True, + "tooltip": "Denormalize output audio" + }), + } + } + + RETURN_TYPES = ("AUDIO",) + RETURN_NAMES = ("audio",) + FUNCTION = "decode_audio" + CATEGORY = "audio/latent" + + def decode_audio(self, latent, vae, denormalize=True): + # Decode with VAE + with torch.no_grad(): + # Ensure latent is on correct device + device = next(vae.parameters()).device + latent = latent.to(device) + + # Decode + waveform = vae.decode(latent) + + # Move to CPU for further processing + waveform = waveform.cpu() + + # Denormalize if needed + if denormalize: + # Ensure audio is in valid range [-1, 1] + waveform = torch.clamp(waveform, -1.0, 1.0) + + # Extract sample rate (default to 44100 if not stored) + sample_rate = 44100 + + # Return as (waveform, sample_rate) tuple + return ((waveform, sample_rate),) + +class AudioLatentInterpolate: + """ + Interpolate between two audio latents for smooth transitions and hybrid sounds + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent_a": ("AUDIO_LATENT", {"tooltip": "First audio latent"}), + "latent_b": ("AUDIO_LATENT", {"tooltip": "Second audio latent"}), + "interpolation_factor": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "Interpolation factor (0=A, 1=B)" + }), + "interpolation_mode": ([ + "linear", + "spherical", + "cubic", + "sine" + ], { + "default": "spherical", + "tooltip": "Interpolation method" + }), + }, + "optional": { + "curve_power": ("FLOAT", { + "default": 1.0, + "min": 0.1, + "max": 10.0, + "step": 0.1, + "tooltip": "Power curve for non-linear interpolation" + }), + } + } + + RETURN_TYPES = ("AUDIO_LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "interpolate_latents" + CATEGORY = "audio/latent" + + def interpolate_latents(self, latent_a, latent_b, interpolation_factor, + interpolation_mode, curve_power=1.0): + # Ensure latents are on the same device + device = latent_a.device + latent_b = latent_b.to(device) + + # Apply curve to interpolation factor if needed + if curve_power != 1.0: + interpolation_factor = interpolation_factor ** curve_power + + if interpolation_mode == "linear": + # Simple linear interpolation + result = (1 - interpolation_factor) * latent_a + interpolation_factor * latent_b + + elif interpolation_mode == "spherical": + # Spherical linear interpolation (SLERP) + # Normalize latents + latent_a_norm = F.normalize(latent_a.flatten(1), dim=1).reshape(latent_a.shape) + latent_b_norm = F.normalize(latent_b.flatten(1), dim=1).reshape(latent_b.shape) + + # Compute angle between latents + dot_product = (latent_a_norm * latent_b_norm).sum() + angle = torch.acos(torch.clamp(dot_product, -1, 1)) + + # SLERP formula + if angle < 1e-6: + result = (1 - interpolation_factor) * latent_a + interpolation_factor * latent_b + else: + sin_angle = torch.sin(angle) + result = (torch.sin((1 - interpolation_factor) * angle) / sin_angle) * latent_a + \ + (torch.sin(interpolation_factor * angle) / sin_angle) * latent_b + + elif interpolation_mode == "cubic": + # Cubic interpolation (smooth start and end) + t = interpolation_factor + t2 = t * t + t3 = t2 * t + factor = 3 * t2 - 2 * t3 + result = (1 - factor) * latent_a + factor * latent_b + + elif interpolation_mode == "sine": + # Sine-based interpolation + factor = 0.5 * (1 + torch.sin(torch.tensor(np.pi * (interpolation_factor - 0.5)))) + result = (1 - factor) * latent_a + factor * latent_b + + return (result,) + +class AudioLatentBlend: + """ + Blend multiple audio latents with customizable weights + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent_a": ("AUDIO_LATENT", {"tooltip": "First audio latent"}), + "latent_b": ("AUDIO_LATENT", {"tooltip": "Second audio latent"}), + "blend_mode": ([ + "add", + "multiply", + "screen", + "overlay", + "soft_light", + "difference", + "exclusion" + ], { + "default": "add", + "tooltip": "Blending mode" + }), + "blend_factor": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 2.0, + "step": 0.01, + "tooltip": "Blend strength" + }), + }, + "optional": { + "latent_c": ("AUDIO_LATENT", {"tooltip": "Optional third audio latent"}), + "latent_d": ("AUDIO_LATENT", {"tooltip": "Optional fourth audio latent"}), + "weight_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_c": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_d": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "normalize": ("BOOLEAN", {"default": True}), + } + } + + RETURN_TYPES = ("AUDIO_LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "blend_latents" + CATEGORY = "audio/latent" + + def blend_latents(self, latent_a, latent_b, blend_mode, blend_factor, + latent_c=None, latent_d=None, + weight_a=1.0, weight_b=1.0, weight_c=1.0, weight_d=1.0, + normalize=True): + + device = latent_a.device + + # Collect all available latents and weights + latents = [latent_a, latent_b] + weights = [weight_a, weight_b] + + if latent_c is not None: + latents.append(latent_c.to(device)) + weights.append(weight_c) + if latent_d is not None: + latents.append(latent_d.to(device)) + weights.append(weight_d) + + # Normalize weights if requested + if normalize: + total_weight = sum(weights) + weights = [w / total_weight for w in weights] + + # Apply blend mode + if blend_mode == "add": + # Weighted addition + result = sum(latent * weight for latent, weight in zip(latents, weights)) + result = result * blend_factor + latents[0] * (1 - blend_factor) + + elif blend_mode == "multiply": + # Multiplicative blending + result = latents[0] + for i in range(1, len(latents)): + blend = latents[i] * weights[i] + result = result * (1 + blend_factor * (blend - 1)) + + elif blend_mode == "screen": + # Screen blending (inverse multiply) + result = latents[0] + for i in range(1, len(latents)): + blend = 1 - (1 - result) * (1 - latents[i] * weights[i]) + result = result * (1 - blend_factor) + blend * blend_factor + + elif blend_mode == "overlay": + # Overlay blending + result = latents[0] + for i in range(1, len(latents)): + mask = (result < 0.5).float() + blend = mask * (2 * result * latents[i] * weights[i]) + \ + (1 - mask) * (1 - 2 * (1 - result) * (1 - latents[i] * weights[i])) + result = result * (1 - blend_factor) + blend * blend_factor + + elif blend_mode == "soft_light": + # Soft light blending + result = latents[0] + for i in range(1, len(latents)): + blend = latents[i] * weights[i] + soft = (1 - 2 * blend) * result * result + 2 * blend * result + result = result * (1 - blend_factor) + soft * blend_factor + + elif blend_mode == "difference": + # Difference blending + result = latents[0] + for i in range(1, len(latents)): + diff = torch.abs(result - latents[i] * weights[i]) + result = result * (1 - blend_factor) + diff * blend_factor + + elif blend_mode == "exclusion": + # Exclusion blending + result = latents[0] + for i in range(1, len(latents)): + excl = result + latents[i] * weights[i] - 2 * result * latents[i] * weights[i] + result = result * (1 - blend_factor) + excl * blend_factor + + return (result,) + +class AudioFeatureExtractor: + """ + Extract various audio features for latent manipulation + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "audio": ("AUDIO", {"tooltip": "Audio to analyze"}), + "feature_type": ([ + "spectral_centroid", + "spectral_rolloff", + "zero_crossing_rate", + "mfcc", + "spectral_bandwidth", + "rms_energy", + "tempo", + "onset_strength" + ], { + "default": "spectral_centroid", + "tooltip": "Type of audio feature to extract" + }), + }, + "optional": { + "n_mfcc": ("INT", { + "default": 13, + "min": 1, + "max": 40, + "tooltip": "Number of MFCC coefficients" + }), + "hop_length": ("INT", { + "default": 512, + "min": 64, + "max": 2048, + "tooltip": "Hop length for feature extraction" + }), + } + } + + RETURN_TYPES = ("AUDIO_FEATURES",) + RETURN_NAMES = ("features",) + FUNCTION = "extract_features" + CATEGORY = "audio/analysis" + + def extract_features(self, audio, feature_type, n_mfcc=13, hop_length=512): + waveform, sample_rate = audio + + # Convert to numpy for librosa + if isinstance(waveform, torch.Tensor): + waveform_np = waveform.squeeze().cpu().numpy() + else: + waveform_np = waveform + + # Extract features based on type + if feature_type == "spectral_centroid": + features = librosa.feature.spectral_centroid( + y=waveform_np, sr=sample_rate, hop_length=hop_length + ) + elif feature_type == "spectral_rolloff": + features = librosa.feature.spectral_rolloff( + y=waveform_np, sr=sample_rate, hop_length=hop_length + ) + elif feature_type == "zero_crossing_rate": + features = librosa.feature.zero_crossing_rate( + waveform_np, hop_length=hop_length + ) + elif feature_type == "mfcc": + features = librosa.feature.mfcc( + y=waveform_np, sr=sample_rate, n_mfcc=n_mfcc, hop_length=hop_length + ) + elif feature_type == "spectral_bandwidth": + features = librosa.feature.spectral_bandwidth( + y=waveform_np, sr=sample_rate, hop_length=hop_length + ) + elif feature_type == "rms_energy": + features = librosa.feature.rms( + y=waveform_np, hop_length=hop_length + ) + elif feature_type == "tempo": + tempo, beats = librosa.beat.beat_track( + y=waveform_np, sr=sample_rate, hop_length=hop_length + ) + features = np.array([[tempo]]) + elif feature_type == "onset_strength": + features = librosa.onset.onset_strength( + y=waveform_np, sr=sample_rate, hop_length=hop_length + ).reshape(1, -1) + + # Convert to tensor + features_tensor = torch.from_numpy(features).float() + + # Create feature dictionary + feature_dict = { + "type": feature_type, + "data": features_tensor, + "sample_rate": sample_rate, + "hop_length": hop_length, + "shape": features_tensor.shape + } + + return (feature_dict,) + +class AudioLatentManipulator: + """ + Manipulate audio latents based on extracted features or manual controls + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "latent": ("AUDIO_LATENT", {"tooltip": "Audio latent to manipulate"}), + "manipulation_type": ([ + "frequency_shift", + "temporal_stretch", + "harmonic_emphasis", + "noise_injection", + "dynamic_range", + "spatial_transform", + "resonance_filter" + ], { + "default": "frequency_shift", + "tooltip": "Type of manipulation to apply" + }), + "strength": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 2.0, + "step": 0.01, + "tooltip": "Manipulation strength" + }), + }, + "optional": { + "features": ("AUDIO_FEATURES", {"tooltip": "Optional features to guide manipulation"}), + "frequency": ("FLOAT", { + "default": 440.0, + "min": 20.0, + "max": 20000.0, + "tooltip": "Target frequency for frequency-based operations" + }), + "time_factor": ("FLOAT", { + "default": 1.0, + "min": 0.25, + "max": 4.0, + "tooltip": "Time stretch factor" + }), + } + } + + RETURN_TYPES = ("AUDIO_LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "manipulate_latent" + CATEGORY = "audio/latent" + + def manipulate_latent(self, latent, manipulation_type, strength, + features=None, frequency=440.0, time_factor=1.0): + device = latent.device + result = latent.clone() + + # Get latent dimensions + batch, channels, *spatial_dims = result.shape + + if manipulation_type == "frequency_shift": + # Shift frequency content in latent space + # Apply FFT-like transformation in latent space + if len(spatial_dims) == 2: + # 2D latent (height, width) + fft_latent = torch.fft.fft2(result) + # Shift frequencies + shift_amount = int(strength * spatial_dims[1] * frequency / 22050) + fft_latent = torch.roll(fft_latent, shifts=shift_amount, dims=-1) + result = torch.fft.ifft2(fft_latent).real + else: + # 1D latent + fft_latent = torch.fft.fft(result) + shift_amount = int(strength * spatial_dims[0] * frequency / 22050) + fft_latent = torch.roll(fft_latent, shifts=shift_amount, dims=-1) + result = torch.fft.ifft(fft_latent).real + + elif manipulation_type == "temporal_stretch": + # Stretch or compress temporal dimension + if len(spatial_dims) >= 1: + time_dim = spatial_dims[-1] + new_time_dim = int(time_dim * time_factor) + + # Interpolate along time dimension + if len(spatial_dims) == 2: + result = F.interpolate(result, size=(spatial_dims[0], new_time_dim), + mode='bilinear', align_corners=False) + else: + result = F.interpolate(result, size=new_time_dim, + mode='linear', align_corners=False) + + # Adjust back to original size with padding or cropping + if new_time_dim > time_dim: + result = result[..., :time_dim] + elif new_time_dim < time_dim: + pad_size = time_dim - new_time_dim + result = F.pad(result, (0, pad_size)) + + elif manipulation_type == "harmonic_emphasis": + # Emphasize harmonic content + # Create harmonic mask in frequency domain + if len(spatial_dims) >= 1: + fft_latent = torch.fft.fft(result, dim=-1) + magnitudes = torch.abs(fft_latent) + + # Find peaks (harmonics) + kernel_size = 5 + maxpool = nn.MaxPool1d(kernel_size, stride=1, padding=kernel_size//2) + if len(spatial_dims) == 2: + mag_flat = magnitudes.reshape(batch * channels * spatial_dims[0], -1) + peaks = (magnitudes == maxpool(mag_flat).reshape(magnitudes.shape)) + else: + peaks = (magnitudes == maxpool(magnitudes)) + + # Emphasize peaks + fft_latent = fft_latent * (1 + strength * peaks.float()) + result = torch.fft.ifft(fft_latent, dim=-1).real + + elif manipulation_type == "noise_injection": + # Inject controlled noise + noise = torch.randn_like(result) * strength * 0.1 + if features is not None and "data" in features: + # Modulate noise by features + feature_data = features["data"].to(device) + if feature_data.numel() > 0: + feature_scalar = feature_data.mean().item() + noise = noise * (1 + feature_scalar) + result = result + noise + + elif manipulation_type == "dynamic_range": + # Compress or expand dynamic range + # Apply soft clipping/expansion + threshold = 0.5 + if strength < 1.0: + # Compression + mask = torch.abs(result) > threshold + result = torch.where(mask, + torch.sign(result) * (threshold + (torch.abs(result) - threshold) * strength), + result) + else: + # Expansion + mask = torch.abs(result) < threshold + result = torch.where(mask, + result * strength, + result) + + elif manipulation_type == "spatial_transform": + # Transform spatial relationships in latent + if len(spatial_dims) == 2: + # Rotate latent space + angle = strength * np.pi + cos_a = torch.cos(torch.tensor(angle)) + sin_a = torch.sin(torch.tensor(angle)) + + # Create rotation matrix + rot_matrix = torch.tensor([[cos_a, -sin_a], [sin_a, cos_a]], device=device) + + # Apply rotation (simplified - works on channel pairs) + if channels >= 2: + for i in range(0, channels-1, 2): + pair = result[:, i:i+2].reshape(batch, 2, -1) + rotated = torch.matmul(rot_matrix, pair) + result[:, i:i+2] = rotated.reshape(batch, 2, *spatial_dims) + + elif manipulation_type == "resonance_filter": + # Apply resonant filtering in latent space + if len(spatial_dims) >= 1: + # Create resonance kernel + center_freq = frequency / 22050 # Normalize frequency + bandwidth = 0.1 * strength + + freqs = torch.linspace(0, 1, spatial_dims[-1], device=device) + resonance = torch.exp(-((freqs - center_freq) ** 2) / (2 * bandwidth ** 2)) + + # Apply in frequency domain + fft_latent = torch.fft.fft(result, dim=-1) + fft_latent = fft_latent * resonance + result = torch.fft.ifft(fft_latent, dim=-1).real + + return (result,) \ No newline at end of file diff --git a/src/network_bending/audio_nodes/audio_style_transfer.py b/src/network_bending/audio_nodes/audio_style_transfer.py new file mode 100644 index 0000000..cc081a7 --- /dev/null +++ b/src/network_bending/audio_nodes/audio_style_transfer.py @@ -0,0 +1,485 @@ +""" +Audio Style Transfer and Advanced Latent Manipulation Nodes +""" + +import torch +import torch.nn as nn +import torch.nn.functional as F +import numpy as np +from typing import Dict, List, Tuple, Any, Optional + +class AudioStyleTransfer: + """ + Transfer style characteristics from one audio to another in latent space + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "content_latent": ("AUDIO_LATENT", {"tooltip": "Content audio latent"}), + "style_latent": ("AUDIO_LATENT", {"tooltip": "Style audio latent"}), + "transfer_mode": ([ + "global", + "frequency_bands", + "temporal_segments", + "adaptive", + "neural_style" + ], { + "default": "adaptive", + "tooltip": "Style transfer mode" + }), + "style_strength": ("FLOAT", { + "default": 0.7, + "min": 0.0, + "max": 1.0, + "step": 0.01, + "tooltip": "How much style to transfer" + }), + }, + "optional": { + "preserve_content": ("FLOAT", { + "default": 0.3, + "min": 0.0, + "max": 1.0, + "tooltip": "How much original content to preserve" + }), + "frequency_bands": ("INT", { + "default": 4, + "min": 2, + "max": 16, + "tooltip": "Number of frequency bands for band-wise transfer" + }), + } + } + + RETURN_TYPES = ("AUDIO_LATENT",) + RETURN_NAMES = ("latent",) + FUNCTION = "transfer_style" + CATEGORY = "audio/style" + + def transfer_style(self, content_latent, style_latent, transfer_mode, + style_strength, preserve_content=0.3, frequency_bands=4): + device = content_latent.device + style_latent = style_latent.to(device) + + if transfer_mode == "global": + # Global style transfer using statistics matching + # Compute statistics + content_mean = content_latent.mean(dim=(2, 3), keepdim=True) if content_latent.dim() == 4 else content_latent.mean(dim=2, keepdim=True) + content_std = content_latent.std(dim=(2, 3), keepdim=True) if content_latent.dim() == 4 else content_latent.std(dim=2, keepdim=True) + + style_mean = style_latent.mean(dim=(2, 3), keepdim=True) if style_latent.dim() == 4 else style_latent.mean(dim=2, keepdim=True) + style_std = style_latent.std(dim=(2, 3), keepdim=True) if style_latent.dim() == 4 else style_latent.std(dim=2, keepdim=True) + + # Normalize content and apply style statistics + normalized = (content_latent - content_mean) / (content_std + 1e-6) + stylized = normalized * style_std + style_mean + + # Blend with original + result = stylized * style_strength + content_latent * preserve_content + + elif transfer_mode == "frequency_bands": + # Transfer style in different frequency bands + result = content_latent.clone() + + # Apply FFT + fft_content = torch.fft.fft(content_latent, dim=-1) + fft_style = torch.fft.fft(style_latent, dim=-1) + + # Split into frequency bands + freq_size = fft_content.shape[-1] + band_size = freq_size // frequency_bands + + for i in range(frequency_bands): + start_idx = i * band_size + end_idx = (i + 1) * band_size if i < frequency_bands - 1 else freq_size + + # Transfer magnitude in this band + content_mag = torch.abs(fft_content[..., start_idx:end_idx]) + content_phase = torch.angle(fft_content[..., start_idx:end_idx]) + style_mag = torch.abs(fft_style[..., start_idx:end_idx]) + + # Blend magnitudes + new_mag = style_mag * style_strength + content_mag * preserve_content + + # Reconstruct with new magnitude + fft_content[..., start_idx:end_idx] = new_mag * torch.exp(1j * content_phase) + + # Inverse FFT + result = torch.fft.ifft(fft_content, dim=-1).real + + elif transfer_mode == "temporal_segments": + # Transfer style in temporal segments + time_dim = content_latent.shape[-1] + segment_size = time_dim // 4 # 4 segments + + result = content_latent.clone() + + for i in range(4): + start_idx = i * segment_size + end_idx = (i + 1) * segment_size if i < 3 else time_dim + + # Get segments + content_seg = content_latent[..., start_idx:end_idx] + style_seg = style_latent[..., start_idx:end_idx] + + # Transfer style in segment + seg_mean_c = content_seg.mean() + seg_std_c = content_seg.std() + seg_mean_s = style_seg.mean() + seg_std_s = style_seg.std() + + normalized_seg = (content_seg - seg_mean_c) / (seg_std_c + 1e-6) + stylized_seg = normalized_seg * seg_std_s + seg_mean_s + + result[..., start_idx:end_idx] = stylized_seg * style_strength + content_seg * preserve_content + + elif transfer_mode == "adaptive": + # Adaptive style transfer using correlation + # Flatten spatial dimensions + content_flat = content_latent.reshape(content_latent.shape[0], content_latent.shape[1], -1) + style_flat = style_latent.reshape(style_latent.shape[0], style_latent.shape[1], -1) + + # Compute correlation matrix + content_cov = torch.matmul(content_flat, content_flat.transpose(1, 2)) + style_cov = torch.matmul(style_flat, style_flat.transpose(1, 2)) + + # Eigen decomposition for whitening and coloring + content_eig_val, content_eig_vec = torch.linalg.eigh(content_cov) + style_eig_val, style_eig_vec = torch.linalg.eigh(style_cov) + + # Whitening transform + content_eig_val = torch.clamp(content_eig_val, min=1e-6) + whitening = torch.matmul(content_eig_vec, torch.diag_embed(1.0 / torch.sqrt(content_eig_val))) + + # Coloring transform + style_eig_val = torch.clamp(style_eig_val, min=1e-6) + coloring = torch.matmul(style_eig_vec, torch.diag_embed(torch.sqrt(style_eig_val))) + + # Apply transforms + whitened = torch.matmul(whitening.transpose(1, 2), content_flat) + stylized = torch.matmul(coloring, whitened) + + # Reshape back + stylized = stylized.reshape(content_latent.shape) + result = stylized * style_strength + content_latent * preserve_content + + elif transfer_mode == "neural_style": + # Neural style transfer inspired approach + # Use convolutional layers to extract features + conv1 = nn.Conv1d(content_latent.shape[1], 64, 3, padding=1).to(device) + conv2 = nn.Conv1d(64, 32, 3, padding=1).to(device) + + # Extract features + content_feat1 = F.relu(conv1(content_latent.squeeze(2) if content_latent.dim() == 4 else content_latent)) + content_feat2 = F.relu(conv2(content_feat1)) + + style_feat1 = F.relu(conv1(style_latent.squeeze(2) if style_latent.dim() == 4 else style_latent)) + style_feat2 = F.relu(conv2(style_feat1)) + + # Compute Gram matrices + def gram_matrix(features): + b, c, w = features.size() + features = features.view(b, c, w) + gram = torch.bmm(features, features.transpose(1, 2)) + return gram / (c * w) + + content_gram = gram_matrix(content_feat2) + style_gram = gram_matrix(style_feat2) + + # Match Gram matrices + target_gram = style_gram * style_strength + content_gram * preserve_content + + # Approximate reconstruction (simplified) + result = content_latent + 0.1 * (target_gram.mean() - content_gram.mean()) + + return (result,) + +class AudioLatentGuidance: + """ + Guide audio generation with reference latents during the diffusion process + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "model": ("MODEL", {"tooltip": "Stable Audio diffusion model"}), + "reference_latent": ("AUDIO_LATENT", {"tooltip": "Reference audio latent for guidance"}), + "guidance_strength": ("FLOAT", { + "default": 0.5, + "min": 0.0, + "max": 2.0, + "step": 0.01, + "tooltip": "Strength of latent guidance" + }), + "guidance_mode": ([ + "additive", + "multiplicative", + "attention", + "gradient" + ], { + "default": "attention", + "tooltip": "How to apply guidance" + }), + }, + "optional": { + "start_at_step": ("INT", { + "default": 0, + "min": 0, + "max": 1000, + "tooltip": "Step to start applying guidance" + }), + "end_at_step": ("INT", { + "default": 1000, + "min": 0, + "max": 1000, + "tooltip": "Step to stop applying guidance" + }), + "features": ("AUDIO_FEATURES", {"tooltip": "Optional features to modulate guidance"}), + } + } + + RETURN_TYPES = ("MODEL",) + RETURN_NAMES = ("model",) + FUNCTION = "apply_guidance" + CATEGORY = "audio/conditioning" + + def apply_guidance(self, model, reference_latent, guidance_strength, guidance_mode, + start_at_step=0, end_at_step=1000, features=None): + # Clone the model to avoid modifying the original + guided_model = model.clone() + + # Store guidance parameters in the model + guided_model.audio_latent_guidance = { + "reference_latent": reference_latent, + "strength": guidance_strength, + "mode": guidance_mode, + "start_step": start_at_step, + "end_step": end_at_step, + "features": features + } + + # Patch the model's forward method to include guidance + original_forward = guided_model.model.forward + + def guided_forward(x, timestep, context=None, **kwargs): + # Get current step from timestep + current_step = timestep[0].item() if isinstance(timestep, torch.Tensor) else timestep + + # Check if we should apply guidance + guidance = guided_model.audio_latent_guidance + if guidance["start_step"] <= current_step <= guidance["end_step"]: + ref_latent = guidance["reference_latent"].to(x.device) + strength = guidance["strength"] + + if guidance["mode"] == "additive": + # Add reference latent influence + x = x + strength * 0.1 * ref_latent + + elif guidance["mode"] == "multiplicative": + # Multiplicative guidance + ref_norm = F.normalize(ref_latent.flatten(1), dim=1).reshape(ref_latent.shape) + x = x * (1 + strength * ref_norm) + + elif guidance["mode"] == "attention": + # Attention-based guidance (modify context) + if context is not None: + # Add reference latent information to context + ref_context = ref_latent.mean(dim=(2, 3), keepdim=True) if ref_latent.dim() == 4 else ref_latent.mean(dim=2, keepdim=True) + ref_context = ref_context.expand_as(context) + context = context + strength * 0.1 * ref_context + + elif guidance["mode"] == "gradient": + # Gradient-based guidance + # Compute similarity between x and reference + similarity = F.cosine_similarity(x.flatten(1), ref_latent.flatten(1), dim=1) + similarity = similarity.mean() + + # Modify x to increase similarity + grad_direction = ref_latent - x + x = x + strength * 0.05 * grad_direction * (1 - similarity) + + # Call original forward + return original_forward(x, timestep, context, **kwargs) + + # Replace forward method + guided_model.model.forward = guided_forward + + return (guided_model,) + +class AudioReferenceEncoder: + """ + Encode multiple audio references into a combined latent representation + """ + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "vae": ("VAE", {"tooltip": "Stable Audio VAE model"}), + "audio_a": ("AUDIO", {"tooltip": "First reference audio"}), + "combination_mode": ([ + "average", + "weighted", + "pca", + "attention" + ], { + "default": "attention", + "tooltip": "How to combine multiple references" + }), + }, + "optional": { + "audio_b": ("AUDIO", {"tooltip": "Second reference audio"}), + "audio_c": ("AUDIO", {"tooltip": "Third reference audio"}), + "audio_d": ("AUDIO", {"tooltip": "Fourth reference audio"}), + "weight_a": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_b": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_c": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "weight_d": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 2.0}), + "extract_segments": ("BOOLEAN", { + "default": False, + "tooltip": "Extract most characteristic segments" + }), + "segment_duration": ("FLOAT", { + "default": 2.0, + "min": 0.5, + "max": 10.0, + "tooltip": "Duration of segments to extract (seconds)" + }), + } + } + + RETURN_TYPES = ("AUDIO_LATENT", "LATENT_INFO") + RETURN_NAMES = ("combined_latent", "info") + FUNCTION = "encode_references" + CATEGORY = "audio/reference" + + def encode_references(self, vae, audio_a, combination_mode, + audio_b=None, audio_c=None, audio_d=None, + weight_a=1.0, weight_b=1.0, weight_c=1.0, weight_d=1.0, + extract_segments=False, segment_duration=2.0): + + device = next(vae.parameters()).device + + # Collect all available audio inputs + audios = [audio_a] + weights = [weight_a] + + if audio_b is not None: + audios.append(audio_b) + weights.append(weight_b) + if audio_c is not None: + audios.append(audio_c) + weights.append(weight_c) + if audio_d is not None: + audios.append(audio_d) + weights.append(weight_d) + + # Normalize weights + total_weight = sum(weights) + weights = [w / total_weight for w in weights] + + # Encode all audio inputs + encoded_latents = [] + + for audio, weight in zip(audios, weights): + waveform, sample_rate = audio + + # Extract characteristic segment if requested + if extract_segments: + segment_samples = int(segment_duration * sample_rate) + + # Find segment with highest energy + if waveform.shape[-1] > segment_samples: + # Compute energy in sliding windows + energy = [] + for i in range(0, waveform.shape[-1] - segment_samples, segment_samples // 4): + segment = waveform[..., i:i+segment_samples] + energy.append(torch.sum(segment ** 2).item()) + + # Get the highest energy segment + best_idx = np.argmax(energy) * (segment_samples // 4) + waveform = waveform[..., best_idx:best_idx+segment_samples] + + # Normalize audio + waveform = waveform / torch.max(torch.abs(waveform)) + + # Ensure correct shape + if waveform.dim() == 2: + waveform = waveform.unsqueeze(0) + + # Encode with VAE + with torch.no_grad(): + waveform = waveform.to(device) + latent = vae.encode(waveform) + encoded_latents.append((latent, weight)) + + # Combine latents based on mode + if combination_mode == "average": + # Simple weighted average + combined = sum(latent * weight for latent, weight in encoded_latents) + + elif combination_mode == "weighted": + # Weighted combination with normalization + combined = torch.zeros_like(encoded_latents[0][0]) + for latent, weight in encoded_latents: + combined += latent * weight + combined = F.normalize(combined, dim=1) + + elif combination_mode == "pca": + # PCA-based combination + # Stack all latents + all_latents = torch.stack([lat for lat, _ in encoded_latents]) + + # Flatten spatial dimensions + b, c, *spatial = all_latents.shape + all_latents_flat = all_latents.reshape(len(encoded_latents), -1) + + # Compute PCA + mean = all_latents_flat.mean(0, keepdim=True) + centered = all_latents_flat - mean + + # Compute covariance + cov = torch.matmul(centered.T, centered) / (len(encoded_latents) - 1) + + # Get first principal component + eigvals, eigvecs = torch.linalg.eigh(cov) + pc1 = eigvecs[:, -1] # Largest eigenvalue + + # Project and combine + projections = torch.matmul(centered, pc1) + combined_flat = mean + projections.mean() * pc1 + combined = combined_flat.reshape(1, c, *spatial) + + elif combination_mode == "attention": + # Attention-based combination + # Use latents as queries and keys + all_latents = torch.stack([lat for lat, _ in encoded_latents]) + + # Compute attention scores + queries = all_latents.mean(dim=(2, 3)) if all_latents.dim() == 4 else all_latents.mean(dim=2) + keys = queries + + # Scaled dot-product attention + scores = torch.matmul(queries, keys.transpose(0, 1)) / np.sqrt(queries.shape[-1]) + attention_weights = F.softmax(scores, dim=-1) + + # Apply attention weights + combined = torch.zeros_like(encoded_latents[0][0]) + for i, (latent, base_weight) in enumerate(encoded_latents): + # Combine base weight with attention weight + final_weight = base_weight * attention_weights[i].mean() + combined += latent * final_weight + + # Create info dictionary + info = { + "num_references": len(audios), + "combination_mode": combination_mode, + "weights": weights, + "latent_shape": combined.shape, + "extracted_segments": extract_segments + } + + return (combined, info) \ No newline at end of file diff --git a/src/network_bending/audio_workflows/audio_style_transfer_workflow.json b/src/network_bending/audio_workflows/audio_style_transfer_workflow.json new file mode 100644 index 0000000..3d88428 --- /dev/null +++ b/src/network_bending/audio_workflows/audio_style_transfer_workflow.json @@ -0,0 +1,219 @@ +{ + "last_node_id": 10, + "last_link_id": 12, + "nodes": [ + { + "id": 1, + "type": "LoadAudio", + "pos": [100, 100], + "size": [300, 100], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [1, 2] + } + ], + "properties": { + "Node name for S&R": "LoadAudio" + }, + "widgets_values": ["content_audio.wav"] + }, + { + "id": 2, + "type": "LoadAudio", + "pos": [100, 250], + "size": [300, 100], + "outputs": [ + { + "name": "AUDIO", + "type": "AUDIO", + "links": [3] + } + ], + "properties": { + "Node name for S&R": "LoadAudio" + }, + "widgets_values": ["style_audio.wav"] + }, + { + "id": 3, + "type": "LoadVAE", + "pos": [100, 400], + "size": [300, 100], + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [4, 5, 6] + } + ], + "properties": { + "Node name for S&R": "LoadVAE" + }, + "widgets_values": ["stable_audio_vae.safetensors"] + }, + { + "id": 4, + "type": "AudioVAEEncode", + "pos": [450, 100], + "size": [300, 150], + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 1 + }, + { + "name": "vae", + "type": "VAE", + "link": 4 + } + ], + "outputs": [ + { + "name": "latent", + "type": "AUDIO_LATENT", + "links": [7] + }, + { + "name": "info", + "type": "LATENT_INFO", + "links": null + } + ], + "properties": { + "Node name for S&R": "AudioVAEEncode" + }, + "widgets_values": [true, 44100] + }, + { + "id": 5, + "type": "AudioVAEEncode", + "pos": [450, 300], + "size": [300, 150], + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 3 + }, + { + "name": "vae", + "type": "VAE", + "link": 5 + } + ], + "outputs": [ + { + "name": "latent", + "type": "AUDIO_LATENT", + "links": [8] + }, + { + "name": "info", + "type": "LATENT_INFO", + "links": null + } + ], + "properties": { + "Node name for S&R": "AudioVAEEncode" + }, + "widgets_values": [true, 44100] + }, + { + "id": 6, + "type": "AudioStyleTransfer", + "pos": [800, 200], + "size": [350, 200], + "inputs": [ + { + "name": "content_latent", + "type": "AUDIO_LATENT", + "link": 7 + }, + { + "name": "style_latent", + "type": "AUDIO_LATENT", + "link": 8 + } + ], + "outputs": [ + { + "name": "latent", + "type": "AUDIO_LATENT", + "links": [9] + } + ], + "properties": { + "Node name for S&R": "AudioStyleTransfer" + }, + "widgets_values": ["adaptive", 0.7, 0.3, 4] + }, + { + "id": 7, + "type": "AudioVAEDecode", + "pos": [1200, 200], + "size": [300, 150], + "inputs": [ + { + "name": "latent", + "type": "AUDIO_LATENT", + "link": 9 + }, + { + "name": "vae", + "type": "VAE", + "link": 6 + } + ], + "outputs": [ + { + "name": "audio", + "type": "AUDIO", + "links": [10] + } + ], + "properties": { + "Node name for S&R": "AudioVAEDecode" + }, + "widgets_values": [true] + }, + { + "id": 8, + "type": "SaveAudio", + "pos": [1550, 200], + "size": [300, 100], + "inputs": [ + { + "name": "audio", + "type": "AUDIO", + "link": 10 + } + ], + "properties": { + "Node name for S&R": "SaveAudio" + }, + "widgets_values": ["styled_output.wav"] + } + ], + "links": [ + [1, 1, 0, 4, 0, "AUDIO"], + [2, 1, 0, 9, 0, "AUDIO"], + [3, 2, 0, 5, 0, "AUDIO"], + [4, 3, 0, 4, 1, "VAE"], + [5, 3, 0, 5, 1, "VAE"], + [6, 3, 0, 7, 1, "VAE"], + [7, 4, 0, 6, 0, "AUDIO_LATENT"], + [8, 5, 0, 6, 1, "AUDIO_LATENT"], + [9, 6, 0, 7, 0, "AUDIO_LATENT"], + [10, 7, 0, 8, 0, "AUDIO"] + ], + "config": {}, + "groups": [], + "version": 1, + "workflow": { + "name": "Audio Style Transfer", + "description": "Transfer audio style characteristics from one audio to another using latent space manipulation" + } +} \ No newline at end of file diff --git a/src/network_bending/nodes.py b/src/network_bending/nodes.py index 72556d0..8ca577b 100644 --- a/src/network_bending/nodes.py +++ b/src/network_bending/nodes.py @@ -5,6 +5,10 @@ import random import numpy as np from typing import Dict, List, Tuple, Any +# Import audio nodes +from .audio_nodes import NODE_CLASS_MAPPINGS as AUDIO_NODE_CLASS_MAPPINGS +from .audio_nodes import NODE_DISPLAY_NAME_MAPPINGS as AUDIO_NODE_DISPLAY_NAME_MAPPINGS + class NetworkBending: """ Network Bending Node - Performs various modifications on neural network models @@ -1363,6 +1367,9 @@ NODE_CLASS_MAPPINGS = { "VAEChannelManipulator": VAEChannelManipulator, } +# Add audio nodes to the mappings +NODE_CLASS_MAPPINGS.update(AUDIO_NODE_CLASS_MAPPINGS) + NODE_DISPLAY_NAME_MAPPINGS = { "NetworkBending": "Network Bending", "NetworkBendingAdvanced": "Network Bending (Advanced)", @@ -1374,6 +1381,9 @@ NODE_DISPLAY_NAME_MAPPINGS = { "VAEChannelManipulator": "VAE Channel Manipulator", } +# Add audio node display names +NODE_DISPLAY_NAME_MAPPINGS.update(AUDIO_NODE_DISPLAY_NAME_MAPPINGS) + # Export web directory for UI components WEB_DIRECTORY = "./js" __all__ = ['NODE_CLASS_MAPPINGS', 'NODE_DISPLAY_NAME_MAPPINGS', 'WEB_DIRECTORY'] \ No newline at end of file