From c920a4e714ab87adcc4e09ff189fe938194f2a4e Mon Sep 17 00:00:00 2001 From: gokayfem Date: Sun, 25 Feb 2024 20:29:16 +0300 Subject: [PATCH] img2music --- __init__.py | 2 + js/playSound.js | 55 +++++++++++++++++++++++++++ nodes/audioldm2.py | 95 ++++++++++++++++++++++++++++++++++++++++++++++ nodes/playmusic.py | 42 ++++++++++++++++++++ requirements.txt | 4 +- 5 files changed, 197 insertions(+), 1 deletion(-) create mode 100644 js/playSound.js create mode 100644 nodes/audioldm2.py create mode 100644 nodes/playmusic.py diff --git a/__init__.py b/__init__.py index ec2ea7f..325f177 100644 --- a/__init__.py +++ b/__init__.py @@ -42,6 +42,8 @@ node_list = [ "internlm", "uform", "kosmos2", + "audioldm2", + "playmusic", ] NODE_CLASS_MAPPINGS = {} diff --git a/js/playSound.js b/js/playSound.js new file mode 100644 index 0000000..33ddbbf --- /dev/null +++ b/js/playSound.js @@ -0,0 +1,55 @@ +import { app } from "/scripts/app.js"; + +app.registerExtension({ + name: "n.PlayMusic", + async beforeRegisterNodeDef(nodeType, nodeData, app) { + if (nodeData.name === "PlayMusic") { + console.warn("PlayMusic"); + const onExecuted = nodeType.prototype.onExecuted; + nodeType.prototype.onExecuted = async function () { + onExecuted?.apply(this, arguments); + + // Check for "on empty queue" condition, if applicable + if (this.widgets[0].value === "on empty queue") { + if (app.ui.lastQueueSize !== 0) { + await new Promise((r) => setTimeout(r, 500)); + } + if (app.ui.lastQueueSize !== 0) { + return; + } + } + + // Assuming that 'arguments[0].a' is the waveform and 'arguments[0].b' is the sample rate + let waveform = arguments[0].a; // An array of floats (-1 to 1) + let sampleRate = arguments[0].b; // The sample rate of the audio + console.log(waveform, sampleRate); + // Create AudioContext + let audioCtx = new (window.AudioContext || window.webkitAudioContext)({sampleRate: sampleRate}); + + // Create AudioBuffer + let buffer = audioCtx.createBuffer(1, waveform[0].length, sampleRate); + + // Fill the AudioBuffer + buffer.getChannelData(0).set(waveform[0]); + + // Create a source and connect it to the buffer + let source = audioCtx.createBufferSource(); + source.buffer = buffer; + source.connect(audioCtx.destination); + + // Set volume, if applicable. Assuming the volume is the second widget's value. + let volume = this.widgets[1].value; + if (volume !== undefined) { + let gainNode = audioCtx.createGain(); + gainNode.gain.value = volume; + source.connect(gainNode); + gainNode.connect(audioCtx.destination); + } + + // Play the sound + source.start(); + }; + } + }, +}); + diff --git a/nodes/audioldm2.py b/nodes/audioldm2.py new file mode 100644 index 0000000..d1f0ac1 --- /dev/null +++ b/nodes/audioldm2.py @@ -0,0 +1,95 @@ +from diffusers import AudioLDM2Pipeline +from huggingface_hub import snapshot_download +from pathlib import Path +import torch +import os +import soundfile as sf +from folder_paths import output_directory + +# Define the directory for saving files related to the audio model +files_for_audio_model = Path(__file__).resolve().parent / "files_for_audioldm2" +files_for_audio_model.mkdir(parents=True, exist_ok=True) # Ensure the directory exists + +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False +base_path = os.path.dirname(os.path.realpath(__file__)) + +# Our any instance wants to be a wildcard string +any = AnyType("*") +class AudioLDM2ModelPredictor: + def __init__(self): + self.device = "cuda" if torch.cuda.is_available() else "cpu" + torch_dtype = torch.float16 if self.device == "cuda" else torch.float32 + + # Use snapshot_download to manage the model download/cache + self.model_path = snapshot_download("cvssp/audioldm2", + local_dir=files_for_audio_model, + force_download=False, # Set to True to always download + local_files_only=False, # Download if not available locally + use_auth_token=False, # Set to True if using a private model + local_dir_use_symlinks="auto", # Auto-manage symlinks + ignore_patterns=["*.bin", "*.jpg", "*.png"]) # Ignore unrelated files + + self.pipeline = AudioLDM2Pipeline.from_pretrained(self.model_path, + torch_dtype=torch_dtype).to(self.device) + self.generator = torch.Generator(self.device) + + def generate_audio(self, text, negative_prompt, duration, guidance_scale, random_seed, sample_rate, n_candidates=1, extension="wav"): + if text is None: + raise ValueError("Please provide a text input.") + + # Manual seed for reproducibility + self.generator.manual_seed(int(random_seed)) + + # Generate audio + waveforms = self.pipeline( + text, + audio_length_in_s=duration, + guidance_scale=guidance_scale, + num_inference_steps=200, + negative_prompt=negative_prompt, + num_waveforms_per_prompt=n_candidates, + generator=self.generator, + )["audios"] + + # Save the generated waveform to a file + audio_path = Path(output_directory) / f"generated_audio_{random_seed}.{extension}" + + sf.write(audio_path , waveforms[0], sample_rate) + final_waveforms = waveforms[0].tolist() + return (final_waveforms, sample_rate) # Return the path of the generated audio file + + +class AudioLDM2Node: + def __init__(self): + self.predictor = AudioLDM2ModelPredictor() + + @classmethod + def INPUT_TYPES(cls): + return { + "required": { + "text": ("STRING",{"default": "", "forceInput": True}), + "negative_prompt": ("STRING",{"default": "", "forceInput": True}), + "duration": ("INT",{"default": 10, "min": 1, "max": 60, "step": 1}), + "guidance_scale": ("FLOAT", {"default": 3.5, "min": 0.1, "max": 20.0, "step": 0.1}), + "seed": ("INT", {"default": 42, "step": 1}), + "n_candidates": ("INT", {"default": 3, "min": 1, "max": 10, "step": 1}), + "sample_rate": ("INT", {"default": 16000, "min": 8000, "max": 48000, "step": 1}), + "extension": (["wav", "mp3", "flac"], {"default": "wav"}), + } + } + + RETURN_NAMES = ("wave_form", "sample_rate", ) + RETURN_TYPES = (any, "INT", ) + OUTPUT_NODE = True + FUNCTION = "generate_audio_final" + + CATEGORY = "VLM Nodes/AudioLDM2" + + def generate_audio_final(self, text, negative_prompt, duration, guidance_scale, sample_rate, seed, n_candidates, extension): + wave_form, sample_rate_final = self.predictor.generate_audio(text, negative_prompt, duration, guidance_scale, seed, sample_rate, n_candidates, extension) + return (wave_form, sample_rate_final, ) + +NODE_CLASS_MAPPINGS = {"AudioLDM2Node": AudioLDM2Node} +NODE_DISPLAY_NAME_MAPPINGS = {"AudioLDM2Node": "AudioLDM-2 Node"} diff --git a/nodes/playmusic.py b/nodes/playmusic.py new file mode 100644 index 0000000..b1a6210 --- /dev/null +++ b/nodes/playmusic.py @@ -0,0 +1,42 @@ +# Hack: string type that is always equal in not equal comparisons +class AnyType(str): + def __ne__(self, __value: object) -> bool: + return False + + +# Our any instance wants to be a wildcard string +any = AnyType("*") + + +class PlayMusic: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "mode": (["always", "on empty queue"], {}), + "volume": ("FLOAT", {"min": 0, "max": 1, "step": 0.1, "default": 0.5}), + "wave_form": ([], {"forceInput": True}), + "sample_rate": ("INT", {"forceInput": True}), + }} + + FUNCTION = "nop" + INPUT_IS_LIST = True + OUTPUT_IS_LIST = (True,) + OUTPUT_NODE = True + RETURN_TYPES = (any,) + + CATEGORY = "VLM Nodes/AudioLDM2" + + def IS_CHANGED(self, **kwargs): + return float("NaN") + + def nop(self, mode, volume, wave_form, sample_rate): + return {"ui": {"a": wave_form, "b": sample_rate}, "result": (any,)} + + +NODE_CLASS_MAPPINGS = { + "PlaySound": PlayMusic, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "PlaySound": "PlaySound Node", +} diff --git a/requirements.txt b/requirements.txt index d3f3f46..3c36d92 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,6 +1,6 @@ openai>=0.27.8 accelerate>=0.25.0 -huggingface-hub==0.20.1 +huggingface-hub==0.20.3 Pillow>=10.1.0 transformers>=4.36.2 torch>=2.0.1,<3.0.0 @@ -19,3 +19,5 @@ pytz six cffi python-dateutil>=2.7.0 +diffusers +soundfile \ No newline at end of file