chatterbox-light 1.0.1__tar.gz

chatterbox_light-1.0.1/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Resemble AI
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

chatterbox_light-1.0.1/PKG-INFO

@@ -0,0 +1,141 @@
+Metadata-Version: 2.4
+Name: chatterbox-light
+Version: 1.0.1
+Summary: Lightweight fork of Resemble AI's Chatterbox text-to-speech and voice conversion, with fewer dependencies and no watermarking
+Author: Resemble AI
+Maintainer: Blair Chintella
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/BBC-Esq/chatterbox-light
+Project-URL: Repository, https://github.com/BBC-Esq/chatterbox-light
+Project-URL: Upstream, https://github.com/resemble-ai/chatterbox
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: torchaudio
+Requires-Dist: transformers>=4.46.0
+Requires-Dist: diffusers
+Requires-Dist: safetensors
+Requires-Dist: s3tokenizer
+Requires-Dist: einops
+Requires-Dist: scipy
+Requires-Dist: pyloudnorm
+Requires-Dist: soundfile
+Requires-Dist: tokenizers
+Requires-Dist: tqdm
+Requires-Dist: huggingface-hub
+Provides-Extra: multilingual
+Requires-Dist: pykakasi; extra == "multilingual"
+Requires-Dist: spacy-pkuseg; extra == "multilingual"
+Dynamic: license-file
+
+## This is a light version of chatterbox that removes:
+* ```librosa``` and all its required dependencies and instead uses ```torchaudio``` and ```scipy```
+* ```perth``` because who the heck needs watermarking anyways
+* ```omegaconf``` and instead use standard python
+* others I forget, but everything works
+
+However, I did add ```soundfile``` because I like it.
+
+# Installation
+>Go through these steps in order.
+```
+python -m venv .
+```
+
+```
+.\Scripts\activate
+```
+
+```
+python.exe -m pip install --upgrade pip
+```
+
+```
+pip install uv
+```
+
+Next, make sure appropriate versions of torch, torchaudio, and CUDA are installed.
+
+```
+uv pip install chatterbox-light
+```
+
+To install the latest code from GitHub instead:
+```
+uv pip install -r requirements.txt
+```
+
+```
+pip install git+https://github.com/BBC-Esq/chatterbox-light.git --no-deps
+```
+
+This package installs the `chatterbox` module, just like Resemble AI's `chatterbox-tts`, so don't install both in one environment. Installs of this fork made from GitHub before version 1.0.1 are also named `chatterbox-tts`; run `pip uninstall chatterbox-tts` before installing this.
+
+Optional: if you'll use the multilingual model with Japanese or Chinese text, also install these (or install `"chatterbox-light[multilingual]"`). Without them, Japanese kanji are dropped and Chinese text isn't split into words. The first time Chinese text is used, spacy-pkuseg downloads a ~35 MB model to `~/.pkuseg`.
+```
+uv pip install pykakasi spacy-pkuseg
+```
+
+# Usage
+Each model's `generate()` returns a `(1, samples)` tensor at `model.sr` (24 kHz). Models download from Hugging Face the first time they're used. Save the audio with `soundfile`:
+```python
+import soundfile as sf
+sf.write("output.wav", wav.squeeze(0).numpy(), model.sr)
+```
+
+### Chatterbox (English)
+```python
+from chatterbox.tts import ChatterboxTTS
+
+model = ChatterboxTTS.from_pretrained(device="cuda")  # or "cpu" / "mps"
+wav = model.generate("Hello there, this is Chatterbox.")
+
+# clone a voice from a reference clip
+wav = model.generate("Hello there, this is Chatterbox.", audio_prompt_path="reference.wav")
+
+# more expressive: raise exaggeration (default 0.5) and lower cfg_weight (default 0.5) for slower pacing
+wav = model.generate("Hello there, this is Chatterbox.", exaggeration=0.7, cfg_weight=0.3)
+```
+
+### Chatterbox Turbo (English)
+Faster, and understands paralinguistic tags such as `[laugh]`, `[chuckle]` and `[cough]`. A reference clip must be longer than 5 seconds. `cfg_weight`, `exaggeration` and `min_p` are ignored.
+```python
+from chatterbox.tts_turbo import ChatterboxTurboTTS
+
+model = ChatterboxTurboTTS.from_pretrained(device="cuda")
+wav = model.generate("Hi there [chuckle], have you got a minute to chat?", audio_prompt_path="reference.wav")
+```
+
+### Chatterbox Nano (English)
+A smaller version of Turbo for CPU and tight memory budgets, loaded through the same class with `nano=True`:
+```python
+from chatterbox.tts_turbo import ChatterboxTurboTTS
+
+model = ChatterboxTurboTTS.from_pretrained(device="cpu", nano=True)
+wav = model.generate("Hi there [chuckle], have you got a minute to chat?")
+```
+
+### Chatterbox Multilingual (23 languages)
+The v2 checkpoint loads by default. Pass `t3_model="v3"` for the newer v3 checkpoint, which upstream recommends (fewer hallucinations, better speaker similarity). For Japanese or Chinese text, install the optional packages above.
+```python
+from chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
+
+print(SUPPORTED_LANGUAGES)  # language codes such as "fr", "de", "ja", "zh"
+model = ChatterboxMultilingualTTS.from_pretrained(device="cuda", t3_model="v3")
+wav = model.generate("Bonjour, comment ça va ?", language_id="fr")
+wav = model.generate("Hola, ¿cómo estás?", language_id="es", audio_prompt_path="reference.wav")
+```
+
+### Voice conversion
+```python
+from chatterbox.vc import ChatterboxVC
+
+model = ChatterboxVC.from_pretrained(device="cuda")
+wav = model.generate("input.wav", target_voice_path="target_voice.wav")
+```
+
+See the `example_*.py` scripts for complete runs.
+
+Enjoy!

