From b97f3c0159ef660aa6b4bb05c12afb76e892f634 Mon Sep 17 00:00:00 2001 From: Abdul Basit Rana Date: Wed, 28 Jan 2026 21:27:26 +0500 Subject: [PATCH 1/2] Enhance device compatibility by auto-detecting available hardware (CUDA/MPS) and updating model loading functions accordingly --- zipvoice/luxvoice.py | 24 +++++++++++++++------ zipvoice/modeling_utils.py | 22 +++++++++---------- zipvoice/models/modules/scaling.py | 32 ++++++++++++++++++++-------- zipvoice/models/modules/zipformer.py | 9 +++++++- 4 files changed, 59 insertions(+), 28 deletions(-) diff --git a/zipvoice/luxvoice.py b/zipvoice/luxvoice.py index 72db595..1ead207 100644 --- a/zipvoice/luxvoice.py +++ b/zipvoice/luxvoice.py @@ -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: diff --git a/zipvoice/modeling_utils.py b/zipvoice/modeling_utils.py index 893e3fe..f4e621d 100644 --- a/zipvoice/modeling_utils.py +++ b/zipvoice/modeling_utils.py @@ -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") diff --git a/zipvoice/models/modules/scaling.py b/zipvoice/models/modules/scaling.py index 580402b..f20864e 100644 --- a/zipvoice/models/modules/scaling.py +++ b/zipvoice/models/modules/scaling.py @@ -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) diff --git a/zipvoice/models/modules/zipformer.py b/zipvoice/models/modules/zipformer.py index 8c00bd1..aed6cd3 100644 --- a/zipvoice/models/modules/zipformer.py +++ b/zipvoice/models/modules/zipformer.py @@ -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) From d236deb1393c93aa4a7bdd0ddc76e09fab6575b1 Mon Sep 17 00:00:00 2001 From: Abdul Basit Rana Date: Wed, 28 Jan 2026 21:30:19 +0500 Subject: [PATCH 2/2] Refactor dependencies in pyproject.toml and ease uv tool installation process --- pyproject.toml | 29 ++++++++++++++++++++++++++++- 1 file changed, 28 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 992bb74..0ead81e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,25 @@ description = "Linacodec is a highly compressive and rapid audio tokenizer for a readme = "README.md" requires-python = ">=3.10" dependencies = [ - "huggingface_hub", + "cn2an", + "huggingface-hub", + "inflect", + "jieba", + "lhotse", + "librosa", + "linacodec", + "numpy", + "onnxruntime", + "piper-phonemize", + "pydub", + "pypinyin", + "safetensors", + "setuptools<81", + "tensorboard", + "torch", + "torchaudio", + "transformers<=4.57.6", + "vocos", ] [project.optional-dependencies] @@ -19,3 +37,12 @@ build-backend = "uv_build" [tool.ruff] line-length = 120 + +[tool.uv] +find-links = ["https://k2-fsa.github.io/icefall/piper_phonemize.html"] + +[tool.uv.sources] +linacodec = { git = "https://github.com/ysharma3501/LinaCodec.git" } + +[tool.uv.build-backend] +module-root = "" # The root module of the project