igna/LuxTTS
1
0
Fork 0
Code Activity

Enhance device compatibility by auto-detecting available hardware (CUDA/MPS) and updating model loading functions accordingly

Browse Source
This commit is contained in:
Abdul Basit Rana
2026-01-28 21:27:26 +05:00
parent fb0c44a30c
commit b97f3c0159
4 changed files with 59 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

24
zipvoice/luxvoice.py
View File

@@ -1,20 +1,30 @@
import torch
from zipvoice.modeling_utils import process_audio, generate, load_models_gpu, load_models_cpu
from zipvoice.onnx_modeling import generate_cpu
class LuxTTS:
"""
LuxTTS class for encoding prompt and generating speech on cpu/cuda.
LuxTTS class for encoding prompt and generating speech on cpu/cuda/mps.
"""
def __init__(self, model_path='YatharthS/LuxTTS', device='cuda', threads=4):
if model_path == 'YatharthS/LuxTTS':
model_path = None
# Auto-detect better device if cuda is requested but not available
if device == 'cuda' and not torch.cuda.is_available():
if torch.backends.mps.is_available():
print("CUDA not available, switching to MPS")
device = 'mps'
else:
print("CUDA not available, switching to CPU")
device = 'cpu'
if device == 'cpu':
model, feature_extractor, vocos, tokenizer, transcriber = load_models_cpu(model_path, threads)
print("Loading model on CPU")
else:
model, feature_extractor, vocos, tokenizer, transcriber = load_models_gpu(model_path)
model, feature_extractor, vocos, tokenizer, transcriber = load_models_gpu(model_path, device=device)
print("Loading model on GPU")
self.model = model
@@ -24,7 +34,7 @@ class LuxTTS:
self.transcriber = transcriber
self.device = device
self.vocos.freq_range = 12000
def encode_prompt(self, prompt_audio, duration=5, rms=0.001):
@@ -33,17 +43,17 @@ class LuxTTS:
encode_dict = {"prompt_tokens": prompt_tokens, 'prompt_features_lens': prompt_features_lens, 'prompt_features': prompt_features, 'prompt_rms': prompt_rms}
return encode_dict
def generate_speech(self, text, encode_dict, num_steps=4, guidance_scale=3.0, t_shift=0.5, speed=1.0, return_smooth=False):
"""encodes text and generates speech using flow matching model according to steps, guidance scale, and t_shift(like temp)"""
prompt_tokens, prompt_features_lens, prompt_features, prompt_rms = encode_dict.values()
if return_smooth == True:
self.vocos.return_48k = False
else:
self.vocos.return_48k = True
if self.device == 'cpu':
final_wav = generate_cpu(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, text, self.model, self.vocos, self.tokenizer, num_step=num_steps, guidance_scale=guidance_scale, t_shift=t_shift, speed=speed)
else:

22
zipvoice/modeling_utils.py
View File

@@ -52,7 +52,7 @@ def process_audio(audio, transcriber, tokenizer, feature_extractor, device, targ
prompt_wav = torch.from_numpy(prompt_wav).unsqueeze(0)
prompt_wav, prompt_rms = rms_norm(prompt_wav, target_rms)
prompt_features = feature_extractor.extract(
prompt_wav, sampling_rate=24000
).to(device)
@@ -90,16 +90,16 @@ def generate(prompt_tokens, prompt_features_lens, prompt_features, prompt_rms, t
return wav
def load_models_gpu(model_path=None):
def load_models_gpu(model_path=None, device="cuda"):
params = LuxTTSConfig()
if model_path is None:
model_path = snapshot_download("YatharthS/LuxTTS")
token_file = f"{model_path}/tokens.txt"
model_ckpt = f"{model_path}/model.pt"
model_config = f"{model_path}/config.json"
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base", device='cuda')
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-base", device=device)
tokenizer = EmiliaTokenizer(token_file=token_file)
tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id}
@@ -111,16 +111,16 @@ def load_models_gpu(model_path=None):
**tokenizer_config,
)
load_checkpoint(filename=model_ckpt, model=model, strict=True)
params.device = torch.device("cuda", 0)
params.device = torch.device(device, 0)
model = model.to(params.device).eval()
feature_extractor = VocosFbank()
vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').cuda()
vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').to(device)
parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight")
parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight")
vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin'))
vocos.load_state_dict(torch.load(f'{model_path}/vocoder/vocos.bin', map_location=params.device))
params.sampling_rate = model_config["feature"]["sampling_rate"]
return model, feature_extractor, vocos, tokenizer, transcriber
@@ -136,15 +136,15 @@ def load_models_cpu(model_path = None, num_thread=2):
model_config = f"{model_path}/config.json"
transcriber = pipeline("automatic-speech-recognition", model="openai/whisper-tiny", device='cpu')
tokenizer = EmiliaTokenizer(token_file=token_file)
tokenizer_config = {"vocab_size": tokenizer.vocab_size, "pad_id": tokenizer.pad_id}
with open(model_config, "r") as f:
model_config = json.load(f)
model = OnnxModel(text_encoder_path, fm_decoder_path, num_thread=num_thread)
vocos = Vocos.from_hparams(f'{model_path}/vocoder/config.yaml').eval()
parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[0], "weight")
parametrize.remove_parametrizations(vocos.upsampler.upsample_layers[1], "weight")

