New Features: Added nodes for advanced video and audio processing
Key updates: Added AddFramesToVideoFolder node to add frames to a video folder Introduced VideoFolderProvider and VideoDefinitionProvider nodes to manage video folders and video definitions Added FadeToInVideoFolder node for fading between frames in a video folder Implemented ListedSVDConditioning node for applying SVD conditioning to a list of parameters Added VideoFolderFFMPEGVideoInterpolate node to interpolate frames in a video folder using FFmpeg Introduced AudioFolderProvider and AudioDefinitionProvider nodes for managing audio folders and definitions Added AddAudioToAudioFolder and AddSilenceToAudioFolder nodes for adding audio files and silence to audio folders Implemented GetAudioFilesLengthsFromAudioFolder node to retrieve lengths of audio files in a folder Added AudioFolderRenderAudio node to join audio files in an audio folder Introduced GenerateMusic node for generating music from descriptions using the AudioCraft library Added XttsNode for text-to-speech conversion using the XTTS library Implemented AdjustAudioVolumeForFiles node to adjust volume for a list of audio files The changes also include the addition of LiteLLM nodes for text generation: Added FuncAnysToText node that runs Python code with multiple input types and outputs a string Introduced FuncAnysToLatent node for running code that outputs a latent representation Added FuncAnysToString node for executing code with various inputs and producing a string output Implemented FuncAnysToList node to run code that takes different input types and returns a list Added FuncAnysToCond node for running code with multiple inputs and generating a conditioning output These new nodes expand the capabilities for video and audio processing, allowing more advanced manipulations and generation of video and audio content within the graph.
This commit is contained in:
Tasha Upchurch
@@ -0,0 +1,404 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import cv2
|
||||
import time
|
||||
from PIL import Image
|
||||
import numpy as np
|
||||
import scipy.ndimage
|
||||
|
||||
|
||||
class MaskExpansionModel(nn.Module):
|
||||
"""
|
||||
A PyTorch model that expands or erodes a mask using the provided function.
|
||||
"""
|
||||
|
||||
def forward(self, mask, expand, tapered_corners):
|
||||
assert mask.dim() == 4, "Input mask must have 3 dimensions (batch_size, height, width)."
|
||||
bs = mask.size(0)
|
||||
c = 0 if tapered_corners else 1
|
||||
kernel = torch.tensor([[c, 1, c],
|
||||
[1, 1, 1],
|
||||
[c, 1, c]], dtype=mask.dtype, device=mask.device)
|
||||
kernel = kernel.view(1, 1, 3, 3)
|
||||
|
||||
if expand < 0:
|
||||
kernel = -kernel
|
||||
expand = -expand
|
||||
mask = 1 - mask
|
||||
|
||||
mask = mask.unsqueeze(1)
|
||||
# print(mask.shape)
|
||||
for _ in range(expand):
|
||||
for i in range(bs):
|
||||
# print(i)
|
||||
mask[i] = torch.nn.functional.conv2d(mask[i], kernel, padding=1)
|
||||
mask[i] = torch.clamp(mask[i], 0, 1) # Clip the mask values to the range [0, 1]
|
||||
|
||||
mask = mask.squeeze(1)
|
||||
if expand < 0:
|
||||
mask = 1 - mask
|
||||
|
||||
return mask
|
||||
|
||||
|
||||
class LanczosInterpolationLayer(nn.Module):
|
||||
"""
|
||||
A PyTorch module that applies Lanczos interpolation to resize a batch of images.
|
||||
|
||||
Args:
|
||||
target_size (tuple): The target size for resizing the images, in the format (width, height).
|
||||
"""
|
||||
|
||||
def __init__(self, target_size):
|
||||
super(LanczosInterpolationLayer, self).__init__()
|
||||
self.target_size = target_size
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Performs the forward pass of the Lanczos interpolation layer.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
|
||||
"""
|
||||
device = torch.device("cpu")
|
||||
interpolated_images = []
|
||||
for image in x:
|
||||
# Convert the PyTorch tensor to a NumPy array
|
||||
image_np = image.permute(1, 2, 0).to("cpu", torch.float32).numpy()
|
||||
|
||||
# Resize the image using OpenCV's resize function with Lanczos interpolation
|
||||
interpolated_image = cv2.resize(image_np, self.target_size, interpolation=cv2.INTER_LANCZOS4)
|
||||
mn = interpolated_image.min()
|
||||
mx = interpolated_image.max()
|
||||
|
||||
color_range = mx - mn
|
||||
|
||||
under_bottom_target_range = -10/255
|
||||
over_top_target_range = 1/255
|
||||
|
||||
if color_range > 1:
|
||||
# here we are going to compress the over the top values
|
||||
over_top = interpolated_image > 1
|
||||
not_over_top = interpolated_image < 1
|
||||
over_top_range = mx - 1
|
||||
# subtract the normal top value
|
||||
interpolated_image[over_top] = interpolated_image[over_top] - 1
|
||||
# scale the over the top values to the target range
|
||||
interpolated_image[over_top] *= (over_top_target_range / over_top_range)
|
||||
# add the target range to the over the top values
|
||||
interpolated_image[over_top] += (1-over_top_target_range)
|
||||
# scale the non-over-the top to make room
|
||||
interpolated_image[not_over_top] *= (1 - over_top_target_range)
|
||||
|
||||
# here we are going to compress the under the bottom values
|
||||
under_bottom = interpolated_image < 0
|
||||
not_under_bottom = interpolated_image > 0
|
||||
under_bottom_range = mn
|
||||
# subtract the normal bottom value
|
||||
interpolated_image[under_bottom] = interpolated_image[under_bottom] - 0
|
||||
# scale the under the bottom values to the target range
|
||||
interpolated_image[under_bottom] *= (under_bottom_target_range / under_bottom_range)
|
||||
# add the target range to the under the bottom values
|
||||
interpolated_image[under_bottom] += under_bottom_target_range
|
||||
# scale the non-under-the-bottom to make room
|
||||
interpolated_image[not_under_bottom] *= (1 - under_bottom_target_range)
|
||||
|
||||
|
||||
|
||||
|
||||
else:
|
||||
pass
|
||||
|
||||
# Convert the resized image back to a PyTorch tensor
|
||||
# interpolated_image = torch.from_numpy(interpolated_image).to(device, dtype=torch.float16).permute(2, 0, 1)
|
||||
interpolated_image = torch.from_numpy(interpolated_image)
|
||||
|
||||
# convert fix range to 0-255
|
||||
interpolated_image = interpolated_image * 255
|
||||
# convert to uint8
|
||||
interpolated_image = interpolated_image.to(device, dtype=torch.uint8).permute(2, 0, 1)
|
||||
|
||||
interpolated_images.append(interpolated_image)
|
||||
del image
|
||||
del image_np
|
||||
del interpolated_image
|
||||
|
||||
# transform this list of tensors into a single tensor of shape (batch_size, channels, height, width)
|
||||
|
||||
# pre allocate a tensor of shape (batch_size, channels, target_height, target_width)
|
||||
y = torch.zeros((x.shape[0], interpolated_images[0].shape[0], self.target_size[1], self.target_size[0]),
|
||||
device=device,
|
||||
dtype=interpolated_images[0].dtype)
|
||||
y = torch.stack(interpolated_images, dim=0, out=y)
|
||||
del interpolated_images
|
||||
|
||||
return y
|
||||
|
||||
|
||||
class LanczosInterpolationModel(nn.Module):
|
||||
"""
|
||||
A PyTorch model that applies Lanczos interpolation to resize a batch of images.
|
||||
|
||||
Args:
|
||||
target_size (tuple): The target size for resizing the images, in the format (width, height).
|
||||
"""
|
||||
|
||||
def __init__(self, target_size):
|
||||
super(LanczosInterpolationModel, self).__init__()
|
||||
self.interpolation_layer = LanczosInterpolationLayer(target_size)
|
||||
|
||||
def forward(self, x):
|
||||
"""
|
||||
Performs the forward pass of the Lanczos interpolation model.
|
||||
|
||||
Args:
|
||||
x (torch.Tensor): A batch of input images with shape (batch_size, channels, height, width).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: A batch of resized images with shape (batch_size, channels, target_height, target_width).
|
||||
"""
|
||||
return self.interpolation_layer(x)
|
||||
|
||||
|
||||
def expand_mask(mask, expand, tapered_corners, device=None, dtype=None):
|
||||
"""
|
||||
Convenience function for expanding or eroding a mask using the MaskExpansionModel.
|
||||
|
||||
Args:
|
||||
mask (torch.Tensor): The input mask tensor.
|
||||
expand (int): The number of iterations to expand (positive) or erode (negative) the mask.
|
||||
tapered_corners (bool): Whether to use tapered corners when expanding or eroding the mask.
|
||||
device (torch.device, optional): The device to use for the model and input tensor.
|
||||
If not provided, it will use the default device (CPU).
|
||||
dtype (torch.dtype, optional): The data type to use for the input tensor.
|
||||
If not provided, it will use the default dtype (float32).
|
||||
|
||||
Returns:
|
||||
torch.Tensor: The expanded or eroded mask tensor.
|
||||
"""
|
||||
# Create an instance of the MaskExpansionModel
|
||||
model = MaskExpansionModel()
|
||||
|
||||
# Move the model to the specified device (or default to CPU)
|
||||
device = device or torch.device("cpu")
|
||||
model = model.to(device)
|
||||
|
||||
# Move the input tensor to the specified device and convert to the specified dtype
|
||||
dtype = dtype or torch.float32
|
||||
mask = mask.to(device, dtype=dtype)
|
||||
|
||||
# Expand or erode the mask using the MaskExpansionModel
|
||||
with torch.no_grad():
|
||||
expanded_mask = model(mask, expand, tapered_corners)
|
||||
|
||||
# Clear the GPU cache if using a GPU
|
||||
if device != torch.device("cpu"):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return expanded_mask
|
||||
|
||||
|
||||
def lanczos_resize(image, target_size, device="cpu", dtype=torch.float32, output_type="tensor"):
|
||||
"""
|
||||
Convenience function for resizing an image using the Lanczos interpolation model.
|
||||
|
||||
Args:
|
||||
image (PIL.Image or torch.Tensor): The input image to be resized.
|
||||
target_size (tuple): The target size for resizing the image, in the format (width, height).
|
||||
device (torch.device, optional): The device to use for the model and input tensor.
|
||||
If not provided, it will use the default device (CPU).
|
||||
dtype (torch.dtype, optional): The data type to use for the input tensor.
|
||||
If not provided, it will use the default dtype (float32).
|
||||
output_type (str, optional): The type of the output image.
|
||||
If "tensor", the function will return a PyTorch tensor.
|
||||
If "pil", the function will return a PIL Image.
|
||||
If not provided, it will return a tensor by default.
|
||||
|
||||
Returns:
|
||||
torch.Tensor or PIL.Image: The resized image.
|
||||
"""
|
||||
import gc
|
||||
import os
|
||||
import sys
|
||||
if sys.platform == "win32":
|
||||
import psutil
|
||||
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
|
||||
else:
|
||||
# works on linux and mac
|
||||
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
|
||||
|
||||
# Create an instance of the LanczosInterpolationModel
|
||||
model = LanczosInterpolationModel(target_size)
|
||||
|
||||
# Move the model to the specified device (or default to CPU)
|
||||
device = device or torch.device("cpu")
|
||||
model = model.to(device)
|
||||
|
||||
# Convert the input image to a PyTorch tensor
|
||||
if isinstance(image, list):
|
||||
if isinstance(image[0], Image.Image):
|
||||
try:
|
||||
images = np.array(image)
|
||||
except:
|
||||
images = np.array([np.array(img) for img in image])
|
||||
|
||||
images = torch.from_numpy(images).to(device, dtype).permute(0, 3, 1, 2)
|
||||
elif isinstance(image, torch.Tensor):
|
||||
images = image.to(device, dtype)
|
||||
|
||||
del image
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Resize the image using the Lanczos interpolation model
|
||||
with torch.no_grad():
|
||||
resized_images = model(images)
|
||||
|
||||
# Convert the resized tensor back to a PIL Image or keep it as a tensor based on the output_type
|
||||
if output_type.lower() == "pil":
|
||||
resized_images = resized_images.squeeze(0).permute(0, 2, 3, 1).cpu().numpy()
|
||||
resized_images = Image.fromarray(resized_images.astype(np.uint8))
|
||||
elif output_type.lower() == "tensor":
|
||||
pass # keep the resized_image as a tensor
|
||||
else:
|
||||
raise ValueError(f"Invalid output_type: {output_type}. Valid options are 'tensor' or 'pil'.")
|
||||
|
||||
# Clear the GPU cache if using a GPU
|
||||
if device != torch.device("cpu"):
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
del images
|
||||
del model
|
||||
gc.collect()
|
||||
# print current free memory on cpu, in a way that is cross platform 1 way that works on all
|
||||
|
||||
if sys.platform == "win32":
|
||||
import psutil
|
||||
print(f"Free memory: {psutil.virtual_memory().available / 1024 / 1024} MB")
|
||||
else:
|
||||
# works on linux and mac
|
||||
print(f"Free memory: {os.popen('free -m').readlines()[1].split()[3]} MB")
|
||||
|
||||
return resized_images
|
||||
|
||||
|
||||
def test_lanczos_interpolation_model():
|
||||
"""
|
||||
Tests the Lanczos interpolation model by resizing a batch of random images and benchmarking the performance.
|
||||
"""
|
||||
# Set the target sizes for resizing
|
||||
hd_size = (1920, 1080)
|
||||
sd_size = (1280, 720)
|
||||
|
||||
# Generate random PIL images
|
||||
batch_size = 600
|
||||
hd_images = [Image.fromarray(np.random.randint(0, 256, size=(hd_size[1], hd_size[0], 3), dtype=np.uint8))
|
||||
for _ in range(batch_size)]
|
||||
sd_images = [Image.fromarray(np.random.randint(0, 256, size=(sd_size[1], sd_size[0], 3), dtype=np.uint8))
|
||||
for _ in range(batch_size)]
|
||||
|
||||
# Perform the benchmarks
|
||||
num_iterations = 3
|
||||
|
||||
# HD to SD benchmark
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
sd_output_images = lanczos_resize(hd_images, sd_size)
|
||||
end_time = time.time()
|
||||
hd_to_sd_elapsed_time = end_time - start_time
|
||||
hd_to_sd_avg_time_per_iteration = hd_to_sd_elapsed_time / num_iterations
|
||||
|
||||
# SD to HD benchmark
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
hd_output_images = lanczos_resize(sd_images, hd_size)
|
||||
end_time = time.time()
|
||||
sd_to_hd_elapsed_time = end_time - start_time
|
||||
sd_to_hd_avg_time_per_iteration = sd_to_hd_elapsed_time / num_iterations
|
||||
|
||||
# Check the sizes of the output images
|
||||
|
||||
assert sd_output_images.shape[0] == batch_size
|
||||
|
||||
assert hd_output_images.shape[0] == batch_size
|
||||
|
||||
print(f"Lanczos interpolation test passed!")
|
||||
print(f"Benchmark results:")
|
||||
print(f"Number of iterations: {num_iterations}")
|
||||
print(f"HD to SD:")
|
||||
print(f" Total elapsed time: {hd_to_sd_elapsed_time:.4f} seconds")
|
||||
print(f" Average time per iteration: {hd_to_sd_avg_time_per_iteration:.4f} seconds")
|
||||
print(f" Average time per HD image -> SD image: {hd_to_sd_avg_time_per_iteration / batch_size:.4f} seconds")
|
||||
print(f"SD to HD:")
|
||||
print(f" Total elapsed time: {sd_to_hd_elapsed_time:.4f} seconds")
|
||||
print(f" Average time per iteration: {sd_to_hd_avg_time_per_iteration:.4f} seconds")
|
||||
print(f" Average time per SD image -> HD image: {sd_to_hd_avg_time_per_iteration / batch_size:.4f} seconds")
|
||||
|
||||
|
||||
def test_mask_expansion_model():
|
||||
"""
|
||||
Tests the MaskExpansionModel by expanding and eroding random masks and benchmarking the performance.
|
||||
"""
|
||||
# Generate random masks
|
||||
batch_size = 1000
|
||||
mask_size = (512, 512)
|
||||
# masks are 0-1 tensors float
|
||||
masks = torch.rand((batch_size, mask_size[0], mask_size[1])).to("cuda")
|
||||
print(masks.min(), masks.max())
|
||||
|
||||
# Perform the benchmarks
|
||||
num_iterations = 1
|
||||
expand_steps = 5
|
||||
erode_steps = -5
|
||||
|
||||
# Mask expansion benchmark
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
print(masks.mean())
|
||||
expanded_masks = expand_mask(masks, expand_steps, tapered_corners=True, device="cuda")
|
||||
print(masks.mean())
|
||||
end_time = time.time()
|
||||
expansion_elapsed_time = end_time - start_time
|
||||
expansion_avg_time_per_iteration = expansion_elapsed_time / num_iterations
|
||||
|
||||
# Mask erosion benchmark
|
||||
start_time = time.time()
|
||||
for _ in range(num_iterations):
|
||||
print(masks.mean())
|
||||
eroded_masks = expand_mask(masks, erode_steps, tapered_corners=True, device="cuda")
|
||||
print(eroded_masks.mean())
|
||||
end_time = time.time()
|
||||
erosion_elapsed_time = end_time - start_time
|
||||
erosion_avg_time_per_iteration = erosion_elapsed_time / num_iterations
|
||||
|
||||
# Check the sizes of the output masks
|
||||
expected_mask_size = (batch_size, mask_size[0], mask_size[1])
|
||||
assert expanded_masks.shape == expected_mask_size, f"Expected expanded mask size: {expected_mask_size}, but got: {expanded_masks.shape}"
|
||||
assert eroded_masks.shape == expected_mask_size, f"Expected eroded mask size: {expected_mask_size}, but got: {eroded_masks.shape}"
|
||||
|
||||
print(f"Mask expansion/erosion test passed!")
|
||||
print(f"Benchmark results:")
|
||||
print(f"Number of iterations: {num_iterations}")
|
||||
print(f"Mask expansion:")
|
||||
print(f" Total elapsed time: {expansion_elapsed_time:.4f} seconds")
|
||||
print(f" Average time per iteration: {expansion_avg_time_per_iteration:.4f} seconds")
|
||||
print(f" Average time per mask expansion: {expansion_avg_time_per_iteration / batch_size:.4f} seconds")
|
||||
print(f"Mask erosion:")
|
||||
print(f" Total elapsed time: {erosion_elapsed_time:.4f} seconds")
|
||||
print(f" Average time per iteration: {erosion_avg_time_per_iteration:.4f} seconds")
|
||||
print(f" Average time per mask erosion: {erosion_avg_time_per_iteration / batch_size:.4f} seconds")
|
||||
|
||||
|
||||
if __name__ == "image_utils":
|
||||
print("image_utils module loaded.")
|
||||
print("Testing Lanczos interpolation model...")
|
||||
test_lanczos_interpolation_model()
|
||||
print("Testing Mask expansion model...")
|
||||
# test_mask_expansion_model()
|
||||
exit(0)
|
||||
|
||||
if __name__ == "__main__":
|
||||
test_lanczos_interpolation_model()
|
||||
# test_mask_expansion_model()
|
||||
@@ -0,0 +1,224 @@
|
||||
def set_all_seeds(seed):
|
||||
import random
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
old_values = {}
|
||||
old_values['random'] = random.getstate()
|
||||
old_values['numpy'] = np.random.get_state()
|
||||
old_values['torch'] = torch.get_rng_state()
|
||||
old_values['torch_cuda'] = torch.cuda.get_rng_state()
|
||||
old_values['torch_cuda_all'] = torch.cuda.get_rng_state_all()
|
||||
old_values['torch_deterministic'] = torch.backends.cudnn.deterministic
|
||||
|
||||
# From https://gist.github.com/gatheluck/c57e2a40e3122028ceaecc3cb0d152ac
|
||||
random.seed(seed)
|
||||
os.environ["PYTHONHASHSEED"] = str(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
torch.cuda.manual_seed(seed)
|
||||
torch.backends.cudnn.deterministic = True
|
||||
|
||||
return old_values
|
||||
|
||||
|
||||
def reset_seeds(old_values):
|
||||
import random
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
random.setstate(old_values['random'])
|
||||
np.random.set_state(old_values['numpy'])
|
||||
torch.set_rng_state(old_values['torch'])
|
||||
torch.cuda.set_rng_state(old_values['torch_cuda'])
|
||||
torch.cuda.set_rng_state_all(old_values['torch_cuda_all'])
|
||||
torch.backends.cudnn.deterministic = old_values['torch_deterministic']
|
||||
|
||||
import time # Add this at the top of your script
|
||||
|
||||
def generate_music_for_description(description, duration, file, model_name, position):
|
||||
time.sleep(position * 5) # Wait for position * 5 seconds
|
||||
from audiocraft.models import MusicGen
|
||||
from audiocraft.data.audio import audio_write
|
||||
import torch
|
||||
model = MusicGen.get_pretrained(model_name)
|
||||
print(f"Generating music for description: {description}")
|
||||
print(f"Using model: {model_name}")
|
||||
print(f"Using duration: {duration}")
|
||||
model.set_generation_params(duration=duration)
|
||||
wav = model.generate([description])
|
||||
print(f"Writing to {file}")
|
||||
audio_write(file.replace(".wav", ""), wav[0].cpu(), model.sample_rate, strategy="loudness")
|
||||
del model
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
|
||||
def generate_music(**kwargs):
|
||||
from audiocraft.models import MusicGen
|
||||
from audiocraft.data.audio import audio_write
|
||||
import torch
|
||||
import ast
|
||||
import copy
|
||||
import random
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
# Set seed, -1 means random seed
|
||||
seed = kwargs.get('seed', -1)
|
||||
if seed == -1:
|
||||
seed = random.randint(0, 100000)
|
||||
old_seed_values = set_all_seeds(seed)
|
||||
|
||||
kwargs = copy.deepcopy(kwargs)
|
||||
description = kwargs.get('description')
|
||||
|
||||
file = kwargs.get('file')
|
||||
# Add .wav if not there
|
||||
if not file.endswith('.wav'):
|
||||
file = file + ".wav"
|
||||
|
||||
model_name = kwargs.get('model_name', 'melody')
|
||||
duration = kwargs.get('duration', 10)
|
||||
|
||||
descriptions = ast.literal_eval(description)
|
||||
durations = ast.literal_eval(duration)
|
||||
|
||||
|
||||
|
||||
o_files = []
|
||||
batch_size = 3
|
||||
with ProcessPoolExecutor(max_workers=batch_size) as executor:
|
||||
futures = []
|
||||
for i, (description, duration) in enumerate(zip(descriptions, durations)):
|
||||
ii = str(i).zfill(1 + len(str(1 + len(descriptions))))
|
||||
file_to_use = file.replace(".wav", f"_{ii}.wav")
|
||||
# Pass i (the position in the batch) as an additional argument
|
||||
future = executor.submit(generate_music_for_description, description, duration, file_to_use, model_name, i//batch_size)
|
||||
futures.append((future, file_to_use))
|
||||
|
||||
for future, file_to_use in futures:
|
||||
# Wait for the future to complete, no need to get the result since the function handles file writing
|
||||
future.result()
|
||||
o_files.append(file_to_use)
|
||||
|
||||
# Reset seed
|
||||
reset_seeds(old_seed_values)
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return (o_files[0], o_files,)
|
||||
|
||||
|
||||
import argparse
|
||||
import copy
|
||||
|
||||
|
||||
def main():
|
||||
print("Running main function...")
|
||||
# Define the model names as a list for the choices parameter
|
||||
model_names = [
|
||||
"facebook/musicgen-melody-large",
|
||||
"facebook/musicgen-melody-small",
|
||||
"facebook/musicgen-melody-medium",
|
||||
|
||||
"facebook/musicgen-stereo-melody-large",
|
||||
"facebook/musicgen-stereo-melody-small",
|
||||
"facebook/musicgen-stereo-melody-medium",
|
||||
|
||||
"facebook/musicgen-large",
|
||||
"facebook/musicgen-small",
|
||||
"facebook/musicgen-medium",
|
||||
|
||||
"facebook/musicgen-stereo-large",
|
||||
"facebook/musicgen-stereo-small",
|
||||
"facebook/musicgen-stereo-medium",
|
||||
]
|
||||
|
||||
# Parse command line arguments
|
||||
parser = argparse.ArgumentParser(description='Generate music based on a given description.')
|
||||
parser.add_argument('--description', help='Description of the music to generate')
|
||||
parser.add_argument('--file', type=str, default='output.wav', help='Output file name (default: output.wav)')
|
||||
parser.add_argument('--model_name', type=str, default='facebook/musicgen-large', choices=model_names,
|
||||
help='Model name to use for generation')
|
||||
parser.add_argument('--duration', help='Duration of the generated music in seconds (default: 10)')
|
||||
|
||||
parser.add_argument('--seed', type=int, default=-1,
|
||||
help='Seed for random number generation (default: -1 for random seed)')
|
||||
parser.add_argument('--test', action='store_true', help='Run the test suite instead of generating music')
|
||||
args = parser.parse_args()
|
||||
|
||||
# Convert argparse.Namespace to dictionary
|
||||
kwargs = vars(args)
|
||||
|
||||
if args.test:
|
||||
run_tests()
|
||||
else:
|
||||
print(f"Generating music with the following arguments: {kwargs}")
|
||||
# Generate music
|
||||
generated_files = generate_music(**kwargs)
|
||||
print(f"Generated music files: {generated_files}")
|
||||
|
||||
|
||||
def run_tests():
|
||||
print("Running test suite...")
|
||||
# Define your test cases here
|
||||
test_model_name_validation()
|
||||
test_file_generation()
|
||||
# Add more tests as needed
|
||||
print("All tests passed!")
|
||||
|
||||
|
||||
def test_model_name_validation():
|
||||
print("Testing model name validation...")
|
||||
model_names = [
|
||||
"facebook/musicgen-melody-large",
|
||||
"facebook/musicgen-melody-small",
|
||||
"facebook/musicgen-melody-medium",
|
||||
"facebook/musicgen-stereo-melody-large",
|
||||
"facebook/musicgen-stereo-melody-small",
|
||||
"facebook/musicgen-stereo-melody-medium",
|
||||
"facebook/musicgen-large",
|
||||
"facebook/musicgen-small",
|
||||
"facebook/musicgen-medium",
|
||||
"facebook/musicgen-stereo-large",
|
||||
"facebook/musicgen-stereo-small",
|
||||
"facebook/musicgen-stereo-medium",
|
||||
]
|
||||
|
||||
description = "A brief test piece" # Example description for testing
|
||||
duration = 5 # 5 seconds long
|
||||
file_template = "test_output_{}.wav" # Template for generated file names
|
||||
|
||||
for i, model_name in enumerate(model_names):
|
||||
print(f"Testing {model_name}...")
|
||||
kwargs = {
|
||||
'description': description,
|
||||
'duration': duration,
|
||||
'model_name': model_name,
|
||||
'file': file_template.format(i), # Ensure unique file names
|
||||
'seed': 42 # Fixed seed for consistency, if needed
|
||||
}
|
||||
|
||||
try:
|
||||
# Assuming generate_music is adapted to work with these arguments
|
||||
generate_music(**kwargs)
|
||||
print(f"Successfully tested {model_name}")
|
||||
except Exception as e:
|
||||
print(f"Error testing {model_name}: {e}")
|
||||
# Handle the error appropriately - for example, you might want to fail the test or log the error.
|
||||
|
||||
print("Model name validation test completed.")
|
||||
|
||||
|
||||
def test_file_generation():
|
||||
# Example test function
|
||||
print("Testing file generation...")
|
||||
# Implement the logic to test music file generation
|
||||
# This could involve generating a small piece of music and checking if the file exists
|
||||
# ...
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -1,3 +1,4 @@
|
||||
# TODO: replace pyfastnoisesimd with pyfastnoiselite (works with rpi)
|
||||
import pyfastnoisesimd as fns
|
||||
|
||||
|
||||
|
||||
@@ -0,0 +1,94 @@
|
||||
import argparse
|
||||
import torch.multiprocessing as mp
|
||||
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
def generate_xtts_for_text(text, file, lang, speaker):
|
||||
from TTS.api import TTS
|
||||
import torch
|
||||
tts = TTS("tts_models/multilingual/multi-dataset/xtts_v2").to("cuda")
|
||||
tts.tts_to_file(
|
||||
text=text,
|
||||
speaker_wav=speaker,
|
||||
language=lang,
|
||||
file_path=file,
|
||||
# emotion="angry", # doesnt do anything as far as i can tell with xtts2
|
||||
)
|
||||
|
||||
|
||||
def generate_xtts(**kwargs):
|
||||
import torch
|
||||
import ast
|
||||
import uuid
|
||||
from concurrent.futures import ProcessPoolExecutor, as_completed
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
text = kwargs.get('text')
|
||||
folder = kwargs.get('folder', '.')
|
||||
lang = kwargs.get('lang', 'en')
|
||||
speaker = kwargs.get('speaker', '')
|
||||
|
||||
try:
|
||||
texts = ast.literal_eval(text)
|
||||
except SyntaxError:
|
||||
texts = [text]
|
||||
|
||||
o_files = []
|
||||
batch_size = 3
|
||||
with ProcessPoolExecutor(max_workers=batch_size) as executor:
|
||||
futures = []
|
||||
for i, text in enumerate(texts):
|
||||
ii = str(i).zfill(1 + len(str(1 + len(texts))))
|
||||
file_to_use = f"{folder}/tts_{ii}_{str(uuid.uuid4())}.wav"
|
||||
future = executor.submit(generate_xtts_for_text, text, file_to_use, lang, speaker)
|
||||
futures.append((future, file_to_use))
|
||||
|
||||
for future, file_to_use in futures:
|
||||
future.result()
|
||||
o_files.append(file_to_use)
|
||||
|
||||
torch.cuda.empty_cache()
|
||||
|
||||
return (o_files[0], o_files,)
|
||||
|
||||
|
||||
|
||||
|
||||
# def generate_xtts_for_text_advanced(text, file, lang, speaker):
|
||||
# from TTS.api import TTS
|
||||
# import torch
|
||||
#
|
||||
# # Create a SharedMemoryManager
|
||||
# manager = mp.Manager()
|
||||
# # Create a shared list to store the result
|
||||
# result = manager.list()
|
||||
# model = manager.model()
|
||||
#
|
||||
#
|
||||
# # also we want to SHARE the TTS model
|
||||
# something = manager.
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
|
||||
def main():
|
||||
parser = argparse.ArgumentParser(description='Generate from given text.')
|
||||
parser.add_argument('--text', help='Text to generate for')
|
||||
parser.add_argument('--folder', type=str, default='.', help='output folder')
|
||||
parser.add_argument('--lang', type=str, default='en')
|
||||
parser.add_argument('--speaker', type=str, default='', help='Speaker wav use for generation')
|
||||
|
||||
args = parser.parse_args()
|
||||
kwargs = vars(args)
|
||||
|
||||
print(f"Generating with the following arguments: {kwargs}")
|
||||
generated_files = generate_xtts(**kwargs)
|
||||
print(f"Generated files: {generated_files}")
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user