chatterbox_light-1.0.1/README.md

@@ -0,0 +1,109 @@
+## This is a light version of chatterbox that removes:
+* ```librosa``` and all its required dependencies and instead uses ```torchaudio``` and ```scipy```
+* ```perth``` because who the heck needs watermarking anyways
+* ```omegaconf``` and instead use standard python
+* others I forget, but everything works
+
+However, I did add ```soundfile``` because I like it.
+
+# Installation
+>Go through these steps in order.
+```
+python -m venv .
+```
+
+```
+.\Scripts\activate
+```
+
+```
+python.exe -m pip install --upgrade pip
+```
+
+```
+pip install uv
+```
+
+Next, make sure appropriate versions of torch, torchaudio, and CUDA are installed.
+
+```
+uv pip install chatterbox-light
+```
+
+To install the latest code from GitHub instead:
+```
+uv pip install -r requirements.txt
+```
+
+```
+pip install git+https://github.com/BBC-Esq/chatterbox-light.git --no-deps
+```
+
+This package installs the `chatterbox` module, just like Resemble AI's `chatterbox-tts`, so don't install both in one environment. Installs of this fork made from GitHub before version 1.0.1 are also named `chatterbox-tts`; run `pip uninstall chatterbox-tts` before installing this.
+
+Optional: if you'll use the multilingual model with Japanese or Chinese text, also install these (or install `"chatterbox-light[multilingual]"`). Without them, Japanese kanji are dropped and Chinese text isn't split into words. The first time Chinese text is used, spacy-pkuseg downloads a ~35 MB model to `~/.pkuseg`.
+```
+uv pip install pykakasi spacy-pkuseg
+```
+
+# Usage
+Each model's `generate()` returns a `(1, samples)` tensor at `model.sr` (24 kHz). Models download from Hugging Face the first time they're used. Save the audio with `soundfile`:
+```python
+import soundfile as sf
+sf.write("output.wav", wav.squeeze(0).numpy(), model.sr)
+```
+
+### Chatterbox (English)
+```python
+from chatterbox.tts import ChatterboxTTS
+
+model = ChatterboxTTS.from_pretrained(device="cuda")  # or "cpu" / "mps"
+wav = model.generate("Hello there, this is Chatterbox.")
+
+# clone a voice from a reference clip
+wav = model.generate("Hello there, this is Chatterbox.", audio_prompt_path="reference.wav")
+
+# more expressive: raise exaggeration (default 0.5) and lower cfg_weight (default 0.5) for slower pacing
+wav = model.generate("Hello there, this is Chatterbox.", exaggeration=0.7, cfg_weight=0.3)
+```
+
+### Chatterbox Turbo (English)
+Faster, and understands paralinguistic tags such as `[laugh]`, `[chuckle]` and `[cough]`. A reference clip must be longer than 5 seconds. `cfg_weight`, `exaggeration` and `min_p` are ignored.
+```python
+from chatterbox.tts_turbo import ChatterboxTurboTTS
+
+model = ChatterboxTurboTTS.from_pretrained(device="cuda")
+wav = model.generate("Hi there [chuckle], have you got a minute to chat?", audio_prompt_path="reference.wav")
+```
+
+### Chatterbox Nano (English)
+A smaller version of Turbo for CPU and tight memory budgets, loaded through the same class with `nano=True`:
+```python
+from chatterbox.tts_turbo import ChatterboxTurboTTS
+
+model = ChatterboxTurboTTS.from_pretrained(device="cpu", nano=True)
+wav = model.generate("Hi there [chuckle], have you got a minute to chat?")
+```
+
+### Chatterbox Multilingual (23 languages)
+The v2 checkpoint loads by default. Pass `t3_model="v3"` for the newer v3 checkpoint, which upstream recommends (fewer hallucinations, better speaker similarity). For Japanese or Chinese text, install the optional packages above.
+```python
+from chatterbox.mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES
+
+print(SUPPORTED_LANGUAGES)  # language codes such as "fr", "de", "ja", "zh"
+model = ChatterboxMultilingualTTS.from_pretrained(device="cuda", t3_model="v3")
+wav = model.generate("Bonjour, comment ça va ?", language_id="fr")
+wav = model.generate("Hola, ¿cómo estás?", language_id="es", audio_prompt_path="reference.wav")
+```
+
+### Voice conversion
+```python
+from chatterbox.vc import ChatterboxVC
+
+model = ChatterboxVC.from_pretrained(device="cuda")
+wav = model.generate("input.wav", target_voice_path="target_voice.wav")
+```
+
+See the `example_*.py` scripts for complete runs.
+
+Enjoy!

