mirror of
https://github.com/Nighthawk42/mOrpheus.git
synced 2026-08-30 17:32:27 +00:00
93 lines
3.7 KiB
Python
93 lines
3.7 KiB
Python
# modules/snac_decoder.py
|
|
|
|
import time
|
|
import torch
|
|
import numpy as np
|
|
from modules.logging import logger
|
|
from snac import SNAC # Ensure that the snac module is installed
|
|
|
|
# Load SNAC model
|
|
snac_model = SNAC.from_pretrained("hubertsiuzdak/snac_24khz").eval()
|
|
snac_device = "cuda" if torch.cuda.is_available() else "cpu"
|
|
logger.info("Using SNAC on device: %s", snac_device)
|
|
snac_model = snac_model.to(snac_device)
|
|
cuda_stream = torch.cuda.Stream() if snac_device == "cuda" else None
|
|
|
|
def convert_to_audio(multiframe, count):
|
|
if len(multiframe) < 7:
|
|
return None
|
|
num_frames = len(multiframe) // 7
|
|
frame = multiframe[:num_frames * 7]
|
|
codes_0 = torch.zeros(num_frames, dtype=torch.int32, device=snac_device)
|
|
codes_1 = torch.zeros(num_frames * 2, dtype=torch.int32, device=snac_device)
|
|
codes_2 = torch.zeros(num_frames * 4, dtype=torch.int32, device=snac_device)
|
|
frame_tensor = torch.tensor(frame, dtype=torch.int32, device=snac_device)
|
|
for j in range(num_frames):
|
|
idx = j * 7
|
|
codes_0[j] = frame_tensor[idx]
|
|
codes_1[j * 2] = frame_tensor[idx + 1]
|
|
codes_1[j * 2 + 1] = frame_tensor[idx + 4]
|
|
codes_2[j * 4] = frame_tensor[idx + 2]
|
|
codes_2[j * 4 + 1] = frame_tensor[idx + 3]
|
|
codes_2[j * 4 + 2] = frame_tensor[idx + 5]
|
|
codes_2[j * 4 + 3] = frame_tensor[idx + 6]
|
|
codes = [codes_0.unsqueeze(0), codes_1.unsqueeze(0), codes_2.unsqueeze(0)]
|
|
if (torch.any(codes[0] < 0) or torch.any(codes[0] > 4096) or
|
|
torch.any(codes[1] < 0) or torch.any(codes[1] > 4096) or
|
|
torch.any(codes[2] < 0) or torch.any(codes[2] > 4096)):
|
|
return None
|
|
stream_ctx = torch.cuda.stream(cuda_stream) if cuda_stream is not None else torch.no_grad()
|
|
with stream_ctx, torch.inference_mode():
|
|
audio_hat = snac_model.decode(codes)
|
|
audio_slice = audio_hat[:, :, 2048:4096]
|
|
if snac_device == "cuda":
|
|
audio_int16_tensor = (audio_slice * 32767).to(torch.int16)
|
|
audio_bytes = audio_int16_tensor.cpu().numpy().tobytes()
|
|
else:
|
|
audio_np = audio_slice.detach().cpu().numpy()
|
|
audio_int16 = (audio_np * 32767).astype(np.int16)
|
|
audio_bytes = audio_int16.tobytes()
|
|
return audio_bytes
|
|
|
|
def turn_token_into_id(token_string, index):
|
|
token_string = token_string.strip()
|
|
if "<custom_token_" not in token_string:
|
|
return None
|
|
last_token_start = token_string.rfind("<custom_token_")
|
|
if last_token_start == -1 or not token_string.endswith(">"):
|
|
return None
|
|
try:
|
|
number_str = token_string[last_token_start + 14:-1]
|
|
return int(number_str) - 10 - ((index % 7) * 4096)
|
|
except (ValueError, IndexError):
|
|
return None
|
|
|
|
token_cache = {}
|
|
MAX_CACHE_SIZE = 1000
|
|
|
|
def tokens_decoder(token_gen):
|
|
buffer = []
|
|
count = 0
|
|
min_frames_required = 28
|
|
process_every = 7
|
|
for token_text in token_gen:
|
|
cache_key = (token_text, count % 7)
|
|
if cache_key in token_cache:
|
|
token = token_cache[cache_key]
|
|
else:
|
|
token = turn_token_into_id(token_text, count)
|
|
if token is not None and len(token_cache) < MAX_CACHE_SIZE:
|
|
token_cache[cache_key] = token
|
|
if token is not None and token > 0:
|
|
buffer.append(token)
|
|
count += 1
|
|
if count % process_every == 0 and count >= min_frames_required:
|
|
buffer_to_proc = buffer[-min_frames_required:]
|
|
audio_samples = convert_to_audio(buffer_to_proc, count)
|
|
if audio_samples is not None:
|
|
yield audio_samples
|
|
|
|
def tokens_decoder_sync(syn_token_gen):
|
|
audio_segments = list(tokens_decoder(syn_token_gen))
|
|
return b"".join(audio_segments)
|