diff --git a/README.md b/README.md index 5c6cd2e..5f9a393 100644 --- a/README.md +++ b/README.md @@ -2,11 +2,18 @@ The New Stable Audio Open 1.0 Sampler In a ComfyUI Node. Make some beats! ![Screenshot from 2024-06-05 23-09-52](https://github.com/lks-ai/ComfyUI-StableAudioSampler/assets/163685473/037a23a7-0183-45b0-ae07-935664ba6dc7) -## Requirements -- At least 7GB VRAM -- ComfyUI +# Caveats +- The longer your audio, the more VRAM you need to stitch it together +- on a 3060, we've tried up to 10 seconds so far ## Installation + +### Download the Model and Config +1. Go to [Stable Audio Open on HuggingFace](https://huggingface.co/stabilityai/stable-audio-open-1.0/blob/main/) and download the `model.safetensors` and `model.config.json` files. +2. Place the files in the `models/audio_checkpoints` folder. If you don't have one, make one in your comfy folder. +3. Open Comfy and StableAudioLoader will see your model and config + +### With a HuggingFace Token 1. Make sure you have your `HF_TOKEN` environment variable for hugging face because model loading doesn't work just yet directly from a saved file 2. Go ahead and download model from here for when we fix that [Stable Audio Open on HuggingFace](https://huggingface.co/stabilityai/stable-audio-open-1.0/blob/main/model.safetensors) 3. Make sure to run `pip install -r requirements.txt` inside the repo folder if you're not using Manager @@ -17,30 +24,32 @@ There will definitely be issues because this is so new and it was coded quickly This is not an official StableAudioOpen repository. ## Current Features -- Uses HuggingFace -- Generates audio and outputs raw bytes and a sample rate +- Load your own models! +- Runs in half precision optional +- Nodes + - A Sampler Node: now with seed control, positive and negative prompts + - A Pre-Conditioning Node: kind of like empty latent audio with batch option + - A Prompt Node: Pipes conditioning + - A Model Loading Node: Includes repo options and scans `models/audio_checkpoints` for models and config.json files +- `control_after_generate` option +- Audio to Audio (like in the Gradio Example) +- Can still use HF env key if you want +- Generates audio and outputs raw bytes and a sample rate for use with VHS - Includes all of the original Stable Audio Open parameters +- Sampler outputs a Spectrogram image (experimental) - Can save audio to file ## Roadmap and Requested Features Keeping track of requests and ideas as they come in: -- Audio to Audio (like in the Gradio Example) -- Output to VHS Video Encoder Format - Stereo output -- Model loading from the `models/audio_checkpoints` folder -- Seed control with `control_after_generate` option - Nodes - A Mixer Node (mix your audio outputs with some sort of mastering) - A Tiling Sampler (concatenate the audios) - A Model Loader Node (to load audio models separately and pipe to wherever) -- Making the Audio format compatible with other Audio Node packs -- A bit of Refactoring cause this was a quick release - More Sampler Node Options - Gain - Possibly Clipping at some dB - Cleaning up some of the current options with selectors, etc. - - Negative Prompts -- Looking at Putting in InitAudio tomorrow so we can see Melspectograms (cause that looks bad ass, saw the gradio) We are very open to anyone who wants to contribute from the open source community. Make your forks and pull requests. We will build something cool. diff --git a/nodes.py b/nodes.py index 5aae45c..fc4e065 100644 --- a/nodes.py +++ b/nodes.py @@ -18,6 +18,7 @@ from stable_audio_tools.inference.generation import generate_diffusion_cond, gen from stable_audio_tools.inference.utils import prepare_audio from stable_audio_tools.training.utils import copy_state_dict from aeiou.viz import audio_spectrogram_image +from torchaudio import transforms as T # Test current setup # Add in Audio2Audio @@ -64,61 +65,41 @@ def repo_path(repo, filename): instance_path = instance_path.replace('\\', "/") return instance_path -# def generate_audio(cond_batch, steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, device, save, save_prefix, modelinfo, batch_size=1, seed=-1, after_generate="randomize", counter=0, init_noise_level=1.0): -# model, sample_rate, sample_size, _device = modelinfo -# b_pos, b_neg = cond_batch -# p_conditioning, p_batch_size = b_pos -# n_conditioning, n_batch_size = b_neg - -# output = generate_diffusion_cond( -# model, -# steps=steps, -# cfg_scale=cfg_scale, -# conditioning=p_conditioning, -# negative_conditioning=n_conditioning, -# sample_size=sample_size, -# sigma_min=sigma_min, -# sigma_max=sigma_max, -# sampler_type=sampler_type, -# device=_device, -# seed=seed, -# batch_size=p_batch_size, -# init_noise_level=init_noise_level, -# #return_latents=True, -# ) - -# # Rearrange audio batch to a single sequence -# output = rearrange(output, "b d n -> d (b n)") -# # Peak normalize, clip, convert to int16 -# output = output.to(torch.float32).div(torch.max(torch.abs(output))).clamp(-1, 1).mul(32767).to(torch.int16).cpu() - -# # Save Audio to output -# if save: -# save_audio_files(output, sample_rate, save_prefix, counter) - -# # Generate the spectrogram -# spectrogram = audio_spectrogram_image(output, sample_rate=sample_rate) - -# # Convert to bytes -# audio_bytes = output.numpy().tobytes() - -# # GC and memory Cleanup -# del model -# if torch.cuda.is_available(): -# torch.cuda.empty_cache() -# gc.collect() - -# return audio_bytes, sample_rate, spectrogram - -def generate_audio(cond_batch, steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, device, save, save_prefix, modelinfo, batch_size=1, seed=-1, after_generate="randomize", counter=0, init_noise_level=1.0): +def generate_audio(cond_batch, steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, device, save, save_prefix, modelinfo, batch_size=1, seed=-1, after_generate="randomize", counter=0, init_noise_level=1.0, init_audio=None): model, sample_rate, sample_size, _device = modelinfo b_pos, b_neg = cond_batch p_conditioning, p_batch_size = b_pos n_conditioning, n_batch_size = b_neg + sample_size = p_conditioning[0]['seconds_total'] * sample_rate - print("Model Loaded:", model) + #dprint("Model Loaded:", model) print("Positive Conditioning:", p_conditioning) print("Negative Conditioning:", n_conditioning) + print("Sample Size:", sample_size) + print("Sample Rate:", sample_rate) + print("Seconds:", sample_size / sample_rate) + + if init_audio is not None: + in_sr, init_audio = init_audio + # Turn into torch tensor, converting from int16 to float32 + init_audio = torch.from_numpy(init_audio).float().div(32767) + + if init_audio.dim() == 1: + init_audio = init_audio.unsqueeze(0) # [1, n] + elif init_audio.dim() == 2: + init_audio = init_audio.transpose(0, 1) # [n, 2] -> [2, n] + + if in_sr != sample_rate: + resample_tf = T.Resample(in_sr, sample_rate).to(init_audio.device) + init_audio = resample_tf(init_audio) + + audio_length = init_audio.shape[-1] + + if audio_length > sample_size: + + input_sample_size = audio_length + (model.min_input_length - (audio_length % model.min_input_length)) % model.min_input_length + + init_audio = (sample_rate, init_audio) output = generate_diffusion_cond( model, @@ -133,7 +114,8 @@ def generate_audio(cond_batch, steps, cfg_scale, sample_size, sigma_min, sigma_m device=_device, seed=seed, batch_size=p_batch_size, - init_noise_level=init_noise_level, + init_noise_level=init_noise_level, + init_audio=init_audio ) print("Raw Output:", output) @@ -158,6 +140,7 @@ def generate_audio(cond_batch, steps, cfg_scale, sample_size, sigma_min, sigma_m return audio_bytes, sample_rate, spectrogram + def get_model(model_filename=None, config=None, repo=None, half_precision=False, device_override=None): #print(model_filename, config, repo, half_precision) if model_filename: @@ -276,6 +259,9 @@ class StableAudioSampler: "denoise": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 1.0, "step": 0.01}), "save": ("BOOLEAN", {"default": True}), "save_prefix": ("STRING", {"default": "StableAudio"}), + }, + "optional": { + "audio": ("VHS_AUDIO", ) } } @@ -286,8 +272,8 @@ class StableAudioSampler: CATEGORY = "audio/samplers" - def sample(self, audio_model, positive, negative, seed, steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, denoise, save, save_prefix): - audio_bytes, sample_rate, spectrogram = generate_audio((positive, negative), steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, device, save, save_prefix, audio_model, seed=seed, counter=self.counter, init_noise_level=denoise) + def sample(self, audio_model, positive, negative, seed, steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, denoise, save, save_prefix, audio=None): + audio_bytes, sample_rate, spectrogram = generate_audio((positive, negative), steps, cfg_scale, sample_size, sigma_min, sigma_max, sampler_type, device, save, save_prefix, audio_model, seed=seed, counter=self.counter, init_noise_level=denoise, init_audio=audio) spectrograms = create_image_batch([spectrogram], 1) return (audio_bytes, sample_rate, spectrograms) @@ -323,7 +309,7 @@ class StableAudioPrompt: return { "required": { "conditioning": ("SAOCOND", {"forceInput": True}), - "prompt": ("STRING", {"default": "", "multiline": True}), + "prompt": ("STRING", {"multiline": True}), } } @@ -334,15 +320,22 @@ class StableAudioPrompt: CATEGORY = "audio/conditioning" def go(self, conditioning, prompt): + print("PROMPT", prompt) cond, batch_size = conditioning - cond[0]['prompt'] = prompt + print(cond, batch_size) + o = [] + #cond[0]['prompt'] = prompt + for v in cond: + v['prompt'] = prompt + o.append(v.copy()) #c = conditioning[0] # conditioning = [{ # "prompt": prompt, # "seconds_start": seconds_start, # "seconds_total": seconds_total # }] - return ((cond, batch_size), ) + print(o, batch_size) + return ((o, batch_size), ) class StableAudioConditioning: @classmethod