chatterbox_light-1.0.1/pyproject.toml

@@ -0,0 +1,48 @@
+[project]
+name = "chatterbox-light"
+version = "1.0.1"
+description = "Lightweight fork of Resemble AI's Chatterbox text-to-speech and voice conversion, with fewer dependencies and no watermarking"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+license-files = ["LICENSE"]
+authors = [
+    {name = "Resemble AI"},
+]
+maintainers = [
+    {name = "Blair Chintella"},
+]
+dependencies = [
+    "numpy",
+    "torch",
+    "torchaudio",
+    "transformers>=4.46.0",
+    "diffusers",
+    "safetensors",
+    "s3tokenizer",
+    "einops",
+    "scipy",
+    "pyloudnorm",
+    "soundfile",
+    "tokenizers",
+    "tqdm",
+    "huggingface-hub",
+]
+
+[project.optional-dependencies]
+multilingual = [
+    "pykakasi",
+    "spacy-pkuseg",
+]
+
+[project.urls]
+Homepage = "https://github.com/BBC-Esq/chatterbox-light"
+Repository = "https://github.com/BBC-Esq/chatterbox-light"
+Upstream = "https://github.com/resemble-ai/chatterbox"
+
+[build-system]
+requires = ["setuptools>=77.0.0"]
+build-backend = "setuptools.build_meta"
+
+[tool.setuptools.packages.find]
+where = ["src"]

chatterbox_light-1.0.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

chatterbox_light-1.0.1/src/chatterbox/__init__.py

@@ -0,0 +1,11 @@
+try:
+    from importlib.metadata import version
+except ImportError:
+    from importlib_metadata import version  # For Python <3.8
+
+__version__ = version("chatterbox-light")
+
+
+from .tts import ChatterboxTTS
+from .vc import ChatterboxVC
+from .mtl_tts import ChatterboxMultilingualTTS, SUPPORTED_LANGUAGES

chatterbox_light-1.0.1/src/chatterbox/models/s3gen/__init__.py

@@ -0,0 +1,2 @@
+from .s3gen import S3Token2Wav as S3Gen
+from .const import S3GEN_SR

chatterbox_light-1.0.1/src/chatterbox/models/s3gen/configs.py

@@ -0,0 +1,10 @@
+from ..utils import AttrDict
+
+CFM_PARAMS = AttrDict({
+    "sigma_min": 1e-06,
+    "solver": "euler",
+    "t_scheduler": "cosine",
+    "training_cfg_rate": 0.2,
+    "inference_cfg_rate": 0.7,
+    "reg_loss_type": "l1"
+})

chatterbox_light-1.0.1/src/chatterbox/models/s3gen/const.py

@@ -0,0 +1,2 @@
+S3GEN_SR = 24000
+S3GEN_SIL = 4299

chatterbox_light-1.0.1/src/chatterbox/models/s3gen/decoder.py

