Mirrors/bark-with-voice-clone
1
0
Fork 0
mirror of https://github.com/serp-ai/bark-with-voice-clone.git synced 2025-12-15 03:07:58 +01:00
Code Issues Packages Projects Releases Wiki Activity

DRY: refactored mps supporting dtype.

Browse Source
This commit is contained in:
Raf Gemmail
2023-05-04 12:57:48 +12:00
parent 89b362b4a6
commit c95350f13f
1 changed files with 13 additions and 10 deletions
Download Patch File Download Diff File Expand all files Collapse all files

23
bark/generation.py
View File

@@ -335,11 +335,20 @@ def preload_models(
_ = load_codec_model(use_gpu=codec_use_gpu, force_reload=force_reload)
####
# Handle MPS immaturity in Pytorch
####
def _logits_to_device_float(logits):
if GLOBAL_ENABLE_MPS:
logits = logits.clone().detach().to("mps").to(torch.float)
else:
logits = logits.to(logits_device).type(logits_dtype)
return logits
####
# Generation Functionality
####
def _tokenize(tokenizer, text):
return tokenizer.encode(text, add_special_tokens=False)
@@ -351,7 +360,6 @@ def _detokenize(tokenizer, enc_text):
def _normalize_whitespace(text):
return re.sub(r"\s+", " ", text).strip()
TEXT_ENCODING_OFFSET = 10_048
SEMANTIC_PAD_TOKEN = 10_000
TEXT_PAD_TOKEN = 129_595
@@ -464,10 +472,8 @@ def generate_text_semantic(
sorted_indices_to_remove[0] = False
relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
relevant_logits = torch.from_numpy(relevant_logits)
if GLOBAL_ENABLE_MPS:
relevant_logits = torch.tensor(relevant_logits, device="mps").to(torch.float)
else:
relevant_logits = relevant_logits.to(logits_device).type(logits_dtype)
relevant_logits = _logits_to_device_float(relevant_logits)
if top_k is not None:
v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))
relevant_logits[relevant_logits < v[-1]] = -float("Inf")
@@ -663,10 +669,7 @@ def generate_coarse(
sorted_indices_to_remove[0] = False
relevant_logits[sorted_indices[sorted_indices_to_remove]] = -np.inf
relevant_logits = torch.from_numpy(relevant_logits)
if GLOBAL_ENABLE_MPS:
relevant_logits = torch.tensor(relevant_logits, device="mps").to(torch.float)
else:
relevant_logits = relevant_logits.to(logits_device).type(logits_dtype)
relevant_logits = _logits_to_device_float(relevant_logits)
if top_k is not None:
v, _ = torch.topk(relevant_logits, min(top_k, relevant_logits.size(-1)))