Spaces:
Running
Running
File size: 3,293 Bytes
506a2b4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 |
from dataclasses import dataclass
from typing import Optional
import torch
from torch import nn, Tensor
from .perceiver import Perceiver
from .t3_config import T3Config
@dataclass
class T3Cond:
"""
Dataclass container for most / all conditioning info.
TODO: serialization methods aren't used, keeping them around for convenience
"""
speaker_emb: Tensor
clap_emb: Optional[Tensor] = None
cond_prompt_speech_tokens: Optional[Tensor] = None
cond_prompt_speech_emb: Optional[Tensor] = None
emotion_adv: Optional[Tensor] = 0.5
def to(self, *, device=None, dtype=None):
"Cast to a device and dtype. Dtype casting is ignored for long/int tensors."
for k, v in self.__dict__.items():
if torch.is_tensor(v):
is_fp = type(v.view(-1)[0].item()) is not int
setattr(self, k, v.to(device=device, dtype=dtype if is_fp else None))
return self
def save(self, fpath):
torch.save(self.__dict__, fpath)
@staticmethod
def load(fpath, map_location="cpu"):
kwargs = torch.load(fpath, map_location=map_location, weights_only=True)
return T3Cond(**kwargs)
class T3CondEnc(nn.Module):
"""
Handle all non-text conditioning, like speaker embeddings / prompts, CLAP, emotion, etc.
"""
def __init__(self, hp: T3Config):
super().__init__()
self.hp = hp
if hp.encoder_type == "voice_encoder":
self.spkr_enc = nn.Linear(hp.speaker_embed_size, hp.n_channels)
else:
raise NotImplementedError(str(hp.encoder_type))
# emotion adv
self.emotion_adv_fc = None
if hp.emotion_adv:
self.emotion_adv_fc = nn.Linear(1, hp.n_channels, bias=False)
# perceiver resampler
self.perceiver = None
if hp.use_perceiver_resampler:
self.perceiver = Perceiver()
def forward(self, cond: T3Cond):
# Validate
assert (cond.cond_prompt_speech_tokens is None) == (cond.cond_prompt_speech_emb is None), \
"no embeddings for cond_prompt_speech_tokens"
# Speaker embedding projection
cond_spkr = self.spkr_enc(cond.speaker_emb.view(-1, self.hp.speaker_embed_size))[:, None] # (B, 1, dim)
empty = torch.zeros_like(cond_spkr[:, :0]) # (B, 0, dim)
# TODO CLAP
assert cond.clap_emb is None, "clap_embed not implemented"
cond_clap = empty # (B, 0, dim)
# Cond prompt
cond_prompt_speech_emb = cond.cond_prompt_speech_emb
if cond_prompt_speech_emb is None:
cond_prompt_speech_emb = empty # (B, 0, dim)
elif self.hp.use_perceiver_resampler:
cond_prompt_speech_emb = self.perceiver(cond_prompt_speech_emb)
# Emotion Adv: must provide a value if this model uses emotion conditioning
cond_emotion_adv = empty # (B, 0, dim)
if self.hp.emotion_adv:
assert cond.emotion_adv is not None
cond_emotion_adv = self.emotion_adv_fc(cond.emotion_adv.view(-1, 1, 1))
# Concat and return
cond_embeds = torch.cat((
cond_spkr,
cond_clap,
cond_prompt_speech_emb,
cond_emotion_adv,
), dim=1)
return cond_embeds
|