Copy the files in the new repo for public release
This commit is contained in:
gabrieldemarmiesse
@@ -0,0 +1 @@
|
||||
"""Modules used for building the models."""
|
||||
@@ -0,0 +1,161 @@
|
||||
import math
|
||||
import warnings
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
|
||||
|
||||
def get_extra_padding_for_conv1d(
|
||||
x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0
|
||||
) -> int:
|
||||
"""See `pad_for_conv1d`."""
|
||||
length = x.shape[-1]
|
||||
n_frames = (length - kernel_size + padding_total) / stride + 1
|
||||
ideal_length = (math.ceil(n_frames) - 1) * stride + (kernel_size - padding_total)
|
||||
return ideal_length - length
|
||||
|
||||
|
||||
def pad_for_conv1d(x: torch.Tensor, kernel_size: int, stride: int, padding_total: int = 0):
|
||||
"""Pad for a convolution to make sure that the last window is full.
|
||||
Extra padding is added at the end. This is required to ensure that we can rebuild
|
||||
an output of the same length, as otherwise, even with padding, some time steps
|
||||
might get removed.
|
||||
For instance, with total padding = 4, kernel size = 4, stride = 2:
|
||||
0 0 1 2 3 4 5 0 0 # (0s are padding)
|
||||
1 2 3 # (output frames of a convolution, last 0 is never used)
|
||||
0 0 1 2 3 4 5 0 # (output of tr. conv., but pos. 5 is going to get removed as padding)
|
||||
1 2 3 4 # once you removed padding, we are missing one time step !
|
||||
"""
|
||||
extra_padding = get_extra_padding_for_conv1d(x, kernel_size, stride, padding_total)
|
||||
return F.pad(x, (0, extra_padding))
|
||||
|
||||
|
||||
class StreamingConv1d(StatefulModule):
|
||||
"""Conv1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
dilation: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
pad_mode: str = "constant",
|
||||
):
|
||||
super().__init__()
|
||||
assert pad_mode in ["constant", "replicate"], pad_mode
|
||||
self.pad_mode = pad_mode
|
||||
# warn user on unusual setup between dilation and stride
|
||||
if stride > 1 and dilation > 1:
|
||||
warnings.warn(
|
||||
"StreamingConv1d has been initialized with stride > 1 and dilation > 1"
|
||||
f" (kernel_size={kernel_size} stride={stride}, dilation={dilation})."
|
||||
)
|
||||
self.conv = nn.Conv1d(
|
||||
in_channels,
|
||||
out_channels,
|
||||
kernel_size,
|
||||
stride,
|
||||
dilation=dilation,
|
||||
groups=groups,
|
||||
bias=bias,
|
||||
)
|
||||
|
||||
@property
|
||||
def _stride(self) -> int:
|
||||
return self.conv.stride[0]
|
||||
|
||||
@property
|
||||
def _kernel_size(self) -> int:
|
||||
return self.conv.kernel_size[0]
|
||||
|
||||
@property
|
||||
def _effective_kernel_size(self) -> int:
|
||||
dilation = self.conv.dilation[0]
|
||||
return (self._kernel_size - 1) * dilation + 1 # effective kernel size with dilations
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
stride = self._stride
|
||||
# Effective kernel size accounting for dilation.
|
||||
kernel = self._effective_kernel_size
|
||||
previous = torch.zeros(batch_size, self.conv.in_channels, kernel - stride)
|
||||
first = torch.ones(batch_size, dtype=torch.bool)
|
||||
return dict(previous=previous, first=first)
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
B, C, T = x.shape
|
||||
S = self._stride
|
||||
assert T > 0 and T % S == 0, "Steps must be multiple of stride"
|
||||
if model_state is None:
|
||||
state = self.init_state(B, 0)
|
||||
else:
|
||||
state = self.get_state(model_state)
|
||||
TP = state["previous"].shape[-1]
|
||||
if TP and self.pad_mode == "replicate":
|
||||
assert T >= TP, "Not enough content to pad streaming."
|
||||
init = x[..., :1]
|
||||
state["previous"][:] = torch.where(
|
||||
state["first"].view(-1, 1, 1), init, state["previous"]
|
||||
)
|
||||
if TP:
|
||||
x = torch.cat([state["previous"], x], dim=-1)
|
||||
y = self.conv(x)
|
||||
if TP:
|
||||
state["previous"][:] = x[..., -TP:]
|
||||
if self.pad_mode == "replicate":
|
||||
state["first"] = torch.zeros_like(state["first"])
|
||||
return y
|
||||
|
||||
|
||||
class StreamingConvTranspose1d(StatefulModule):
|
||||
"""ConvTranspose1d with some builtin handling of asymmetric or causal padding
|
||||
and normalization.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels: int,
|
||||
out_channels: int,
|
||||
kernel_size: int,
|
||||
stride: int = 1,
|
||||
groups: int = 1,
|
||||
bias: bool = True,
|
||||
):
|
||||
super().__init__()
|
||||
self.convtr = nn.ConvTranspose1d(
|
||||
in_channels, out_channels, kernel_size, stride, groups=groups, bias=bias
|
||||
)
|
||||
|
||||
@property
|
||||
def _stride(self) -> int:
|
||||
return self.convtr.stride[0]
|
||||
|
||||
@property
|
||||
def _kernel_size(self) -> int:
|
||||
return self.convtr.kernel_size[0]
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
K = self._kernel_size
|
||||
S = self._stride
|
||||
return dict(partial=torch.zeros(batch_size, self.convtr.out_channels, K - S))
|
||||
|
||||
def forward(self, x, mimi_state: dict):
|
||||
layer_state = self.get_state(mimi_state)["partial"]
|
||||
y = self.convtr(x)
|
||||
PT = layer_state.shape[-1]
|
||||
if PT > 0:
|
||||
y[..., :PT] += layer_state
|
||||
bias = self.convtr.bias
|
||||
for_partial = y[..., -PT:]
|
||||
if bias is not None:
|
||||
for_partial -= bias[:, None]
|
||||
layer_state[:] = for_partial
|
||||
y = y[..., :-PT]
|
||||
return y
|
||||
@@ -0,0 +1,18 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class DummyQuantizer(nn.Module):
|
||||
"""Simplified quantizer that only provides output projection for TTS.
|
||||
|
||||
This removes all unnecessary quantization logic since we don't use actual quantization.
|
||||
"""
|
||||
|
||||
def __init__(self, dimension: int, output_dimension: int):
|
||||
super().__init__()
|
||||
self.dimension = dimension
|
||||
self.output_dimension = output_dimension
|
||||
self.output_proj = torch.nn.Conv1d(self.dimension, self.output_dimension, 1, bias=False)
|
||||
|
||||
def forward(self, x: torch.Tensor) -> torch.Tensor:
|
||||
return self.output_proj(x)
|
||||
@@ -0,0 +1,11 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class LayerScale(nn.Module):
|
||||
def __init__(self, channels: int, init: float):
|
||||
super().__init__()
|
||||
self.scale = nn.Parameter(torch.full((channels,), init))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return self.scale * x
|
||||
@@ -0,0 +1,285 @@
|
||||
from typing import NamedTuple
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from einops import rearrange
|
||||
from torch.nn import functional as F
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.modules.layer_scale import LayerScale
|
||||
from pocket_tts.modules.rope import RotaryEmbedding
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
from pocket_tts.modules.transformer import StreamingMultiheadAttention
|
||||
from pocket_tts.utils.config import FlowLMTransformerConfig
|
||||
|
||||
|
||||
class KVCacheResult(NamedTuple):
|
||||
keys: torch.Tensor
|
||||
values: torch.Tensor
|
||||
positions: torch.Tensor
|
||||
|
||||
@staticmethod
|
||||
def from_kv(keys: torch.Tensor, values: torch.Tensor) -> "KVCacheResult":
|
||||
B, H, T, D = keys.shape
|
||||
assert tuple(values.shape[:-1]) == (B, H, T)
|
||||
positions = torch.arange(T, device=keys.device, dtype=torch.long)
|
||||
return KVCacheResult(keys, values, positions.expand(B, -1))
|
||||
|
||||
|
||||
def complete(
|
||||
cache: torch.Tensor, end_offset: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
||||
) -> KVCacheResult:
|
||||
capacity = cache.shape[3]
|
||||
assert k.shape[:-1] == v.shape[:-1], (k.shape, v.shape)
|
||||
B, H, T, D = k.shape
|
||||
assert T > 0
|
||||
indexes = torch.arange(T, device=end_offset.device, dtype=end_offset.dtype)
|
||||
indexes = indexes + end_offset.view(-1, 1)
|
||||
indexes = indexes % capacity
|
||||
# indexes is [B, T]
|
||||
# k is [B, H, T, D]
|
||||
# cache is [B, H, T', D]
|
||||
this_indexes = indexes.view(B, 1, T, 1)
|
||||
this_indexes = this_indexes.expand(-1, H, T, D)
|
||||
cache[0].scatter_(2, this_indexes, k)
|
||||
cache[1].scatter_(2, this_indexes, v)
|
||||
|
||||
keys = cache[0]
|
||||
values = cache[1]
|
||||
|
||||
indexes = torch.arange(capacity, device=end_offset.device, dtype=torch.long)
|
||||
|
||||
# end_index correspond to the actual index where the last value was written.
|
||||
last_offset = end_offset.view(-1, 1) + T - 1
|
||||
end_index = last_offset % capacity
|
||||
delta = indexes - end_index
|
||||
|
||||
positions = torch.where(delta <= 0, last_offset + delta, last_offset + delta - capacity)
|
||||
end_offset[:] = end_offset + T
|
||||
invalid = indexes >= end_offset.view(-1, 1)
|
||||
positions = torch.where(invalid, torch.full_like(positions, -1), positions)
|
||||
|
||||
return KVCacheResult(keys, values, positions)
|
||||
|
||||
|
||||
class MimiStreamingMultiheadAttention(StatefulModule):
|
||||
def __init__(self, embed_dim: int, num_heads: int, context: int, rope: RotaryEmbedding):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.context = context
|
||||
self.rope = rope
|
||||
self.num_heads = num_heads
|
||||
out_dim = 3 * embed_dim
|
||||
|
||||
self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
|
||||
self.in_proj = nn.Linear(embed_dim, out_dim, bias=False)
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
dim_per_head = self.embed_dim // self.num_heads
|
||||
|
||||
state = {}
|
||||
state["offset"] = torch.zeros(batch_size, dtype=torch.long)
|
||||
state["cache"] = torch.zeros((2, batch_size, self.num_heads, sequence_length, dim_per_head))
|
||||
state["end_offset"] = torch.zeros(batch_size, dtype=torch.long)
|
||||
return state
|
||||
|
||||
def increment_step(self, state, increment: int = 1):
|
||||
state["offset"] += increment
|
||||
|
||||
def _complete_kv(self, k, v, model_state: dict | None) -> KVCacheResult:
|
||||
if model_state is None:
|
||||
return KVCacheResult.from_kv(k, v)
|
||||
else:
|
||||
layer_state = self.get_state(model_state)
|
||||
return complete(layer_state["cache"], layer_state["end_offset"], k, v)
|
||||
|
||||
def forward(self, query: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
B, T = query.shape[:2]
|
||||
|
||||
if model_state is None:
|
||||
offset = torch.zeros(B, device=query.device, dtype=torch.long)
|
||||
else:
|
||||
offset = self.get_state(model_state)["offset"]
|
||||
|
||||
projected = self.in_proj(query)
|
||||
|
||||
q, k, v = rearrange(projected, "b t (p h d) -> p b h t d", p=3, h=self.num_heads)
|
||||
|
||||
# Permute from [b, h, t, d] to [b, t, h, d] for rope
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
q, k = self.rope(q, k, offset)
|
||||
# Permute back from [b, t, h, d] to [b, h, t, d]
|
||||
q = q.permute(0, 2, 1, 3)
|
||||
k = k.permute(0, 2, 1, 3)
|
||||
|
||||
k, v, pos_k = self._complete_kv(k, v, model_state)
|
||||
pos_k = pos_k[:, None]
|
||||
pos_q = offset.view(-1, 1, 1) + torch.arange(T, device=q.device, dtype=torch.long).view(
|
||||
-1, 1
|
||||
)
|
||||
delta = pos_q - pos_k
|
||||
attn_bias = (pos_k >= 0) & (delta >= 0)
|
||||
attn_bias = attn_bias & (delta < self.context)
|
||||
attn_bias = attn_bias[:, None]
|
||||
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_bias, dropout_p=0.0)
|
||||
|
||||
x = rearrange(x, "b h t d -> b t (h d)")
|
||||
x = self.out_proj(x)
|
||||
return x
|
||||
|
||||
|
||||
class StreamingTransformerLayer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
dim_feedforward: int,
|
||||
context: int | None,
|
||||
rope: RotaryEmbedding,
|
||||
layer_scale: float | None = None,
|
||||
attention_kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
# Redefine self_attn to our streaming multi-head attention
|
||||
if attention_kind == "mimi":
|
||||
# TODO: we should actually use StreamingMultiheadAttention here and add context window
|
||||
# support. And we should then delete MimiStreamingMultiheadAttention.
|
||||
# The implementation is really close.
|
||||
self.self_attn = MimiStreamingMultiheadAttention(
|
||||
context=context, rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
else:
|
||||
self.self_attn = StreamingMultiheadAttention(
|
||||
rope=rope, embed_dim=d_model, num_heads=num_heads
|
||||
)
|
||||
self.norm1 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
self.norm2 = nn.LayerNorm(d_model, eps=1e-5)
|
||||
|
||||
self.linear1 = nn.Linear(d_model, dim_feedforward, bias=False)
|
||||
self.linear2 = nn.Linear(dim_feedforward, d_model, bias=False)
|
||||
|
||||
if layer_scale is None:
|
||||
self.layer_scale_1 = nn.Identity()
|
||||
self.layer_scale_2 = nn.Identity()
|
||||
else:
|
||||
self.layer_scale_1 = LayerScale(d_model, layer_scale)
|
||||
self.layer_scale_2 = LayerScale(d_model, layer_scale)
|
||||
|
||||
def _ff_block(self, x: torch.Tensor) -> torch.Tensor:
|
||||
x_orig = x
|
||||
x = self.norm2(x)
|
||||
update = self.linear2(F.gelu(self.linear1(x)))
|
||||
return x_orig.to(update) + self.layer_scale_2(update)
|
||||
|
||||
def _sa_block(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
x_orig = x
|
||||
x = self.norm1(x)
|
||||
update = self.self_attn(x, model_state)
|
||||
return x_orig.to(update) + self.layer_scale_1(update)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None) -> torch.Tensor:
|
||||
x = self._sa_block(x, model_state)
|
||||
x = self._ff_block(x)
|
||||
return x
|
||||
|
||||
|
||||
class StreamingTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
num_layers: int,
|
||||
layer_scale: float | None = None,
|
||||
dim_feedforward: int | list[int] = 2048,
|
||||
context: int | None = None,
|
||||
max_period: float = 10_000.0,
|
||||
kind: str = "mimi",
|
||||
):
|
||||
super().__init__()
|
||||
assert d_model % num_heads == 0
|
||||
self.max_period = max_period
|
||||
|
||||
self.rope = RotaryEmbedding(max_period=max_period)
|
||||
|
||||
self.layers = nn.ModuleList()
|
||||
for _ in range(num_layers):
|
||||
self.layers.append(
|
||||
StreamingTransformerLayer(
|
||||
d_model=d_model,
|
||||
num_heads=num_heads,
|
||||
dim_feedforward=dim_feedforward,
|
||||
context=context,
|
||||
rope=self.rope,
|
||||
layer_scale=layer_scale,
|
||||
attention_kind=kind,
|
||||
)
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def from_pydantic_config(cls, config: FlowLMTransformerConfig) -> Self:
|
||||
dim_feedforward = int(config.d_model * config.hidden_scale)
|
||||
return cls(
|
||||
d_model=config.d_model,
|
||||
num_heads=config.num_heads,
|
||||
num_layers=config.num_layers,
|
||||
dim_feedforward=dim_feedforward,
|
||||
max_period=float(config.max_period),
|
||||
kind="flow_lm",
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
for layer in self.layers:
|
||||
x = layer(x, model_state)
|
||||
return x
|
||||
|
||||
|
||||
class ProjectedTransformer(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
input_dimension: int,
|
||||
output_dimensions: tuple[int, ...],
|
||||
d_model: int,
|
||||
num_heads: int,
|
||||
num_layers: int,
|
||||
layer_scale: float,
|
||||
context: int,
|
||||
max_period: float,
|
||||
dim_feedforward: int,
|
||||
):
|
||||
super().__init__()
|
||||
self.transformer = StreamingTransformer(
|
||||
d_model=d_model,
|
||||
num_heads=num_heads,
|
||||
num_layers=num_layers,
|
||||
layer_scale=layer_scale,
|
||||
context=context,
|
||||
max_period=max_period,
|
||||
dim_feedforward=dim_feedforward,
|
||||
)
|
||||
self.input_dimension = input_dimension
|
||||
self.output_dimensions = output_dimensions
|
||||
self.input_proj = None
|
||||
if d_model != input_dimension:
|
||||
self.input_proj = nn.Linear(input_dimension, d_model, bias=False)
|
||||
|
||||
self.output_projs = nn.ModuleList()
|
||||
for output_dimension in output_dimensions:
|
||||
if d_model == output_dimension:
|
||||
self.output_projs.append(nn.Identity())
|
||||
else:
|
||||
self.output_projs.append(nn.Linear(d_model, output_dimension, bias=False))
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
x = x.transpose(1, 2)
|
||||
if self.input_proj is not None:
|
||||
x = self.input_proj(x)
|
||||
z = self.transformer(x, model_state)
|
||||
ys = []
|
||||
for output_proj in self.output_projs:
|
||||
y = output_proj(z)
|
||||
y = y.transpose(1, 2)
|
||||
ys.append(y)
|
||||
return ys
|
||||
@@ -0,0 +1,215 @@
|
||||
"""
|
||||
Taken from
|
||||
https://github.com/LTH14/mar/blob/fe470ac24afbee924668d8c5c83e9fec60af3a73/models/diffloss.py
|
||||
|
||||
"""
|
||||
|
||||
import math
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from typing_extensions import Self
|
||||
|
||||
from pocket_tts.utils.config import FlowLMConfig
|
||||
|
||||
|
||||
def modulate(x, shift, scale):
|
||||
return x * (1 + scale) + shift
|
||||
|
||||
|
||||
def _rms_norm(x: torch.Tensor, alpha: torch.Tensor, eps: float):
|
||||
assert x.dim() >= alpha.dim()
|
||||
x_dtype = x.dtype
|
||||
var = eps + x.var(dim=-1, keepdim=True)
|
||||
y = (x * (alpha.to(var) * torch.rsqrt(var))).to(x_dtype)
|
||||
return y
|
||||
|
||||
|
||||
class RMSNorm(nn.Module):
|
||||
def __init__(self, dim: int, eps: float = 1e-5):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
alpha_shape = (dim,)
|
||||
self.alpha = nn.Parameter(torch.full(alpha_shape, 1.0, requires_grad=True))
|
||||
|
||||
def forward(self, x: torch.Tensor):
|
||||
return _rms_norm(x, self.alpha, self.eps)
|
||||
|
||||
|
||||
class LayerNorm(nn.Module):
|
||||
"""Reimplementation of LayerNorm because the default one doesn't support jvp."""
|
||||
|
||||
def __init__(self, channels, eps=1e-6, elementwise_affine=True):
|
||||
super().__init__()
|
||||
self.eps = eps
|
||||
if elementwise_affine:
|
||||
self.weight = nn.Parameter(torch.ones(channels))
|
||||
self.bias = nn.Parameter(torch.zeros(channels))
|
||||
|
||||
def forward(self, x):
|
||||
mean = x.mean(dim=-1, keepdim=True)
|
||||
var = x.var(dim=-1, unbiased=False, keepdim=True)
|
||||
x = (x - mean) / torch.sqrt(var + self.eps)
|
||||
if hasattr(self, "weight"):
|
||||
x = x * self.weight + self.bias
|
||||
return x
|
||||
|
||||
|
||||
class TimestepEmbedder(nn.Module):
|
||||
"""Embeds scalar timesteps into vector representations."""
|
||||
|
||||
def __init__(
|
||||
self, hidden_size: int, frequency_embedding_size: int = 256, max_period: int = 10000
|
||||
):
|
||||
super().__init__()
|
||||
blocks = [
|
||||
nn.Linear(frequency_embedding_size, hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(hidden_size, hidden_size, bias=True),
|
||||
]
|
||||
blocks.append(RMSNorm(hidden_size))
|
||||
self.mlp = nn.Sequential(*blocks)
|
||||
self.frequency_embedding_size = frequency_embedding_size
|
||||
half = frequency_embedding_size // 2
|
||||
self.register_buffer(
|
||||
"freqs", torch.exp(-math.log(max_period) * torch.arange(start=0, end=half) / half)
|
||||
)
|
||||
|
||||
def forward(self, t):
|
||||
args = t * self.freqs.to(t.dtype)
|
||||
embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
|
||||
assert not (self.frequency_embedding_size % 2)
|
||||
t_emb = self.mlp(embedding)
|
||||
return t_emb
|
||||
|
||||
|
||||
class ResBlock(nn.Module):
|
||||
"""
|
||||
A residual block that can optionally change the number of channels.
|
||||
:param channels: the number of input channels.
|
||||
"""
|
||||
|
||||
def __init__(self, channels):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
|
||||
self.in_ln = LayerNorm(channels, eps=1e-6)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(channels, channels, bias=True),
|
||||
)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(channels, 3 * channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, y):
|
||||
shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(y).chunk(3, dim=-1)
|
||||
h = modulate(self.in_ln(x), shift_mlp, scale_mlp)
|
||||
h = self.mlp(h)
|
||||
return x + gate_mlp * h
|
||||
|
||||
|
||||
class FinalLayer(nn.Module):
|
||||
"""
|
||||
The final layer adopted from DiT.
|
||||
"""
|
||||
|
||||
def __init__(self, model_channels, out_channels):
|
||||
super().__init__()
|
||||
self.norm_final = LayerNorm(model_channels, elementwise_affine=False, eps=1e-6)
|
||||
self.linear = nn.Linear(model_channels, out_channels, bias=True)
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(), nn.Linear(model_channels, 2 * model_channels, bias=True)
|
||||
)
|
||||
|
||||
def forward(self, x, c):
|
||||
shift, scale = self.adaLN_modulation(c).chunk(2, dim=-1)
|
||||
x = modulate(self.norm_final(x), shift, scale)
|
||||
x = self.linear(x)
|
||||
return x
|
||||
|
||||
|
||||
class SimpleMLPAdaLN(nn.Module):
|
||||
"""Taken from https://arxiv.org/abs/2406.11838.
|
||||
|
||||
The MLP for Diffusion Loss.
|
||||
:param in_channels: channels in the input Tensor.
|
||||
:param model_channels: base channel count for the model.
|
||||
:param out_channels: channels in the output Tensor.
|
||||
:param cond_channels: channels in the condition.
|
||||
:param num_res_blocks: number of residual blocks per downsample.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
in_channels,
|
||||
model_channels,
|
||||
out_channels,
|
||||
cond_channels,
|
||||
num_res_blocks,
|
||||
num_time_conds=1,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.in_channels = in_channels
|
||||
self.model_channels = model_channels
|
||||
self.out_channels = out_channels
|
||||
self.num_res_blocks = num_res_blocks
|
||||
self.num_time_conds = num_time_conds
|
||||
|
||||
assert num_time_conds != 1
|
||||
self.time_embed = nn.ModuleList(
|
||||
[TimestepEmbedder(model_channels) for _ in range(num_time_conds)]
|
||||
)
|
||||
self.cond_embed = nn.Linear(cond_channels, model_channels)
|
||||
|
||||
self.input_proj = nn.Linear(in_channels, model_channels)
|
||||
|
||||
res_blocks = []
|
||||
for i in range(num_res_blocks):
|
||||
res_blocks.append(ResBlock(model_channels))
|
||||
|
||||
self.res_blocks = nn.ModuleList(res_blocks)
|
||||
self.final_layer = FinalLayer(model_channels, out_channels)
|
||||
|
||||
@classmethod
|
||||
def from_pydantic_config(cls, cfg: FlowLMConfig, latent_dim: int, cond_dim: int) -> Self:
|
||||
config = cfg.flow
|
||||
|
||||
flow_dim = config.dim
|
||||
flow_depth = config.depth
|
||||
num_time_conds = 2
|
||||
return SimpleMLPAdaLN(
|
||||
latent_dim, flow_dim, latent_dim, cond_dim, flow_depth, num_time_conds=num_time_conds
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, c: torch.Tensor, s: torch.Tensor, t: torch.Tensor, x: torch.Tensor
|
||||
) -> torch.Tensor:
|
||||
"""
|
||||
Apply the model to an input batch.
|
||||
:param c: conditioning from AR transformer.
|
||||
:param s: start time tensor.
|
||||
:param t: target time tensor.
|
||||
:param x: an [N x C] Tensor of inputs.
|
||||
:return: an [N x C] Tensor of outputs.
|
||||
"""
|
||||
# Combine time conditions
|
||||
ts = [s, t]
|
||||
x = self.input_proj(x)
|
||||
assert len(ts) == self.num_time_conds, (
|
||||
f"Expected {self.num_time_conds} time conditions, got {len(ts)}"
|
||||
)
|
||||
assert self.num_time_conds != 1
|
||||
t_combined = (
|
||||
sum(self.time_embed[i](ts[i]) for i in range(self.num_time_conds)) / self.num_time_conds
|
||||
)
|
||||
c = self.cond_embed(c)
|
||||
y = t_combined + c
|
||||
|
||||
for block in self.res_blocks:
|
||||
x = block(x, y)
|
||||
|
||||
return self.final_layer(x, y)
|
||||
@@ -0,0 +1,46 @@
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from pocket_tts.modules.conv import StreamingConv1d, StreamingConvTranspose1d
|
||||
|
||||
|
||||
class ConvDownsample1d(nn.Module):
|
||||
"""
|
||||
Downsampling by some integer amount `stride` using convolutions
|
||||
with a kernel size of twice the stride.
|
||||
"""
|
||||
|
||||
def __init__(self, stride: int, dimension: int):
|
||||
super().__init__()
|
||||
self.conv = StreamingConv1d(
|
||||
dimension,
|
||||
dimension,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
groups=1,
|
||||
bias=False,
|
||||
pad_mode="replicate",
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
return self.conv(x, model_state)
|
||||
|
||||
|
||||
class ConvTrUpsample1d(nn.Module):
|
||||
"""
|
||||
Upsample by some integer amount `stride` using transposed convolutions.
|
||||
"""
|
||||
|
||||
def __init__(self, stride: int, dimension: int):
|
||||
super().__init__()
|
||||
self.convtr = StreamingConvTranspose1d(
|
||||
dimension,
|
||||
dimension,
|
||||
kernel_size=2 * stride,
|
||||
stride=stride,
|
||||
groups=dimension,
|
||||
bias=False,
|
||||
)
|
||||
|
||||
def forward(self, x: torch.Tensor, model_state: dict | None):
|
||||
return self.convtr(x, model_state)
|
||||
@@ -0,0 +1,74 @@
|
||||
import math
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def apply_rope(
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
offset: int | torch.Tensor = 0,
|
||||
max_period: int | float = 10_000,
|
||||
):
|
||||
"""
|
||||
Args:
|
||||
q (torch.Tensor): Queries, shape `[B, T, H, D]`.
|
||||
k (torch.Tensor): Keys, shape `[B, T, H, D]`.
|
||||
offset (int): Current offset, e.g. when streaming.
|
||||
max_period (float): Maximum period for the cos and sin.
|
||||
"""
|
||||
|
||||
B, T, H, D = q.shape
|
||||
Bk, Tk, Hk, Dk = k.shape
|
||||
assert (B, T, D) == (Bk, Tk, Dk)
|
||||
assert D > 0
|
||||
assert D % 2 == 0
|
||||
assert max_period > 0
|
||||
|
||||
ds = torch.arange(D // 2, device=q.device, dtype=torch.float32)
|
||||
freqs = torch.exp(ds * (-math.log(max_period) * 2 / D))
|
||||
|
||||
# could be optimized in one call
|
||||
ts = torch.arange(T, device=q.device, dtype=torch.float32)
|
||||
ts += offset
|
||||
ts = ts.view(-1, 1, 1)
|
||||
|
||||
q = q.view(B, T, H, D // 2, 2)
|
||||
k = k.view(B, T, Hk, D // 2, 2)
|
||||
|
||||
# convention is `r` suffix is real part, `i` is imaginary.
|
||||
qr = q[..., 0].float()
|
||||
qi = q[..., 1].float()
|
||||
|
||||
kr = k[..., 0].float()
|
||||
ki = k[..., 1].float()
|
||||
|
||||
rotr = torch.cos(freqs * ts)
|
||||
roti = torch.sin(freqs * ts)
|
||||
qor = qr * rotr - qi * roti
|
||||
qoi = qr * roti + qi * rotr
|
||||
|
||||
kor = kr * rotr - ki * roti
|
||||
koi = kr * roti + ki * rotr
|
||||
|
||||
dtype = q.dtype
|
||||
qo = torch.stack([qor.to(dtype), qoi.to(dtype)], dim=-1)
|
||||
ko = torch.stack([kor.to(dtype), koi.to(dtype)], dim=-1)
|
||||
|
||||
return qo.view(B, T, H, D), ko.view(B, T, Hk, D)
|
||||
|
||||
|
||||
class RotaryEmbedding(nn.Module):
|
||||
"""Rotary positional embedding (RoPE) from [Su et al 2022](https://arxiv.org/abs/2104.09864).
|
||||
|
||||
Args:
|
||||
max_period (float): Maximum period of the rotation frequencies.
|
||||
"""
|
||||
|
||||
def __init__(self, max_period: float | int = 10000.0):
|
||||
super().__init__()
|
||||
self.max_period = max_period
|
||||
|
||||
def forward(self, q: torch.Tensor, k: torch.Tensor, offset: torch.Tensor | int):
|
||||
"""Apply rope rotation to query or key tensor."""
|
||||
return apply_rope(q, k, offset, self.max_period)
|
||||
@@ -0,0 +1,180 @@
|
||||
import numpy as np
|
||||
import torch.nn as nn
|
||||
|
||||
from .conv import StreamingConv1d, StreamingConvTranspose1d
|
||||
|
||||
|
||||
class SEANetResnetBlock(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
dim: int,
|
||||
kernel_sizes: list[int] = [3, 1],
|
||||
dilations: list[int] = [1, 1],
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
assert len(kernel_sizes) == len(dilations), (
|
||||
"Number of kernel sizes should match number of dilations"
|
||||
)
|
||||
hidden = dim // compress
|
||||
block = nn.ModuleList([])
|
||||
for i, (kernel_size, dilation) in enumerate(zip(kernel_sizes, dilations)):
|
||||
in_chs = dim if i == 0 else hidden
|
||||
out_chs = dim if i == len(kernel_sizes) - 1 else hidden
|
||||
block += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(
|
||||
in_chs, out_chs, kernel_size=kernel_size, dilation=dilation, pad_mode=pad_mode
|
||||
),
|
||||
]
|
||||
self.block = block
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
v = x
|
||||
for layer in self.block:
|
||||
if isinstance(layer, StreamingConv1d):
|
||||
v = layer(v, model_state)
|
||||
else:
|
||||
v = layer(v)
|
||||
assert x.shape == v.shape, (x.shape, v.shape, x.shape)
|
||||
return x + v
|
||||
|
||||
|
||||
class SEANetEncoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int = 1,
|
||||
dimension: int = 128,
|
||||
n_filters: int = 32,
|
||||
n_residual_layers: int = 3,
|
||||
ratios: list[int] = [8, 5, 4, 2],
|
||||
kernel_size: int = 7,
|
||||
last_kernel_size: int = 7,
|
||||
residual_kernel_size: int = 3,
|
||||
dilation_base: int = 2,
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.channels = channels
|
||||
self.dimension = dimension
|
||||
self.n_filters = n_filters
|
||||
self.ratios = list(reversed(ratios))
|
||||
del ratios
|
||||
self.n_residual_layers = n_residual_layers
|
||||
self.hop_length = int(np.prod(self.ratios))
|
||||
self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks
|
||||
|
||||
mult = 1
|
||||
model = nn.ModuleList(
|
||||
[StreamingConv1d(channels, mult * n_filters, kernel_size, pad_mode=pad_mode)]
|
||||
)
|
||||
# Downsample to raw audio scale
|
||||
for i, ratio in enumerate(self.ratios):
|
||||
# Add residual layers
|
||||
for j in range(n_residual_layers):
|
||||
model += [
|
||||
SEANetResnetBlock(
|
||||
mult * n_filters,
|
||||
kernel_sizes=[residual_kernel_size, 1],
|
||||
dilations=[dilation_base**j, 1],
|
||||
pad_mode=pad_mode,
|
||||
compress=compress,
|
||||
)
|
||||
]
|
||||
|
||||
# Add downsampling layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(
|
||||
mult * n_filters,
|
||||
mult * n_filters * 2,
|
||||
kernel_size=ratio * 2,
|
||||
stride=ratio,
|
||||
pad_mode=pad_mode,
|
||||
),
|
||||
]
|
||||
mult *= 2
|
||||
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(mult * n_filters, dimension, last_kernel_size, pad_mode=pad_mode),
|
||||
]
|
||||
|
||||
self.model = model
|
||||
|
||||
def forward(self, x, model_state: dict | None):
|
||||
for layer in self.model:
|
||||
if isinstance(layer, (StreamingConv1d, SEANetResnetBlock)):
|
||||
x = layer(x, model_state)
|
||||
else:
|
||||
x = layer(x)
|
||||
return x
|
||||
|
||||
|
||||
class SEANetDecoder(nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
channels: int = 1,
|
||||
dimension: int = 128,
|
||||
n_filters: int = 32,
|
||||
n_residual_layers: int = 3,
|
||||
ratios: list[int] = [8, 5, 4, 2],
|
||||
kernel_size: int = 7,
|
||||
last_kernel_size: int = 7,
|
||||
residual_kernel_size: int = 3,
|
||||
dilation_base: int = 2,
|
||||
pad_mode: str = "reflect",
|
||||
compress: int = 2,
|
||||
):
|
||||
super().__init__()
|
||||
self.dimension = dimension
|
||||
self.channels = channels
|
||||
self.n_filters = n_filters
|
||||
self.ratios = ratios
|
||||
del ratios
|
||||
self.n_residual_layers = n_residual_layers
|
||||
self.hop_length = int(np.prod(self.ratios))
|
||||
self.n_blocks = len(self.ratios) + 2 # first and last conv + residual blocks
|
||||
mult = int(2 ** len(self.ratios))
|
||||
model = nn.ModuleList(
|
||||
[StreamingConv1d(dimension, mult * n_filters, kernel_size, pad_mode=pad_mode)]
|
||||
)
|
||||
# Upsample to raw audio scale
|
||||
for _, ratio in enumerate(self.ratios):
|
||||
# Add upsampling layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConvTranspose1d(
|
||||
mult * n_filters, mult * n_filters // 2, kernel_size=ratio * 2, stride=ratio
|
||||
),
|
||||
]
|
||||
# Add residual layers
|
||||
for j in range(n_residual_layers):
|
||||
model += [
|
||||
SEANetResnetBlock(
|
||||
mult * n_filters // 2,
|
||||
kernel_sizes=[residual_kernel_size, 1],
|
||||
dilations=[dilation_base**j, 1],
|
||||
pad_mode=pad_mode,
|
||||
compress=compress,
|
||||
)
|
||||
]
|
||||
|
||||
mult //= 2
|
||||
|
||||
# Add final layers
|
||||
model += [
|
||||
nn.ELU(alpha=1.0),
|
||||
StreamingConv1d(n_filters, channels, last_kernel_size, pad_mode=pad_mode),
|
||||
]
|
||||
self.model = model
|
||||
|
||||
def forward(self, z, model_state: dict | None):
|
||||
for layer in self.model:
|
||||
if isinstance(layer, (StreamingConvTranspose1d, SEANetResnetBlock, StreamingConv1d)):
|
||||
z = layer(z, model_state)
|
||||
else:
|
||||
z = layer(z)
|
||||
return z
|
||||
@@ -0,0 +1,45 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
def init_states(
|
||||
model: nn.Module, batch_size: int, sequence_length: int
|
||||
) -> dict[str, dict[str, torch.Tensor]]:
|
||||
result = {}
|
||||
for module_name, module in model.named_modules():
|
||||
if not isinstance(module, StatefulModule):
|
||||
continue
|
||||
module._module_absolute_name = module_name
|
||||
module_state = module.init_state(batch_size, sequence_length=sequence_length)
|
||||
result[module_name] = module_state
|
||||
return result
|
||||
|
||||
|
||||
def increment_steps(
|
||||
module: nn.Module, model_state: dict[str, dict[str, torch.Tensor]], increment: int = 1
|
||||
):
|
||||
# print("incrementing steps by", increment)
|
||||
for module_name, module in module.named_modules():
|
||||
if not isinstance(module, StatefulModule):
|
||||
continue
|
||||
module.increment_step(model_state[module_name], increment)
|
||||
|
||||
|
||||
class StatefulModule(ABC, nn.Module):
|
||||
def __init__(self, *args, **kwds):
|
||||
self._module_absolute_name = None
|
||||
return super().__init__(*args, **kwds)
|
||||
|
||||
@abstractmethod
|
||||
def init_state(self, batch_size: int, sequence_length: int):
|
||||
"""Initialize the state."""
|
||||
raise NotImplementedError
|
||||
|
||||
def increment_step(self, state: dict, increment: int = 1):
|
||||
pass
|
||||
|
||||
def get_state(self, model_state: dict[str, dict[str, torch.Tensor]]) -> dict[str, torch.Tensor]:
|
||||
"""Get the state for this module from the model state."""
|
||||
return model_state[self._module_absolute_name]
|
||||
@@ -0,0 +1,126 @@
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
from torch.nn import functional as F
|
||||
|
||||
from pocket_tts.modules.rope import RotaryEmbedding
|
||||
from pocket_tts.modules.stateful_module import StatefulModule
|
||||
|
||||
|
||||
def complete_kv(
|
||||
cache: torch.Tensor, current_end: torch.Tensor, k: torch.Tensor, v: torch.Tensor
|
||||
) -> tuple[torch.Tensor, torch.Tensor]:
|
||||
current_end = current_end.shape[0]
|
||||
|
||||
cache[0, :, current_end : current_end + k.shape[1]] = k
|
||||
cache[1, :, current_end : current_end + v.shape[1]] = v
|
||||
valid = cache[:, :, : current_end + k.shape[1]]
|
||||
return valid[0], valid[1]
|
||||
|
||||
|
||||
def _materialize_causal_mask(
|
||||
shape: tuple[int, ...], shift: int, device: str | torch.device = "cpu"
|
||||
) -> torch.Tensor:
|
||||
dtype = torch.float32
|
||||
|
||||
num_queries, num_keys = shape[-2:]
|
||||
shift = num_keys - num_queries
|
||||
|
||||
tensor = torch.full(shape, dtype=dtype, fill_value=1, device=device)
|
||||
mask = torch.tril(tensor, diagonal=shift).to(dtype)
|
||||
mask = torch.log(mask)
|
||||
return mask.to(dtype)
|
||||
|
||||
|
||||
class StreamingMultiheadAttention(StatefulModule):
|
||||
"""Similar to `nn.MultiheadAttention` but with support for streaming.
|
||||
|
||||
Args:
|
||||
embed_dim (int): Dimension to project to.
|
||||
num_heads (int): Number of heads.
|
||||
context (int, optional): Number of time steps the attention can access to.
|
||||
Can access `context` time steps into the past.
|
||||
rope (`RotaryEmbedding`, optional): Rope embedding to use.
|
||||
device (torch.device, optional): Device on which to initialize.
|
||||
dtype (torch.dtype, optional): dtype to use.
|
||||
"""
|
||||
|
||||
def __init__(self, embed_dim: int, num_heads: int, rope: RotaryEmbedding):
|
||||
super().__init__()
|
||||
|
||||
self.embed_dim = embed_dim
|
||||
self.rope = rope
|
||||
self.num_heads = num_heads
|
||||
|
||||
out_dim = embed_dim
|
||||
num_kv = num_heads
|
||||
kv_dim = (embed_dim // num_heads) * num_kv
|
||||
out_dim += 2 * kv_dim
|
||||
mult = 1
|
||||
in_proj = nn.Linear(embed_dim, mult * out_dim, bias=False)
|
||||
# We try to follow the default PyTorch MHA convention, to easily compare results.
|
||||
self.in_proj_weight = in_proj.weight
|
||||
self.out_proj = nn.Linear(embed_dim, mult * embed_dim, bias=False)
|
||||
|
||||
def _get_mask(self, shape: tuple[int, int], shift: int, device: torch.device) -> torch.Tensor:
|
||||
return _materialize_causal_mask(shape, shift=shift, device=device)
|
||||
|
||||
def init_state(self, batch_size: int, sequence_length: int) -> dict[str, torch.Tensor]:
|
||||
dim_per_head = self.embed_dim // self.num_heads
|
||||
initial_current_end = torch.zeros((0,)).to(self.in_proj_weight.device)
|
||||
return dict(
|
||||
current_end=initial_current_end,
|
||||
cache=torch.full(
|
||||
(2, batch_size, sequence_length, self.num_heads, dim_per_head),
|
||||
float("NaN"),
|
||||
device=self.in_proj_weight.device,
|
||||
dtype=self.in_proj_weight.dtype,
|
||||
),
|
||||
)
|
||||
|
||||
def increment_step(self, state: dict, increment: int = 1):
|
||||
new_size = state["current_end"].shape[0] + increment
|
||||
state["current_end"] = torch.zeros((new_size,)).to(state["current_end"].device)
|
||||
|
||||
def _complete_kv(self, k, v, state: dict | None):
|
||||
k, v = complete_kv(state["cache"], state["current_end"], k, v)
|
||||
return k, v
|
||||
|
||||
def _apply_rope(self, query: torch.Tensor, key: torch.Tensor, state: dict | None):
|
||||
# Apply rope embeddings to query and key tensors.
|
||||
streaming_offset = self._streaming_offset(state)
|
||||
return self.rope(query, key, offset=streaming_offset)
|
||||
|
||||
def _streaming_offset(self, state: dict | None) -> torch.Tensor | int:
|
||||
return state["current_end"].shape[0]
|
||||
|
||||
def check_model_state(self, model_state: dict):
|
||||
if model_state is None:
|
||||
raise ValueError("model_state must be provided")
|
||||
return self.get_state(model_state)
|
||||
|
||||
def forward(self, query: torch.Tensor, model_state: dict | None):
|
||||
state = self.check_model_state(model_state)
|
||||
|
||||
projected = nn.functional.linear(query, self.in_proj_weight)
|
||||
# Reshape from (b, t, p*h*d) to (b, t, p, h, d) where p=3, h=num_heads
|
||||
b, t, _ = projected.shape
|
||||
d = self.embed_dim // self.num_heads
|
||||
packed = projected.view(b, t, 3, self.num_heads, d)
|
||||
q, k, v = torch.unbind(packed, dim=2)
|
||||
q, k = self._apply_rope(q, k, state)
|
||||
k, v = self._complete_kv(k, v, state)
|
||||
|
||||
mask_shape = (query.shape[1], query.shape[1] + state["current_end"].shape[0])
|
||||
shift = state["current_end"].shape[0]
|
||||
|
||||
attn_mask = self._get_mask(mask_shape, shift=shift, device=q.device)
|
||||
|
||||
q, k, v = [x.transpose(1, 2) for x in (q, k, v)]
|
||||
x = F.scaled_dot_product_attention(q, k, v, attn_mask)
|
||||
x = x.transpose(1, 2)
|
||||
# Reshape from (b, t, h, d) to (b, t, h*d)
|
||||
b, t, h, d = x.shape
|
||||
x = x.reshape(b, t, h * d)
|
||||
x = self.out_proj(x)
|
||||
|
||||
return x
|
||||
Reference in New Issue
Block a user