pocket-tts 1.0.2__py3-none-any.whl

pocket_tts/modules/__init__.py

@@ -0,0 +1 @@
+"""Modules used for building the models."""

pocket_tts/modules/conv.py

@@ -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

pocket_tts/modules/dummy_quantizer.py

@@ -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)

pocket_tts/modules/layer_scale.py

@@ -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

pocket_tts/modules/mimi_transformer.py

@@ -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

pocket_tts/modules/mlp.py

@@ -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)