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

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)