Add ABot-World

docs/source/en/_toctree.yml

@@ -489,6 +489,8 @@
         title: Oobleck AutoEncoder
       - local: api/models/autoencoder_tiny
         title: Tiny AutoEncoder
+      - local: api/models/autoencoder_tiny_video
+        title: Tiny Video AutoEncoder
       - local: api/models/vq
         title: VQModel
       title: VAEs

docs/source/en/api/models/autoencoder_tiny_video.md

@@ -0,0 +1,37 @@
+<!--Copyright 2026 The HuggingFace Team. All rights reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except in compliance with
+the License. You may obtain a copy of the License at
+
+http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software distributed under the License is distributed on
+an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
+specific language governing permissions and limitations under the License.
+-->
+
+# Tiny Video AutoEncoder
+
+Tiny AutoEncoder for Hunyuan Video (TAEHV) was introduced in [madebyollin/taehv](https://github.com/madebyollin/taehv) by Ollin Boer Bohan. It is a family of tiny causal video autoencoders distilled from full video VAEs — `taew2_2` decodes the Wan 2.2 latent space of [`AutoencoderKLWan`] about 50× faster than the full model — for previews and real-time decoding. Latents are the normalized (roughly unit Gaussian) latents of the full VAE.
+
+Decode a video chunk by chunk with a [`TinyVideoDecodeCache`]: each call decodes only the new latent frames, continuing from the previous calls, and the result is identical to a single decode of all frames.
+
+```python
+import torch
+from diffusers import AutoencoderTinyVideo
+from diffusers.models.autoencoders.autoencoder_tiny_video import TinyVideoDecodeCache
+
+vae = AutoencoderTinyVideo.from_pretrained("YiYiXu/taew2_2-diffusers", dtype=torch.bfloat16).to("cuda")
+
+cache = TinyVideoDecodeCache()
+for latents in latent_chunks:  # [B, 48, T, h, w], normalized Wan 2.2 latents
+    frames = vae.decode(latents, cache=cache).sample  # [B, 3, 4 * T, 16 * h, 16 * w] in [-1, 1]
+```
+
+## AutoencoderTinyVideo
+
+[[autodoc]] AutoencoderTinyVideo
+
+## TinyVideoDecodeCache
+
+[[autodoc]] models.autoencoders.autoencoder_tiny_video.TinyVideoDecodeCache

scripts/convert_abot_world_to_diffusers.py

@@ -0,0 +1,57 @@
+# Convert the ABot-World checkpoint (https://huggingface.co/acvlab/ABot-World-0-5B-LF) to diffusers format.
+#
+#   python scripts/convert_abot_world_to_diffusers.py \
+#       --checkpoint_path <repo>/diffusion_pytorch_model.safetensors --output_path <out_dir> [--dtype bf16]
+import argparse
+
+import torch
+from safetensors.torch import load_file
+
+from diffusers import ABotWorldTransformer3DModel
+
+
+def convert_abot_world_transformer(state_dict):
+    """Map the reference CausalWanModel state dict to ABotWorldTransformer3DModel naming."""
+    converted = {}
+    for key, value in state_dict.items():
+        new_key = key
+        new_key = new_key.replace("text_embedding.0.", "condition_embedder.text_embedder.0.")
+        new_key = new_key.replace("text_embedding.2.", "condition_embedder.text_embedder.2.")
+        new_key = new_key.replace("time_embedding.0.", "condition_embedder.time_embedder.0.")
+        new_key = new_key.replace("time_embedding.2.", "condition_embedder.time_embedder.2.")
+        new_key = new_key.replace("time_projection.1.", "condition_embedder.time_proj.1.")
+        if ".self_attn." in new_key or ".cross_attn." in new_key:
+            new_key = new_key.replace(".self_attn.", ".attn1.").replace(".cross_attn.", ".attn2.")
+            new_key = new_key.replace(".q.", ".to_q.").replace(".k.", ".to_k.").replace(".v.", ".to_v.")
+            new_key = new_key.replace(".o.", ".to_out.0.")
+        new_key = new_key.replace(".norm3.", ".norm2.")  # the cross-attn LayerNorm
+        if new_key.endswith(".modulation"):
+            new_key = new_key.replace("head.modulation", "scale_shift_table")
+            new_key = new_key.replace(".modulation", ".scale_shift_table")
+        new_key = new_key.replace("head.head.", "proj_out.")
+        converted[new_key] = value
+    return converted
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint_path", type=str, required=True)
+    parser.add_argument("--output_path", type=str, required=True)
+    parser.add_argument("--dtype", type=str, default="bf16", choices=["bf16", "fp32"])
+    args = parser.parse_args()
+
+    state_dict = convert_abot_world_transformer(load_file(args.checkpoint_path))
+
+    transformer = ABotWorldTransformer3DModel()
+    transformer.load_state_dict(state_dict, strict=True)
+    if args.dtype == "bf16":
+        transformer = transformer.to(torch.bfloat16)
+    transformer.save_pretrained(args.output_path)
+
+    # round-trip check
+    ABotWorldTransformer3DModel.from_pretrained(args.output_path)
+    print(f"saved and round-trip loaded: {args.output_path}")
+
+
+if __name__ == "__main__":
+    main()

scripts/convert_taehv_to_diffusers.py

@@ -0,0 +1,103 @@
+"""
+Convert a TAEHV checkpoint (https://github.com/madebyollin/taehv, e.g. `taew2_2.pth` for the Wan 2.2 VAE) to an
+`AutoencoderTinyVideo`:
+
+    python scripts/convert_taehv_to_diffusers.py --checkpoint_path taew2_2.pth --variant taew2_2 --output_path ./taew2_2
+"""
+
+import argparse
+
+import torch
+
+from diffusers import AutoencoderTinyVideo
+
+
+# TAEHV model configs, keyed by checkpoint name
+VARIANTS = {
+    "taehv": {"latent_channels": 16, "patch_size": 1},  # Hunyuan Video
+    "taew2_1": {"latent_channels": 16, "patch_size": 1},  # Wan 2.1
+    "taew2_2": {"latent_channels": 48, "patch_size": 2},  # Wan 2.2
+    "taehv1_5": {"latent_channels": 32, "patch_size": 2},  # Hunyuan Video 1.5
+    "taeltx": {  # LTX-2 / LTX-2.3
+        "latent_channels": 128,
+        "patch_size": 4,
+        "encoder_time_downscale": (True, True, True),
+        "decoder_time_upscale": (True, True, True),
+    },
+}
+
+# the reference builds the encoder/decoder as `nn.Sequential`; these are the module names at each index
+ENCODER_LAYERS = {
+    0: "conv_in",
+    2: "blocks.0.time_pool",
+    3: "blocks.0.conv_down",
+    4: "blocks.0.mem_blocks.0",
+    5: "blocks.0.mem_blocks.1",
+    6: "blocks.0.mem_blocks.2",
+    7: "blocks.1.time_pool",
+    8: "blocks.1.conv_down",
+    9: "blocks.1.mem_blocks.0",
+    10: "blocks.1.mem_blocks.1",
+    11: "blocks.1.mem_blocks.2",
+    12: "blocks.2.time_pool",
+    13: "blocks.2.conv_down",
+    14: "blocks.2.mem_blocks.0",
+    15: "blocks.2.mem_blocks.1",
+    16: "blocks.2.mem_blocks.2",
+    17: "conv_out",
+}
+DECODER_LAYERS = {
+    1: "conv_in",
+    3: "blocks.0.mem_blocks.0",
+    4: "blocks.0.mem_blocks.1",
+    5: "blocks.0.mem_blocks.2",
+    7: "blocks.0.time_grow",
+    8: "blocks.0.conv_out",
+    9: "blocks.1.mem_blocks.0",
+    10: "blocks.1.mem_blocks.1",
+    11: "blocks.1.mem_blocks.2",
+    13: "blocks.1.time_grow",
+    14: "blocks.1.conv_out",
+    15: "blocks.2.mem_blocks.0",
+    16: "blocks.2.mem_blocks.1",
+    17: "blocks.2.mem_blocks.2",
+    19: "blocks.2.time_grow",
+    20: "blocks.2.conv_out",
+    22: "conv_out",
+}
+# inside a MemBlock / TPool / TGrow
+PARAM_RENAMES = {".conv.0.": ".conv1.", ".conv.2.": ".conv2.", ".conv.4.": ".conv3.", ".conv.": "."}
+
+
+def convert_taehv_state_dict(state_dict: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
+    converted = {}
+    for key, value in state_dict.items():
+        part, index, rest = key.split(".", 2)
+        layers = ENCODER_LAYERS if part == "encoder" else DECODER_LAYERS
+        new_key = f"{part}.{layers[int(index)]}.{rest}"
+        for old, new in PARAM_RENAMES.items():
+            if old in new_key:
+                new_key = new_key.replace(old, new)
+                break
+        converted[new_key] = value
+    return converted
+
+
+def convert_taehv(checkpoint_path: str, variant: str) -> AutoencoderTinyVideo:
+    state_dict = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
+    model = AutoencoderTinyVideo(**VARIANTS[variant])
+    model.load_state_dict(convert_taehv_state_dict(state_dict), strict=True)
+    return model
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--checkpoint_path", type=str, required=True)
+    parser.add_argument("--variant", type=str, choices=sorted(VARIANTS), required=True)
+    parser.add_argument("--output_path", type=str, required=True)
+    args = parser.parse_args()
+
+    model = convert_taehv(args.checkpoint_path, args.variant)
+    model.save_pretrained(args.output_path)
+    AutoencoderTinyVideo.from_pretrained(args.output_path)
+    print(f"saved to {args.output_path}")

src/diffusers/__init__.py

@@ -224,6 +224,7 @@
     ]
     _import_structure["models"].extend(
         [
+            "ABotWorldTransformer3DModel",
             "AceStepTransformer1DModel",
             "AllegroTransformer3DModel",
             "AnimaTextConditioner",
@@ -258,6 +259,7 @@
             "AutoencoderRAE",
             "AutoencoderSAME",
             "AutoencoderTiny",
+            "AutoencoderTinyVideo",
             "AutoencoderVidTok",
             "AutoModel",
             "BriaFiboTransformer2DModel",
@@ -513,6 +515,8 @@
 else:
     _import_structure["modular_pipelines"].extend(
         [
+            "ABotWorldBlocks",
+            "ABotWorldModularPipeline",
             "AnimaAutoBlocks",
             "AnimaModularPipeline",
             "Cosmos3DistilledBlocks",
@@ -1098,6 +1102,7 @@
             VaeImageProcessorLDM3D,
         )
         from .models import (
+            ABotWorldTransformer3DModel,
             AceStepTransformer1DModel,
             AllegroTransformer3DModel,
             AnimaTextConditioner,
@@ -1132,6 +1137,7 @@
             AutoencoderRAE,
             AutoencoderSAME,
             AutoencoderTiny,
+            AutoencoderTinyVideo,
             AutoencoderVidTok,
             AutoModel,
             BriaFiboTransformer2DModel,
@@ -1366,6 +1372,8 @@
         from .utils.dummy_torch_and_transformers_objects import *  # noqa F403
     else:
         from .modular_pipelines import (
+            ABotWorldBlocks,
+            ABotWorldModularPipeline,
             AnimaAutoBlocks,
             AnimaModularPipeline,
             Cosmos3DistilledBlocks,

src/diffusers/models/__init__.py

@@ -57,6 +57,7 @@
     _import_structure["autoencoders.autoencoder_rae"] = ["AutoencoderRAE"]
     _import_structure["autoencoders.autoencoder_same"] = ["AutoencoderSAME"]
     _import_structure["autoencoders.autoencoder_tiny"] = ["AutoencoderTiny"]
+    _import_structure["autoencoders.autoencoder_tiny_video"] = ["AutoencoderTinyVideo"]
     _import_structure["autoencoders.autoencoder_vidtok"] = ["AutoencoderVidTok"]
     _import_structure["autoencoders.consistency_decoder_vae"] = ["ConsistencyDecoderVAE"]
     _import_structure["autoencoders.ltx2_diffusion_decoder"] = ["LTX2VideoDiffusionDecoderModel"]
@@ -103,6 +104,7 @@
     _import_structure["transformers.t5_film_transformer"] = ["T5FilmDecoder"]
     _import_structure["transformers.transformer_2d"] = ["Transformer2DModel"]
     _import_structure["transformers.transformer_2d_dreamlite"] = ["DreamLiteTransformer2DModel"]
+    _import_structure["transformers.transformer_abot_world"] = ["ABotWorldTransformer3DModel"]
     _import_structure["transformers.transformer_allegro"] = ["AllegroTransformer3DModel"]
     _import_structure["transformers.transformer_anyflow"] = ["AnyFlowTransformer3DModel"]
     _import_structure["transformers.transformer_anyflow_far"] = ["AnyFlowFARTransformer3DModel"]
@@ -201,6 +203,7 @@
             AutoencoderRAE,
             AutoencoderSAME,
             AutoencoderTiny,
+            AutoencoderTinyVideo,
             AutoencoderVidTok,
             ConsistencyDecoderVAE,
             Cosmos3AVAEAudioTokenizer,
@@ -234,6 +237,7 @@
         from .embeddings import ImageProjection
         from .modeling_utils import ModelMixin
         from .transformers import (
+            ABotWorldTransformer3DModel,
             AceStepTransformer1DModel,
             AllegroTransformer3DModel,
             AnyFlowFARTransformer3DModel,

src/diffusers/models/autoencoders/__init__.py

@@ -27,6 +27,7 @@
 from .autoencoder_rae import AutoencoderRAE
 from .autoencoder_same import AutoencoderSAME
 from .autoencoder_tiny import AutoencoderTiny
+from .autoencoder_tiny_video import AutoencoderTinyVideo
 from .autoencoder_vidtok import AutoencoderVidTok
 from .consistency_decoder_vae import ConsistencyDecoderVAE
 from .ltx2_diffusion_decoder import LTX2VideoDiffusionDecoderModel

src/diffusers/models/autoencoders/autoencoder_kl_wan.py

@@ -31,6 +31,19 @@
 CACHE_T = 2
 
 
+class WanDecodeCache:
+    """
+    Causal-convolution feature cache for decoding a video chunk by chunk with [`AutoencoderKLWan.decode`].
+
+    Pass the same cache to consecutive `decode(z, cache=cache)` calls: each call decodes only the latent frames in `z`,
+    continuing from the frames decoded by the previous calls, and the concatenated result is identical to decoding all
+    frames in a single call. Create a new cache for every new video.
+    """
+
+    def __init__(self):
+        self.feat_map: list | None = None
+
+
 class AvgDown3D(nn.Module):
     def __init__(
         self,
@@ -1184,24 +1197,34 @@ def encode(
             return (posterior,)
         return AutoencoderKLOutput(latent_dist=posterior)
 
-    def _decode(self, z: torch.Tensor, return_dict: bool = True):
+    def _decode(self, z: torch.Tensor, return_dict: bool = True, cache: WanDecodeCache | None = None):
         _, _, num_frame, height, width = z.shape
         tile_latent_min_height = self.tile_sample_min_height // self.spatial_compression_ratio
         tile_latent_min_width = self.tile_sample_min_width // self.spatial_compression_ratio
 
         if self.use_tiling and (width > tile_latent_min_width or height > tile_latent_min_height):
             return self.tiled_decode(z, return_dict=return_dict)
 
-        self.clear_cache()
+        if cache is None:
+            self.clear_cache()
+            feat_map = self._feat_map
+            first_chunk = True
+        else:
+            # a fresh cache starts a new video; a used one continues the previous call's video
+            if cache.feat_map is None:
+                cache.feat_map = [None] * self._cached_conv_counts["decoder"]
+            feat_map = cache.feat_map
+            first_chunk = feat_map[0] is None
+
         x = self.post_quant_conv(z)
         for i in range(num_frame):
-            self._conv_idx = [0]
+            conv_idx = [0]
             if i == 0:
                 out = self.decoder(
-                    x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx, first_chunk=True
+                    x[:, :, i : i + 1, :, :], feat_cache=feat_map, feat_idx=conv_idx, first_chunk=first_chunk
                 )
             else:
-                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=self._feat_map, feat_idx=self._conv_idx)
+                out_ = self.decoder(x[:, :, i : i + 1, :, :], feat_cache=feat_map, feat_idx=conv_idx)
                 out = torch.cat([out, out_], 2)
 
         if self.config.patch_size is not None:
@@ -1216,25 +1239,32 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True):
         return DecoderOutput(sample=out)
 
     @apply_forward_hook
-    def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | torch.Tensor:
+    def decode(
+        self, z: torch.Tensor, return_dict: bool = True, cache: WanDecodeCache | None = None
+    ) -> DecoderOutput | torch.Tensor:
         r"""
         Decode a batch of images.
 
         Args:
             z (`torch.Tensor`): Input batch of latent vectors.
             return_dict (`bool`, *optional*, defaults to `True`):
                 Whether to return a [`~models.vae.DecoderOutput`] instead of a plain tuple.
+            cache (`WanDecodeCache`, *optional*):
+                Decode a video chunk by chunk: pass the same cache to consecutive calls and each call decodes only the
+                frames in `z`, continuing from the previous calls. Not supported together with slicing or tiling.
 
         Returns:
             [`~models.vae.DecoderOutput`] or `tuple`:
                 If return_dict is True, a [`~models.vae.DecoderOutput`] is returned, otherwise a plain `tuple` is
                 returned.
         """
+        if cache is not None and (self.use_slicing or self.use_tiling):
+            raise ValueError("Decoding with a `cache` does not support slicing or tiling.")
         if self.use_slicing and z.shape[0] > 1:
             decoded_slices = [self._decode(z_slice).sample for z_slice in z.split(1)]
             decoded = torch.cat(decoded_slices)
         else:
-            decoded = self._decode(z).sample
+            decoded = self._decode(z, cache=cache).sample
 
         if not return_dict:
             return (decoded,)