32
zipvoice/models/modules/scaling.py
View File

@@ -33,10 +33,24 @@ import torch
import torch.nn as nn
from torch import Tensor
if torch.cuda.is_available():
DEVICE_TYPE = "cuda"
elif torch.backends.mps.is_available():
DEVICE_TYPE = "mps"
else:
DEVICE_TYPE = "cpu"
def get_memory_allocated():
if DEVICE_TYPE == "cuda":
return torch.cuda.memory_allocated()
elif DEVICE_TYPE == "mps":
return torch.mps.current_allocated_memory()
else:
return 0
def custom_amp_decorator(dec, cuda_amp_deprecated):
def decorator(func):
return dec(func) if not cuda_amp_deprecated else dec(device_type="cuda")(func)
return dec(func) if not cuda_amp_deprecated else dec(device_type=DEVICE_TYPE)(func)
return decorator
@@ -319,7 +333,7 @@ class SoftmaxFunction(torch.autograd.Function):
@staticmethod
def backward(ctx, ans_grad: Tensor):
(ans,) = ctx.saved_tensors
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
ans_grad = ans_grad.to(torch.float32)
ans = ans.to(torch.float32)
x_grad = ans_grad * ans
@@ -535,7 +549,7 @@ class BalancerFunction(torch.autograd.Function):
try:
with torch.enable_grad():
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
x = x.to(torch.float32)
x = x.detach()
x.requires_grad = True
@@ -648,7 +662,7 @@ class Balancer(torch.nn.Module):
if (
torch.jit.is_scripting()
or not x.requires_grad
or (x.is_cuda and self.mem_cutoff(torch.cuda.memory_allocated()))
or ((x.is_cuda or x.device.type == "mps") and self.mem_cutoff(get_memory_allocated()))
):
return _no_op(x)
@@ -802,7 +816,7 @@ class WhiteningPenaltyFunction(torch.autograd.Function):
try:
with torch.enable_grad():
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
x_detached = x_orig.to(torch.float32).detach()
x_detached.requires_grad = True
@@ -1046,7 +1060,7 @@ class SwooshLFunction(torch.autograd.Function):
coeff = -0.08
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
with torch.enable_grad():
x = x.detach()
x.requires_grad = True
@@ -1124,7 +1138,7 @@ class SwooshRFunction(torch.autograd.Function):
zero = torch.tensor(0.0, dtype=x.dtype, device=x.device)
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
with torch.enable_grad():
x = x.detach()
x.requires_grad = True
@@ -1187,7 +1201,7 @@ class SwooshROnnx(torch.nn.Module):
# simple version of SwooshL that does not redefine the backprop, used in
# ActivationDropoutAndLinearFunction.
def SwooshLForward(x: Tensor):
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
x = x.to(torch.float32)
x_offset = x - 4.0
log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)
@@ -1198,7 +1212,7 @@ def SwooshLForward(x: Tensor):
# simple version of SwooshR that does not redefine the backprop, used in
# ActivationDropoutAndLinearFunction.
def SwooshRForward(x: Tensor):
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
x = x.to(torch.float32)
x_offset = x - 1.0
log_sum = (1.0 + x_offset.exp()).log().to(x.dtype)

9
zipvoice/models/modules/zipformer.py
View File

@@ -27,6 +27,13 @@ from typing import Optional, Tuple, Union
import torch
from torch import Tensor, nn
if torch.cuda.is_available():
DEVICE_TYPE = "cuda"
elif torch.backends.mps.is_available():
DEVICE_TYPE = "mps"
else:
DEVICE_TYPE = "cpu"
from zipvoice.models.modules.scaling import (
ActivationDropoutAndLinear,
Balancer,
@@ -1310,7 +1317,7 @@ class RelPositionMultiheadAttentionWeights(nn.Module):
(num_heads, batch_size, seq_len, seq_len) = attn_weights.shape
with torch.no_grad():
with torch.amp.autocast("cuda", enabled=False):
with torch.amp.autocast(DEVICE_TYPE, enabled=False):
attn_weights = attn_weights.to(torch.float32)
attn_weights_entropy = (
-((attn_weights + 1.0e-20).log() * attn_weights)