Merge upstream main: keep audio nodes, upstream README, and add VAE/Latent nodes
This commit is contained in:
David Piazza
@@ -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.
|
||||
|
||||
##
|
||||
|
||||
@@ -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.
|
||||
@@ -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']
|
||||
@@ -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,)
|
||||
@@ -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)
|
||||
@@ -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"
|
||||
}
|
||||
}
|
||||
@@ -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']
|
||||
Reference in New Issue
Block a user