201 lines
6.3 KiB
Python
201 lines
6.3 KiB
Python
import torch
|
|
import random
|
|
import json
|
|
import os
|
|
import numpy as np
|
|
|
|
node_dir = os.path.dirname(os.path.abspath(__file__))
|
|
|
|
def decode_audio(
|
|
latents: torch.Tensor,
|
|
vae_model: torch.nn.Module,
|
|
chunked: bool = False,
|
|
overlap: int = 32,
|
|
chunk_size: int = 128
|
|
):
|
|
downsampling_ratio = 2048
|
|
io_channels = 2
|
|
if not chunked:
|
|
try:
|
|
output = vae_model.decode_export(latents)
|
|
return output
|
|
except Exception as e:
|
|
raise
|
|
else:
|
|
# Chunked decoding logic
|
|
hop_size = chunk_size - overlap
|
|
total_size = latents.shape[2]
|
|
batch_size = latents.shape[0]
|
|
chunks = []
|
|
i = 0
|
|
for i in range(0, total_size - chunk_size + 1, hop_size):
|
|
chunk = latents[:, :, i : i + chunk_size]
|
|
chunks.append(chunk)
|
|
if i + chunk_size != total_size:
|
|
# Final chunk
|
|
chunk = latents[:, :, -chunk_size:]
|
|
chunks.append(chunk)
|
|
chunks = torch.stack(chunks)
|
|
num_chunks = chunks.shape[0]
|
|
# samples_per_latent is just the downsampling ratio
|
|
samples_per_latent = downsampling_ratio
|
|
# Create an empty waveform, we will populate it with chunks as decode them
|
|
y_size = total_size * samples_per_latent
|
|
y_final = torch.zeros((batch_size, io_channels, y_size)).to(latents.device)
|
|
for i in range(num_chunks):
|
|
x_chunk = chunks[i, :]
|
|
try:
|
|
y_chunk = vae_model.decode_export(x_chunk)
|
|
except Exception as e:
|
|
raise
|
|
# figure out where to put the audio along the time domain
|
|
if i == num_chunks - 1:
|
|
# final chunk always goes at the end
|
|
t_end = y_size
|
|
t_start = t_end - y_chunk.shape[2]
|
|
else:
|
|
t_start = i * hop_size * samples_per_latent
|
|
t_end = t_start + chunk_size * samples_per_latent
|
|
# remove the edges of the overlaps
|
|
ol = (overlap // 2) * samples_per_latent
|
|
chunk_start = 0
|
|
chunk_end = y_chunk.shape[2]
|
|
if i > 0:
|
|
# no overlap for the start of the first chunk
|
|
t_start += ol
|
|
chunk_start += ol
|
|
if i < num_chunks - 1:
|
|
# no overlap for the end of the last chunk
|
|
t_end -= ol
|
|
chunk_end -= ol
|
|
# paste the chunked audio into our y_final output audio
|
|
y_final[:, :, t_start:t_end] = y_chunk[:, :, chunk_start:chunk_end]
|
|
return y_final
|
|
|
|
def get_reference_latent(device: torch.device, max_frames: int):
|
|
return torch.zeros(1, max_frames, 64).to(device)
|
|
|
|
def get_negative_style_prompt(device: torch.device):
|
|
file_path = f"{node_dir}/vocal.npy"
|
|
try:
|
|
vocal_style = np.load(file_path)
|
|
except Exception as e:
|
|
raise
|
|
|
|
vocal_style = torch.from_numpy(vocal_style).to(device) # [1, 512]
|
|
return vocal_style.half()
|
|
|
|
def parse_lyrics(lyrics: str):
|
|
lyrics_with_time = []
|
|
lyrics = lyrics.strip()
|
|
for line in lyrics.split("\n"):
|
|
try:
|
|
time, lyric = line[1:9], line[10:]
|
|
mins, secs = time.split(":")
|
|
secs = int(mins) * 60 + float(secs)
|
|
lyrics_with_time.append((secs, lyric.strip()))
|
|
except ValueError:
|
|
continue
|
|
return lyrics_with_time
|
|
|
|
class CNENTokenizer:
|
|
def __init__(self):
|
|
vocab_path = f"{node_dir}/g2p/g2p/vocab.json"
|
|
try:
|
|
with open(vocab_path, "r", encoding="utf-8") as file:
|
|
self.phone2id: dict = json.load(file)["vocab"]
|
|
except Exception as e:
|
|
raise
|
|
|
|
self.id2phone = {v: k for k, v in self.phone2id.items()}
|
|
|
|
try:
|
|
from g2p.g2p_generation import chn_eng_g2p
|
|
self.tokenizer = chn_eng_g2p
|
|
except Exception as e:
|
|
raise
|
|
|
|
def encode(self, text: str):
|
|
try:
|
|
phone, token = self.tokenizer(text)
|
|
return [x + 1 for x in token]
|
|
except Exception as e:
|
|
print(f"Text encoding failed: {str(e)}")
|
|
raise
|
|
|
|
def decode(self, token: list):
|
|
try:
|
|
return "|".join([self.id2phone[x - 1] for x in token])
|
|
except Exception as e:
|
|
raise
|
|
|
|
def get_lrc_token(
|
|
max_frames: int,
|
|
text: str,
|
|
tokenizer: CNENTokenizer,
|
|
device: torch.device
|
|
):
|
|
# Audio processing parameters
|
|
lyrics_shift = 0
|
|
sampling_rate = 44100
|
|
downsample_rate = 2048
|
|
max_secs = max_frames / (sampling_rate / downsample_rate)
|
|
|
|
# Token configuration
|
|
comma_token_id = 1
|
|
period_token_id = 2
|
|
|
|
lrc_with_time = parse_lyrics(text)
|
|
|
|
modified_lrc_with_time = []
|
|
for i in range(len(lrc_with_time)):
|
|
time, line = lrc_with_time[i]
|
|
try:
|
|
line_token = tokenizer.encode(line)
|
|
modified_lrc_with_time.append((time, line_token))
|
|
except Exception as e:
|
|
raise
|
|
lrc_with_time = modified_lrc_with_time
|
|
|
|
lrc_with_time = [
|
|
(time_start, line)
|
|
for (time_start, line) in lrc_with_time
|
|
if time_start < max_secs
|
|
]
|
|
# lrc_with_time = lrc_with_time[:-1] if len(lrc_with_time) >= 1 else lrc_with_time
|
|
|
|
normalized_start_time = 0.0
|
|
|
|
lrc = torch.zeros((max_frames,), dtype=torch.long)
|
|
|
|
tokens_count = 0
|
|
last_end_pos = 0
|
|
for time_start, line in lrc_with_time:
|
|
tokens = [
|
|
token if token != period_token_id else comma_token_id for token in line
|
|
] + [period_token_id]
|
|
tokens = torch.tensor(tokens, dtype=torch.long)
|
|
num_tokens = tokens.shape[0]
|
|
|
|
gt_frame_start = int(time_start * sampling_rate / downsample_rate)
|
|
|
|
frame_shift = random.randint(int(lyrics_shift), int(lyrics_shift))
|
|
|
|
frame_start = max(gt_frame_start - frame_shift, last_end_pos)
|
|
frame_len = min(num_tokens, max_frames - frame_start)
|
|
|
|
lrc[frame_start : frame_start + frame_len] = tokens[:frame_len]
|
|
|
|
tokens_count += num_tokens
|
|
last_end_pos = frame_start + frame_len
|
|
|
|
lrc_emb = lrc.unsqueeze(0).to(device)
|
|
|
|
normalized_start_time = torch.tensor(normalized_start_time).unsqueeze(0).to(device)
|
|
if device == "cuda":
|
|
normalized_start_time = normalized_start_time.half()
|
|
else:
|
|
normalized_start_time = normalized_start_time.float()
|
|
|
|
return lrc_emb, normalized_start_time
|