@@ -0,0 +1,292 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import pack, rearrange, repeat
+
+from .utils.mask import add_optional_chunk_mask
+from .matcha.decoder import SinusoidalPosEmb, Block1D, ResnetBlock1D, Downsample1D, \
+    TimestepEmbedding, Upsample1D
+from .matcha.transformer import BasicTransformerBlock
+from .utils.intmeanflow import get_intmeanflow_time_mixer
+
+
+def mask_to_bias(mask: torch.Tensor, dtype: torch.dtype) -> torch.Tensor:
+    assert mask.dtype == torch.bool
+    assert dtype in [torch.float32, torch.bfloat16, torch.float16]
+    mask = mask.to(dtype)
+    mask = (1.0 - mask) * -1.0e+10
+    return mask
+
+
+
+class Transpose(torch.nn.Module):
+    def __init__(self, dim0: int, dim1: int):
+        super().__init__()
+        self.dim0 = dim0
+        self.dim1 = dim1
+
+    def forward(self, x: torch.Tensor):
+        x = torch.transpose(x, self.dim0, self.dim1)
+        return x
+
+
+class CausalBlock1D(Block1D):
+    def __init__(self, dim: int, dim_out: int):
+        super(CausalBlock1D, self).__init__(dim, dim_out)
+        self.block = torch.nn.Sequential(
+            CausalConv1d(dim, dim_out, 3),
+            Transpose(1, 2),
+            nn.LayerNorm(dim_out),
+            Transpose(1, 2),
+            nn.Mish(),
+        )
+
+    def forward(self, x: torch.Tensor, mask: torch.Tensor):
+        output = self.block(x * mask)
+        return output * mask
+
+
+class CausalResnetBlock1D(ResnetBlock1D):
+    def __init__(self, dim: int, dim_out: int, time_emb_dim: int, groups: int = 8):
+        super(CausalResnetBlock1D, self).__init__(dim, dim_out, time_emb_dim, groups)
+        self.block1 = CausalBlock1D(dim, dim_out)
+        self.block2 = CausalBlock1D(dim_out, dim_out)
+
+
+class CausalConv1d(torch.nn.Conv1d):
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        kernel_size: int,
+        stride: int = 1,
+        dilation: int = 1,
+        groups: int = 1,
+        bias: bool = True,
+        padding_mode: str = 'zeros',
+        device=None,
+        dtype=None
+    ) -> None:
+        super(CausalConv1d, self).__init__(in_channels, out_channels,
+                                           kernel_size, stride,
+                                           padding=0, dilation=dilation,
+                                           groups=groups, bias=bias,
+                                           padding_mode=padding_mode,
+                                           device=device, dtype=dtype)
+        assert stride == 1
+        self.causal_padding = (kernel_size - 1, 0)
+
+    def forward(self, x: torch.Tensor):
+        x = F.pad(x, self.causal_padding)
+        x = super(CausalConv1d, self).forward(x)
+        return x
+class ConditionalDecoder(nn.Module):
+    def __init__(
+        self,
+        in_channels=320,
+        out_channels=80,
+        causal=True,
+        channels=[256],
+        dropout=0.0,
+        attention_head_dim=64,
+        n_blocks=4,
+        num_mid_blocks=12,
+        num_heads=8,
+        act_fn="gelu",
+        meanflow=False,
+    ):
+        super().__init__()
+        channels = tuple(channels)
+        self.meanflow = meanflow
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+        self.causal = causal
+        self.time_embeddings = SinusoidalPosEmb(in_channels)
+        time_embed_dim = channels[0] * 4
+        self.time_mlp = TimestepEmbedding(
+            in_channels=in_channels,
+            time_embed_dim=time_embed_dim,
+            act_fn="silu",
+        )
+
+        self.down_blocks = nn.ModuleList([])
+        self.mid_blocks = nn.ModuleList([])
+        self.up_blocks = nn.ModuleList([])
+
+        self.static_chunk_size = 0
+
+        output_channel = in_channels
+        for i in range(len(channels)):
+            input_channel = output_channel
+            output_channel = channels[i]
+            is_last = i == len(channels) - 1
+            resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) if self.causal else \
+                ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            downsample = (
+                Downsample1D(output_channel) if not is_last else
+                CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.down_blocks.append(nn.ModuleList([resnet, transformer_blocks, downsample]))
+
+        for _ in range(num_mid_blocks):
+            input_channel = channels[-1]
+            out_channels = channels[-1]
+            resnet = CausalResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim) if self.causal else \
+                ResnetBlock1D(dim=input_channel, dim_out=output_channel, time_emb_dim=time_embed_dim)
+
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+
+            self.mid_blocks.append(nn.ModuleList([resnet, transformer_blocks]))
+
+        channels = channels[::-1] + (channels[0],)
+        for i in range(len(channels) - 1):
+            input_channel = channels[i] * 2
+            output_channel = channels[i + 1]
+            is_last = i == len(channels) - 2
+            resnet = CausalResnetBlock1D(
+                dim=input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            ) if self.causal else ResnetBlock1D(
+                dim=input_channel,
+                dim_out=output_channel,
+                time_emb_dim=time_embed_dim,
+            )
+            transformer_blocks = nn.ModuleList(
+                [
+                    BasicTransformerBlock(
+                        dim=output_channel,
+                        num_attention_heads=num_heads,
+                        attention_head_dim=attention_head_dim,
+                        dropout=dropout,
+                        activation_fn=act_fn,
+                    )
+                    for _ in range(n_blocks)
+                ]
+            )
+            upsample = (
+                Upsample1D(output_channel, use_conv_transpose=True)
+                if not is_last
+                else CausalConv1d(output_channel, output_channel, 3) if self.causal else nn.Conv1d(output_channel, output_channel, 3, padding=1)
+            )
+            self.up_blocks.append(nn.ModuleList([resnet, transformer_blocks, upsample]))
+        self.final_block = CausalBlock1D(channels[-1], channels[-1]) if self.causal else Block1D(channels[-1], channels[-1])
+        self.final_proj = nn.Conv1d(channels[-1], self.out_channels, 1)
+        self.initialize_weights()
+        self.time_embed_mixer = None
+        if self.meanflow:
+            self.time_embed_mixer = get_intmeanflow_time_mixer(time_embed_dim)
+
+
+    @property
+    def dtype(self):
+        return self.final_proj.weight.dtype
+
+    def initialize_weights(self):
+        for m in self.modules():
+            if isinstance(m, nn.Conv1d):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.GroupNorm):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+            elif isinstance(m, nn.Linear):
+                nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+                if m.bias is not None:
+                    nn.init.constant_(m.bias, 0)
+
+    def forward(self, x, mask, mu, t, spks=None, cond=None, r=None):
+        t = self.time_embeddings(t).to(t.dtype)
+        t = self.time_mlp(t)
+
+        if self.meanflow:
+            r = self.time_embeddings(r).to(t.dtype)
+            r = self.time_mlp(r)
+            concat_embed = torch.cat([t, r], dim=1)
+            t = self.time_embed_mixer(concat_embed)
+
+        x = pack([x, mu], "b * t")[0]
+
+        if spks is not None:
+            spks = repeat(spks, "b c -> b c t", t=x.shape[-1])
+            x = pack([x, spks], "b * t")[0]
+        if cond is not None:
+            x = pack([x, cond], "b * t")[0]
+
+        hiddens = []
+        masks = [mask]
+        for resnet, transformer_blocks, downsample in self.down_blocks:
+            mask_down = masks[-1]
+            x = resnet(x, mask_down, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = add_optional_chunk_mask(x, mask_down.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            hiddens.append(x)
+            x = downsample(x * mask_down)
+            masks.append(mask_down[:, :, ::2])
+        masks = masks[:-1]
+        mask_mid = masks[-1]
+
+        for resnet, transformer_blocks in self.mid_blocks:
+            x = resnet(x, mask_mid, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = add_optional_chunk_mask(x, mask_mid.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+
+        for resnet, transformer_blocks, upsample in self.up_blocks:
+            mask_up = masks.pop()
+            skip = hiddens.pop()
+            x = pack([x[:, :, :skip.shape[-1]], skip], "b * t")[0]
+            x = resnet(x, mask_up, t)
+            x = rearrange(x, "b c t -> b t c").contiguous()
+            attn_mask = add_optional_chunk_mask(x, mask_up.bool(), False, False, 0, self.static_chunk_size, -1)
+            attn_mask = mask_to_bias(attn_mask == 1, x.dtype)
+            for transformer_block in transformer_blocks:
+                x = transformer_block(
+                    hidden_states=x,
+                    attention_mask=attn_mask,
+                    timestep=t,
+                )
+            x = rearrange(x, "b t c -> b c t").contiguous()
+            x = upsample(x * mask_up)
+        x = self.final_block(x, mask_up)
+        output = self.final_proj(x * mask_up)
+        return output * mask