diff --git a/docs/source/en/_toctree.yml b/docs/source/en/_toctree.yml index 6c6e3a8e7882..2646d77c6c71 100644 --- a/docs/source/en/_toctree.yml +++ b/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 diff --git a/docs/source/en/api/models/autoencoder_tiny_video.md b/docs/source/en/api/models/autoencoder_tiny_video.md new file mode 100644 index 000000000000..67ed4bc445d8 --- /dev/null +++ b/docs/source/en/api/models/autoencoder_tiny_video.md @@ -0,0 +1,37 @@ + + +# 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 diff --git a/scripts/convert_abot_world_to_diffusers.py b/scripts/convert_abot_world_to_diffusers.py new file mode 100644 index 000000000000..ae64c48761eb --- /dev/null +++ b/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 /diffusion_pytorch_model.safetensors --output_path [--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() diff --git a/scripts/convert_taehv_to_diffusers.py b/scripts/convert_taehv_to_diffusers.py new file mode 100644 index 000000000000..e141fd355362 --- /dev/null +++ b/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}") diff --git a/src/diffusers/__init__.py b/src/diffusers/__init__.py index 88ede2fd009a..a39267d1e618 100644 --- a/src/diffusers/__init__.py +++ b/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, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index 8ba17d896434..583ea01c5ca4 100755 --- a/src/diffusers/models/__init__.py +++ b/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, diff --git a/src/diffusers/models/autoencoders/__init__.py b/src/diffusers/models/autoencoders/__init__.py index 607704343743..05e20b2d27d8 100644 --- a/src/diffusers/models/autoencoders/__init__.py +++ b/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 diff --git a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py b/src/diffusers/models/autoencoders/autoencoder_kl_wan.py index de8a56edc20e..a0cef4b2d2bf 100644 --- a/src/diffusers/models/autoencoders/autoencoder_kl_wan.py +++ b/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,7 +1197,7 @@ 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 @@ -1192,16 +1205,26 @@ def _decode(self, z: torch.Tensor, return_dict: bool = True): 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,7 +1239,9 @@ 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. @@ -1224,17 +1249,22 @@ def decode(self, z: torch.Tensor, return_dict: bool = True) -> DecoderOutput | t 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,) diff --git a/src/diffusers/models/autoencoders/autoencoder_tiny_video.py b/src/diffusers/models/autoencoders/autoencoder_tiny_video.py new file mode 100644 index 000000000000..0126d4b7fee7 --- /dev/null +++ b/src/diffusers/models/autoencoders/autoencoder_tiny_video.py @@ -0,0 +1,285 @@ +# Copyright 2026 Ollin Boer Bohan and 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. + +import torch +import torch.nn as nn +import torch.nn.functional as F + +from ...configuration_utils import ConfigMixin, register_to_config +from ...utils.accelerate_utils import apply_forward_hook +from ..modeling_utils import ModelMixin +from .autoencoder_tiny import AutoencoderTinyOutput +from .vae import AutoencoderMixin, DecoderOutput + + +class TinyVideoDecodeCache: + """ + Per-layer memory for decoding a video chunk by chunk with [`AutoencoderTinyVideo.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.memory: list[torch.Tensor] | None = None + + +class TinyVideoMemBlock(nn.Module): + """ + Residual block that sees the previous frame: its input is concatenated with the block input of the frame before it + (zeros for the very first frame of a video). Frames are stacked along the batch dimension. + """ + + def __init__(self, channels: int): + super().__init__() + self.conv1 = nn.Conv2d(channels * 2, channels, 3, padding=1) + self.conv2 = nn.Conv2d(channels, channels, 3, padding=1) + self.conv3 = nn.Conv2d(channels, channels, 3, padding=1) + self.act = nn.ReLU() + + def forward( + self, hidden_states: torch.Tensor, batch_size: int, last_frame: torch.Tensor | None = None + ) -> tuple[torch.Tensor, torch.Tensor]: + # memory of frame t is the block input of frame t-1; `last_frame` carries it over from the previous chunk + frames = hidden_states.unflatten(0, (batch_size, -1)) + if last_frame is None: + last_frame = torch.zeros_like(frames[:, :1]) + memory = torch.cat([last_frame, frames[:, :-1]], dim=1).flatten(0, 1) + + residual = hidden_states + hidden_states = self.act(self.conv1(torch.cat([hidden_states, memory], dim=1))) + hidden_states = self.act(self.conv2(hidden_states)) + hidden_states = self.conv3(hidden_states) + return self.act(hidden_states + residual), frames[:, -1:] + + +class TinyVideoEncoderBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, time_downscale: bool): + super().__init__() + self.time_stride = 2 if time_downscale else 1 + self.time_pool = nn.Conv2d(in_channels * self.time_stride, in_channels, 1, bias=False) + self.conv_down = nn.Conv2d(in_channels, out_channels, 3, stride=2, padding=1, bias=False) + self.mem_blocks = nn.ModuleList([TinyVideoMemBlock(out_channels) for _ in range(3)]) + + def forward(self, hidden_states: torch.Tensor, batch_size: int) -> torch.Tensor: + # fold `time_stride` consecutive frames into the channels, then pool them with a 1x1 conv + hidden_states = hidden_states.unflatten(0, (-1, self.time_stride)).flatten(1, 2) + hidden_states = self.time_pool(hidden_states) + hidden_states = self.conv_down(hidden_states) + for mem_block in self.mem_blocks: + hidden_states, _ = mem_block(hidden_states, batch_size) + return hidden_states + + +class TinyVideoDecoderBlock(nn.Module): + def __init__(self, in_channels: int, out_channels: int, time_upscale: bool): + super().__init__() + self.mem_blocks = nn.ModuleList([TinyVideoMemBlock(in_channels) for _ in range(3)]) + self.upsample = nn.Upsample(scale_factor=2) + self.time_stride = 2 if time_upscale else 1 + self.time_grow = nn.Conv2d(in_channels, in_channels * self.time_stride, 1, bias=False) + self.conv_out = nn.Conv2d(in_channels, out_channels, 3, padding=1, bias=False) + + def forward( + self, hidden_states: torch.Tensor, batch_size: int, memory: list[torch.Tensor | None] + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + new_memory = [] + for mem_block, last_frame in zip(self.mem_blocks, memory): + hidden_states, last_frame = mem_block(hidden_states, batch_size, last_frame) + new_memory.append(last_frame) + hidden_states = self.upsample(hidden_states) + # grow every frame into `time_stride` consecutive frames with a 1x1 conv + hidden_states = self.time_grow(hidden_states) + hidden_states = hidden_states.unflatten(1, (self.time_stride, -1)).flatten(0, 1) + hidden_states = self.conv_out(hidden_states) + return hidden_states, new_memory + + +class TinyVideoEncoder(nn.Module): + def __init__( + self, + in_channels: int, + latent_channels: int, + block_out_channels: tuple[int, ...], + time_downscale: tuple[bool, ...], + ): + super().__init__() + self.conv_in = nn.Conv2d(in_channels, block_out_channels[0], 3, padding=1) + self.act = nn.ReLU() + self.blocks = nn.ModuleList( + [ + TinyVideoEncoderBlock(block_out_channels[max(i - 1, 0)], block_out_channels[i], time_downscale[i]) + for i in range(len(block_out_channels)) + ] + ) + self.conv_out = nn.Conv2d(block_out_channels[-1], latent_channels, 3, padding=1) + + def forward(self, hidden_states: torch.Tensor, batch_size: int) -> torch.Tensor: + hidden_states = self.act(self.conv_in(hidden_states)) + for block in self.blocks: + hidden_states = block(hidden_states, batch_size) + return self.conv_out(hidden_states) + + +class TinyVideoDecoder(nn.Module): + def __init__( + self, + latent_channels: int, + out_channels: int, + block_out_channels: tuple[int, ...], + time_upscale: tuple[bool, ...], + ): + super().__init__() + self.conv_in = nn.Conv2d(latent_channels, block_out_channels[0], 3, padding=1) + self.act = nn.ReLU() + self.blocks = nn.ModuleList( + [ + TinyVideoDecoderBlock(block_out_channels[i], block_out_channels[i + 1], time_upscale[i]) + for i in range(len(block_out_channels) - 1) + ] + ) + self.conv_out = nn.Conv2d(block_out_channels[-1], out_channels, 3, padding=1) + + def forward( + self, hidden_states: torch.Tensor, batch_size: int, memory: list[torch.Tensor | None] + ) -> tuple[torch.Tensor, list[torch.Tensor]]: + hidden_states = torch.tanh(hidden_states / 3) * 3 + hidden_states = self.act(self.conv_in(hidden_states)) + new_memory = [] + for i, block in enumerate(self.blocks): + hidden_states, block_memory = block(hidden_states, batch_size, memory[3 * i : 3 * i + 3]) + new_memory.extend(block_memory) + hidden_states = self.act(hidden_states) + return self.conv_out(hidden_states), new_memory + + +class AutoencoderTinyVideo(ModelMixin, AutoencoderMixin, ConfigMixin): + r""" + A tiny causal video autoencoder (TAEHV, [madebyollin/taehv](https://github.com/madebyollin/taehv)) that encodes to + and decodes from the latent space of a full video VAE — e.g. `taew2_2` for the Wan 2.2 VAE — orders of magnitude + faster than the full model, for previews and real-time decoding. Latents are the *normalized* (roughly unit + Gaussian) latents of the full VAE. + + This model inherits from [`ModelMixin`]. Check the superclass documentation for its generic methods implemented for + all models (such as downloading or saving). + + Parameters: + in_channels (`int`, defaults to `3`): Number of channels in the input video. + latent_channels (`int`, defaults to `48`): Number of channels in the latent space. + patch_size (`int`, defaults to `2`): + Pixel-(un)shuffle factor applied to the frames before the encoder and after the decoder. + encoder_block_out_channels (`tuple[int, ...]`, defaults to `(64, 64, 64)`): + Output channels of the encoder blocks; each block halves the spatial size. + decoder_block_out_channels (`tuple[int, ...]`, defaults to `(256, 128, 64, 64)`): + Channels of the decoder: the first entry is the width after the input conv, each following block doubles + the spatial size and outputs the next entry. + encoder_time_downscale (`tuple[bool, ...]`, defaults to `(True, True, False)`): + Whether each encoder block halves the number of frames. + decoder_time_upscale (`tuple[bool, ...]`, defaults to `(False, True, True)`): + Whether each decoder block doubles the number of frames. + """ + + _skip_keys = ["memory"] + + @register_to_config + def __init__( + self, + in_channels: int = 3, + latent_channels: int = 48, + patch_size: int = 2, + encoder_block_out_channels: tuple[int, ...] = (64, 64, 64), + decoder_block_out_channels: tuple[int, ...] = (256, 128, 64, 64), + encoder_time_downscale: tuple[bool, ...] = (True, True, False), + decoder_time_upscale: tuple[bool, ...] = (False, True, True), + ): + super().__init__() + self.encoder = TinyVideoEncoder( + in_channels * patch_size**2, latent_channels, encoder_block_out_channels, encoder_time_downscale + ) + self.decoder = TinyVideoDecoder( + latent_channels, in_channels * patch_size**2, decoder_block_out_channels, decoder_time_upscale + ) + self.temporal_compression_ratio = 2 ** sum(encoder_time_downscale) + self.temporal_upsampling_ratio = 2 ** sum(decoder_time_upscale) + self.spatial_compression_ratio = patch_size * 2 ** len(encoder_block_out_channels) + + @apply_forward_hook + def encode(self, x: torch.Tensor, return_dict: bool = True) -> AutoencoderTinyOutput | tuple[torch.Tensor]: + r""" + Encode a batch of videos `[B, C, T, H, W]` in `[-1, 1]`. The frames are padded at the end, by repeating the + last one, to a multiple of `temporal_compression_ratio`. + """ + batch_size, _, num_frames = x.shape[:3] + if num_frames % self.temporal_compression_ratio != 0: + num_pad = self.temporal_compression_ratio - num_frames % self.temporal_compression_ratio + x = torch.cat([x, x[:, :, -1:].repeat_interleave(num_pad, dim=2)], dim=2) + + frames = x.permute(0, 2, 1, 3, 4).flatten(0, 1) + frames = F.pixel_unshuffle(frames.add(1).div(2), self.config.patch_size) + latents = self.encoder(frames, batch_size) + latents = latents.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + + if not return_dict: + return (latents,) + return AutoencoderTinyOutput(latents=latents) + + @apply_forward_hook + def decode( + self, z: torch.Tensor, return_dict: bool = True, cache: TinyVideoDecodeCache | None = None + ) -> DecoderOutput | tuple[torch.Tensor]: + r""" + Decode a batch of latents `[B, C, T, h, w]` to videos in `[-1, 1]`. `T` latent frames decode to `T * + temporal_upsampling_ratio - (temporal_upsampling_ratio - 1)` frames: the first frames produced by the decoder + are warm-up frames and are dropped. + + 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 (`TinyVideoDecodeCache`, *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. + """ + batch_size = z.shape[0] + # a fresh cache (or no cache) starts a new video; a used one continues the previous call's video + first_chunk = cache is None or cache.memory is None + memory = cache.memory if not first_chunk else [None] * (3 * len(self.decoder.blocks)) + + latents = z.permute(0, 2, 1, 3, 4).flatten(0, 1) + frames, memory = self.decoder(latents, batch_size, memory) + if cache is not None: + cache.memory = memory + + frames = F.pixel_shuffle(frames, self.config.patch_size).clamp(0, 1).mul(2).sub(1) + frames = frames.unflatten(0, (batch_size, -1)).permute(0, 2, 1, 3, 4) + if first_chunk: + frames = frames[:, :, self.temporal_upsampling_ratio - 1 :] + + if not return_dict: + return (frames,) + return DecoderOutput(sample=frames) + + def forward(self, sample: torch.Tensor, return_dict: bool = True) -> DecoderOutput | tuple[torch.Tensor]: + r""" + Args: + sample (`torch.Tensor`): Input video `[B, C, T, H, W]` in `[-1, 1]`. + return_dict (`bool`, *optional*, defaults to `True`): + Whether or not to return a [`DecoderOutput`] instead of a plain tuple. + """ + latents = self.encode(sample).latents + decoded = self.decode(latents).sample + if not return_dict: + return (decoded,) + return DecoderOutput(sample=decoded) diff --git a/src/diffusers/models/transformers/__init__.py b/src/diffusers/models/transformers/__init__.py index 0e167812ad88..29d4170e3a93 100755 --- a/src/diffusers/models/transformers/__init__.py +++ b/src/diffusers/models/transformers/__init__.py @@ -19,6 +19,7 @@ from .t5_film_transformer import T5FilmDecoder from .transformer_2d import Transformer2DModel from .transformer_2d_dreamlite import DreamLiteTransformer2DModel + from .transformer_abot_world import ABotWorldTransformer3DModel from .transformer_allegro import AllegroTransformer3DModel from .transformer_anyflow import AnyFlowTransformer3DModel from .transformer_anyflow_far import AnyFlowFARTransformer3DModel diff --git a/src/diffusers/models/transformers/transformer_abot_world.py b/src/diffusers/models/transformers/transformer_abot_world.py new file mode 100644 index 000000000000..078d8f62dfaa --- /dev/null +++ b/src/diffusers/models/transformers/transformer_abot_world.py @@ -0,0 +1,699 @@ +# 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. + +import math + +import torch +import torch.nn as nn + +from ...configuration_utils import ConfigMixin, register_to_config +from ...loaders import PeftAdapterMixin +from ..attention import AttentionMixin, AttentionModuleMixin +from ..attention_dispatch import dispatch_attention_fn +from ..modeling_outputs import Transformer2DModelOutput +from ..modeling_utils import ModelMixin +from ..normalization import FP32LayerNorm + + +def rope_params(max_seq_len, dim, theta=10000): + freqs = torch.outer( + torch.arange(max_seq_len), + 1.0 / torch.pow(theta, torch.arange(0, dim, 2).to(torch.float64).div(dim)), + ) + return torch.polar(torch.ones_like(freqs), freqs) + + +def rope_apply(x, grid_sizes, freqs, start_frame=0): + """Apply 3D rotary embeddings with the temporal band offset by `start_frame`. + + Computes in float64/complex128 and returns float32, matching the reference `causal_rope_apply`. + """ + num_heads, c = x.size(2), x.size(3) // 2 + freqs = freqs.split([c - 2 * (c // 3), c // 3, c // 3], dim=1) + + output = [] + for i, (f, h, w) in enumerate(grid_sizes.tolist()): + seq_len = f * h * w + x_i = torch.view_as_complex(x[i, :seq_len].to(torch.float64).reshape(seq_len, num_heads, -1, 2)) + freqs_i = torch.cat( + [ + freqs[0][start_frame : start_frame + f].view(f, 1, 1, -1).expand(f, h, w, -1), + freqs[1][:h].view(1, h, 1, -1).expand(f, h, w, -1), + freqs[2][:w].view(1, 1, w, -1).expand(f, h, w, -1), + ], + dim=-1, + ).reshape(seq_len, 1, -1) + x_i = torch.view_as_real(x_i * freqs_i).flatten(2) + x_i = torch.cat([x_i, x[i, seq_len:]]) + output.append(x_i) + return torch.stack(output).float() + + +def reference_rope_freqs(freqs, num_slots, tokens_per_slot, ref_grid, device): + """Rotary frequencies for the reference-image tokens. + + Reference slots sit at large *negative* temporal positions (one stride of `max(tokens_per_slot, 256)` frames per + slot) so they never collide with the rolling video window's temporal ids. + """ + patch_t, patch_h, patch_w = ref_grid + freq_dim = freqs.shape[1] + f_band = freq_dim - 2 * (freq_dim // 3) + + temporal_step = max(tokens_per_slot, 256) + neg_temporal = torch.tensor( + [-(num_slots - i) * temporal_step for i in range(num_slots)], dtype=torch.float64, device=device + ) + t_freqs = torch.outer( + neg_temporal, + 1.0 / torch.pow(10000, torch.arange(0, 2 * f_band, 2, device=device, dtype=torch.float64).div(2 * f_band)), + ) + t_freqs = torch.polar(torch.ones_like(t_freqs), t_freqs) + + freqs_split = freqs.split([f_band, freq_dim // 3, freq_dim // 3], dim=1) + h_freqs = freqs_split[1][:patch_h].to(device) + w_freqs = freqs_split[2][:patch_w].to(device) + ref_freqs = torch.cat( + [ + t_freqs[:, None, None, None, :].expand(num_slots, patch_t, patch_h, patch_w, f_band), + h_freqs[None, None, :, None, :].expand(num_slots, patch_t, patch_h, patch_w, freq_dim // 3), + w_freqs[None, None, None, :, :].expand(num_slots, patch_t, patch_h, patch_w, freq_dim // 3), + ], + dim=-1, + ).reshape(num_slots * tokens_per_slot, 1, -1) + return ref_freqs.to(torch.complex64) + + +def reference_rope_apply(x, freqs): + """Apply the reference-slot rotary embeddings (float32/complex64, matching `rope_apply_with_refimg`).""" + x_out = torch.view_as_complex(x.to(torch.float32).reshape(x.shape[0], x.shape[1], x.shape[2], -1, 2)) + x_out = torch.view_as_real(x_out * freqs.to(x.device)).flatten(3) + return x_out.to(x.dtype) + + +class ABotWorldLayerKVCache: + """Rolling K/V cache for one self-attention layer plus the layer's one-shot cross-attention cache. + + `key_raw` holds pre-RoPE keys (reference-token slots first, then the rolling video window), `key_roped` the + post-RoPE video keys used for attention (re-based periodically so temporal positions stay within the rotary table), + `value` the values. Tensor format: `(batch_size, num_tokens, num_heads, head_dim)`. + """ + + def __init__(self, batch_size, num_tokens, num_heads, head_dim, device, dtype): + self.key_raw = torch.zeros(batch_size, num_tokens, num_heads, head_dim, device=device, dtype=dtype) + self.key_roped = torch.zeros_like(self.key_raw) + self.value = torch.zeros_like(self.key_raw) + self.global_end_index = 0 + self.local_end_index = 0 + self.rope_base_frame = 0 + self.cross_key: torch.Tensor | None = None + self.cross_value: torch.Tensor | None = None + + def reset(self): + self.global_end_index = 0 + self.local_end_index = 0 + self.rope_base_frame = 0 + self.cross_key = None + self.cross_value = None + + +class ABotWorldKVCache: + """Container holding one [`ABotWorldLayerKVCache`] per transformer layer. + + Args: + num_layers: Number of transformer layers. + batch_size: Batch size of the rollout. + num_tokens: Cache length in tokens: `ref_token_len + local_attn_size * tokens_per_frame`. + ref_token_len: Number of reference-image tokens pinned at the start of the cache (never evicted). + num_heads / head_dim / device / dtype: K/V tensor layout. + """ + + def __init__(self, num_layers, batch_size, num_tokens, ref_token_len, num_heads, head_dim, device, dtype): + self.layer_caches = [ + ABotWorldLayerKVCache(batch_size, num_tokens, num_heads, head_dim, device, dtype) + for _ in range(num_layers) + ] + self.ref_token_len = ref_token_len + + def get(self, layer_idx: int) -> ABotWorldLayerKVCache: + return self.layer_caches[layer_idx] + + def reset(self): + for cache in self.layer_caches: + cache.reset() + + +class ABotWorldSelfAttnProcessor: + r""" + Causal windowed self-attention over a rolling K/V cache. + + Each forward writes the new block's keys/values into the cache (evicting the oldest video tokens once the + `local_attn_size`-frame window is full; reference tokens at the head of the cache are never evicted) and attends + the new block's queries over `[reference tokens | visible video window]`. Temporal rotary positions use an absolute + counter relative to `rope_base_frame`, re-based whenever positions approach the rotary table limit — the attention + logits only depend on position differences within the window, so re-basing does not change them. + + On the first block of a stream (`current_start == 0`) the reference tokens ride along in `hidden_states` (and in + the queries); their pre-RoPE keys/values are pinned into the cache prefix. + """ + + _attention_backend = None + _parallel_config = None + + # keep roped temporal positions well below the rotary-table length before re-basing + _REBASE_MAX_POS = 256 + + def __call__( + self, + attn: "ABotWorldAttention", + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor, + grid_sizes: torch.Tensor, + kv_cache: ABotWorldLayerKVCache, + current_start: int, + query_ref_token_len: int, + ref_token_len: int, + ref_rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + query = attn.norm_q(attn.to_q(hidden_states)) + key = attn.norm_k(attn.to_k(hidden_states)) + value = attn.to_v(hidden_states) + + query = query.unflatten(2, (attn.heads, -1)) + key = key.unflatten(2, (attn.heads, -1)) + value = value.unflatten(2, (attn.heads, -1)) + + frame_seqlen = int(math.prod(grid_sizes[0][1:]).item()) + video_token_len = query.shape[1] - query_ref_token_len + num_video_frames = video_token_len // frame_seqlen + video_grid_sizes = grid_sizes.clone() + video_grid_sizes[:, 0] = num_video_frames + + cache_size = kv_cache.key_raw.shape[1] + cache_current_end = ref_token_len + current_start + video_token_len + + # a new stream reuses the cache tensors; reset the rope base with the indices + if kv_cache.global_end_index == 0: + kv_cache.rope_base_frame = 0 + + # roll the window: evict the oldest video tokens (reference tokens are pinned at the head) + if cache_current_end > kv_cache.global_end_index and video_token_len + kv_cache.local_end_index > cache_size: + num_evicted = video_token_len + kv_cache.local_end_index - cache_size + num_rolled = kv_cache.local_end_index - num_evicted - ref_token_len + src = slice(ref_token_len + num_evicted, ref_token_len + num_evicted + num_rolled) + dst = slice(ref_token_len, ref_token_len + num_rolled) + kv_cache.key_raw[:, dst] = kv_cache.key_raw[:, src].clone() + kv_cache.key_roped[:, dst] = kv_cache.key_roped[:, src].clone() + kv_cache.value[:, dst] = kv_cache.value[:, src].clone() + local_end_index = kv_cache.local_end_index + cache_current_end - kv_cache.global_end_index - num_evicted + else: + local_end_index = kv_cache.local_end_index + cache_current_end - kv_cache.global_end_index + local_start_index = local_end_index - video_token_len + + if query_ref_token_len > 0: + kv_cache.key_raw[:, :ref_token_len] = key[:, :query_ref_token_len] + kv_cache.value[:, :ref_token_len] = value[:, :query_ref_token_len] + kv_cache.key_raw[:, local_start_index:local_end_index] = key[:, query_ref_token_len:] + kv_cache.value[:, local_start_index:local_end_index] = value[:, query_ref_token_len:] + + # the visible video window, frame-aligned + max_attention_tokens = attn.local_attn_size * frame_seqlen + recent_start = max(ref_token_len, local_end_index - max_attention_tokens) + recent_start += (local_end_index - recent_start) % frame_seqlen + visible_video_frames = (local_end_index - recent_start) // frame_seqlen + + # temporal rotary positions: absolute counter relative to the rope base, re-based before it + # approaches the rotary table limit (logits depend only on position differences in the window) + abs_frame_start = current_start // frame_seqlen + new_start_pos = abs_frame_start - kv_cache.rope_base_frame + rebase_limit = min(self._REBASE_MAX_POS, rotary_emb.shape[0] - attn.local_attn_size - num_video_frames) + if visible_video_frames > 0 and (new_start_pos + num_video_frames > rebase_limit or new_start_pos < 0): + kv_cache.rope_base_frame = abs_frame_start + num_video_frames - visible_video_frames + window_grid = grid_sizes.clone() + window_grid[:, 0] = visible_video_frames + kv_cache.key_roped[:, recent_start:local_end_index] = rope_apply( + kv_cache.key_raw[:, recent_start:local_end_index], window_grid, rotary_emb, start_frame=0 + ).type_as(value) + new_start_pos = abs_frame_start - kv_cache.rope_base_frame + else: + kv_cache.key_roped[:, local_start_index:local_end_index] = rope_apply( + key[:, query_ref_token_len:], video_grid_sizes, rotary_emb, start_frame=new_start_pos + ).type_as(value) + + roped_query = rope_apply( + query[:, query_ref_token_len:], video_grid_sizes, rotary_emb, start_frame=new_start_pos + ).type_as(value) + + attn_key = kv_cache.key_roped[:, recent_start:local_end_index] + attn_value = kv_cache.value[:, recent_start:local_end_index] + if ref_token_len > 0: + ref_key = reference_rope_apply(kv_cache.key_raw[:, :ref_token_len], ref_rotary_emb).type_as(value) + attn_key = torch.cat([ref_key, attn_key], dim=1) + attn_value = torch.cat([kv_cache.value[:, :ref_token_len], attn_value], dim=1) + if query_ref_token_len > 0: + ref_query = reference_rope_apply(query[:, :query_ref_token_len], ref_rotary_emb).type_as(value) + roped_query = torch.cat([ref_query, roped_query], dim=1) + + hidden_states = dispatch_attention_fn( + roped_query, + attn_key, + attn_value, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + + kv_cache.global_end_index = cache_current_end + kv_cache.local_end_index = local_end_index + + hidden_states = attn.to_out[0](hidden_states.flatten(2, 3)) + return hidden_states + + +class ABotWorldCrossAttnProcessor: + r""" + Text cross-attention with per-stream K/V caching: the text keys/values are projected once on the first block of a + stream and reused for every subsequent block. + """ + + _attention_backend = None + _parallel_config = None + + def __call__( + self, + attn: "ABotWorldAttention", + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + kv_cache: ABotWorldLayerKVCache, + ) -> torch.Tensor: + query = attn.norm_q(attn.to_q(hidden_states)).unflatten(2, (attn.heads, -1)) + + if kv_cache.cross_key is None: + kv_cache.cross_key = attn.norm_k(attn.to_k(encoder_hidden_states)).unflatten(2, (attn.heads, -1)) + kv_cache.cross_value = attn.to_v(encoder_hidden_states).unflatten(2, (attn.heads, -1)) + + hidden_states = dispatch_attention_fn( + query, + kv_cache.cross_key, + kv_cache.cross_value, + attn_mask=None, + backend=self._attention_backend, + parallel_config=self._parallel_config, + ) + hidden_states = attn.to_out[0](hidden_states.flatten(2, 3)) + return hidden_states + + +class ABotWorldAttention(torch.nn.Module, AttentionModuleMixin): + _default_processor_cls = ABotWorldSelfAttnProcessor + _available_processors = [ABotWorldSelfAttnProcessor, ABotWorldCrossAttnProcessor] + + def __init__( + self, + dim: int, + heads: int, + eps: float, + local_attn_size: int | None = None, + is_cross_attention: bool = False, + processor=None, + ): + super().__init__() + self.heads = heads + self.local_attn_size = local_attn_size + self.is_cross_attention = is_cross_attention + + self.to_q = nn.Linear(dim, dim, bias=True) + self.to_k = nn.Linear(dim, dim, bias=True) + self.to_v = nn.Linear(dim, dim, bias=True) + self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=True)]) + self.norm_q = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + self.norm_k = nn.RMSNorm(dim, eps=eps, elementwise_affine=True) + + if processor is None: + processor = self._default_processor_cls() + self.set_processor(processor) + + def forward(self, hidden_states: torch.Tensor, **kwargs) -> torch.Tensor: + return self.processor(self, hidden_states, **kwargs) + + +class ABotWorldResidualBlock(nn.Module): + def __init__(self, dim: int): + super().__init__() + self.conv1 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + self.relu = nn.ReLU(inplace=True) + self.conv2 = nn.Conv2d(dim, dim, kernel_size=3, padding=1) + + def forward(self, x): + return x + self.conv2(self.relu(self.conv1(x))) + + +class ABotWorldActionAdapter(nn.Module): + """Encodes the broadcast action planes to the patch-token grid. + + Input `(B, action_in_channels, F, H_pix, W_pix)`; PixelUnshuffle + a stride-2 conv bring the spatial dims to the + latent patch grid (`H_pix / (downscale_factor * 2)`), producing `(B, dim, F, H_patch, W_patch)` — added directly + onto the patch-embedded video tokens. + """ + + def __init__(self, in_channels: int, dim: int, downscale_factor: int): + super().__init__() + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=downscale_factor) + self.conv = nn.Conv2d( + in_channels * downscale_factor * downscale_factor, dim, kernel_size=(2, 2), stride=(2, 2), padding=0 + ) + self.residual_blocks = nn.Sequential(ABotWorldResidualBlock(dim)) + + def forward(self, x): + batch_size, channels, num_frames, height, width = x.size() + x = x.permute(0, 2, 1, 3, 4).contiguous().view(batch_size * num_frames, channels, height, width) + x = self.residual_blocks(self.conv(self.pixel_unshuffle(x))) + x = x.view(batch_size, num_frames, x.size(1), x.size(2), x.size(3)) + return x.permute(0, 2, 1, 3, 4) + + +class ABotWorldTimeTextEmbedding(nn.Module): + def __init__(self, dim: int, time_freq_dim: int, text_embed_dim: int): + super().__init__() + self.time_freq_dim = time_freq_dim + self.time_embedder = nn.Sequential(nn.Linear(time_freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_proj = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + self.text_embedder = nn.Sequential( + nn.Linear(text_embed_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim) + ) + + def sinusoidal_embedding(self, timestep: torch.Tensor) -> torch.Tensor: + # matches the reference `sinusoidal_embedding_1d`: half-dim sin/cos over a 10000 theta, float64 + half = self.time_freq_dim // 2 + timestep = timestep.type(torch.float64) + sinusoid = torch.outer( + timestep, torch.pow(10000, -torch.arange(half, device=timestep.device).to(timestep.dtype).div(half)) + ) + return torch.cat([torch.cos(sinusoid), torch.sin(sinusoid)], dim=1) + + def forward(self, timestep: torch.Tensor, dtype: torch.dtype): + # timestep: any shape; returns (temb [N, dim], temb_proj [*timestep.shape, 6, dim]) + temb = self.time_embedder(self.sinusoidal_embedding(timestep.flatten()).to(dtype)) + temb_proj = self.time_proj(temb).unflatten(1, (6, -1)).unflatten(0, timestep.shape) + return temb, temb_proj + + +class ABotWorldTransformerBlock(nn.Module): + def __init__(self, dim: int, ffn_dim: int, num_heads: int, eps: float, local_attn_size: int): + super().__init__() + self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.attn1 = ABotWorldAttention(dim, num_heads, eps, local_attn_size=local_attn_size) + self.attn2 = ABotWorldAttention( + dim, num_heads, eps, is_cross_attention=True, processor=ABotWorldCrossAttnProcessor() + ) + self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb_proj: torch.Tensor, + rotary_emb: torch.Tensor, + grid_sizes: torch.Tensor, + kv_cache: ABotWorldLayerKVCache, + current_start: int, + query_ref_token_len: int, + ref_token_len: int, + ref_rotary_emb: torch.Tensor | None, + ) -> torch.Tensor: + # temb_proj is [B, F, 6, C] frame-level modulation, or [B, L, 6, C] token-level on the first + # block of a stream (where reference tokens with zero-timestep modulation ride along) + token_level = temb_proj.shape[1] == hidden_states.shape[1] + if token_level: + num_frames, frame_seqlen = hidden_states.shape[1], 1 + else: + num_frames, frame_seqlen = temb_proj.shape[1], hidden_states.shape[1] // temb_proj.shape[1] + + e = (self.scale_shift_table.unsqueeze(0).float() + temb_proj.float()).chunk(6, dim=2) + + def modulate(normed, shift, scale): + if token_level: + return normed * (1 + scale.squeeze(2)) + shift.squeeze(2) + return ((normed.unflatten(1, (num_frames, frame_seqlen)) * (1 + scale)) + shift).flatten(1, 2) + + def gate(value, g): + if token_level: + return value * g.squeeze(2) + return (value.unflatten(1, (num_frames, frame_seqlen)) * g).flatten(1, 2) + + norm_hidden_states = modulate(self.norm1(hidden_states.float()), e[0], e[1]).type_as(hidden_states) + attn_output = self.attn1( + norm_hidden_states, + rotary_emb=rotary_emb, + grid_sizes=grid_sizes, + kv_cache=kv_cache, + current_start=current_start, + query_ref_token_len=query_ref_token_len, + ref_token_len=ref_token_len, + ref_rotary_emb=ref_rotary_emb, + ) + hidden_states = (hidden_states.float() + gate(attn_output.float(), e[2])).type_as(hidden_states) + + attn_output = self.attn2( + self.norm2(hidden_states.float()).type_as(hidden_states), + encoder_hidden_states=encoder_hidden_states, + kv_cache=kv_cache, + ) + hidden_states = hidden_states + attn_output + + norm_hidden_states = modulate(self.norm3(hidden_states.float()), e[3], e[4]).type_as(hidden_states) + ffn_output = self.ffn(norm_hidden_states) + hidden_states = (hidden_states.float() + gate(ffn_output.float(), e[5])).type_as(hidden_states) + return hidden_states + + +class ABotWorldTransformer3DModel(ModelMixin, ConfigMixin, AttentionMixin, PeftAdapterMixin): + r""" + The causal, action-conditioned video transformer from [ABot-World](https://github.com/amap-cvlab/ABot-World), a + Wan2.2-TI2V-5B finetune for real-time interactive world generation. + + The model denoises one block of latent frames at a time: self-attention is windowed over the last `local_attn_size` + frames through a rolling K/V cache ([`ABotWorldKVCache`]), reference-image tokens are pinned at the head of the + cache, and keyboard-action planes are injected through a learned adapter added onto the patch tokens. Timesteps are + per latent frame (`(batch, frames)`). + """ + + _repeated_blocks = ["ABotWorldTransformerBlock"] + _no_split_modules = ["ABotWorldTransformerBlock"] + _skip_layerwise_casting_patterns = ["patch_embedding", "condition_embedder", "norm"] + _skip_keys = ["kv_cache"] + + @register_to_config + def __init__( + self, + patch_size: tuple[int] = (1, 2, 2), + num_attention_heads: int = 24, + attention_head_dim: int = 128, + in_channels: int = 48, + out_channels: int = 48, + text_dim: int = 4096, + text_len: int = 512, + freq_dim: int = 256, + ffn_dim: int = 14336, + num_layers: int = 30, + eps: float = 1e-6, + local_attn_size: int = 21, + action_in_channels: int = 32, + action_downscale_factor: int = 16, + rope_max_seq_len: int = 1024, + ): + super().__init__() + inner_dim = num_attention_heads * attention_head_dim + + self.patch_embedding = nn.Conv3d(in_channels, inner_dim, kernel_size=patch_size, stride=patch_size) + self.act_control_adapter = ABotWorldActionAdapter(action_in_channels, inner_dim, action_downscale_factor) + self.condition_embedder = ABotWorldTimeTextEmbedding(inner_dim, freq_dim, text_dim) + + self.blocks = nn.ModuleList( + [ + ABotWorldTransformerBlock(inner_dim, ffn_dim, num_attention_heads, eps, local_attn_size) + for _ in range(num_layers) + ] + ) + + self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) + self.scale_shift_table = nn.Parameter(torch.randn(1, 2, inner_dim) / inner_dim**0.5) + + # kept as a plain float64/complex128 attribute (not a buffer) so `model.to(dtype)` never downcasts it; + # moved to the execution device in `forward`, matching the reference + head_dim = attention_head_dim + self.rotary_freqs = torch.cat( + [ + rope_params(rope_max_seq_len, head_dim - 4 * (head_dim // 6)), + rope_params(rope_max_seq_len, 2 * (head_dim // 6)), + rope_params(rope_max_seq_len, 2 * (head_dim // 6)), + ], + dim=1, + ) + + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + action_hidden_states: torch.Tensor | None = None, + action_scale: float = 1.0, + reference_hidden_states: torch.Tensor | None = None, + reference_mask: torch.Tensor | None = None, + kv_cache: ABotWorldKVCache = None, + current_start: int = 0, + return_dict: bool = True, + ): + r""" + Args: + hidden_states: Noisy latents for one block, `(batch, in_channels, frames, height, width)`. + timestep: Per-latent-frame timesteps, `(batch, frames)`. + encoder_hidden_states: Text embeddings `(batch, seq_len, text_dim)`; zero-padded to `text_len` inside. + action_hidden_states: Broadcast action planes `(batch, action_in_channels, frames, pixel_h, pixel_w)`. + reference_hidden_states: Reference-image latents `(batch, num_slots, in_channels, 1, ref_h, ref_w)`. + Only consumed on the first block of a stream (`current_start == 0`), where the reference tokens are + pinned into the K/V cache; later blocks attend to them from the cache. + reference_mask: Per-slot validity mask `(batch, num_slots)`. + kv_cache: The stream's [`ABotWorldKVCache`], allocated with room for the reference tokens. + current_start: Token offset of this block in the rollout: `start_frame * tokens_per_frame`. + """ + if kv_cache is None: + raise ValueError("`kv_cache` is required: this model only runs block-causal rollout with a K/V cache.") + if self.rotary_freqs.device != hidden_states.device: + self.rotary_freqs = self.rotary_freqs.to(hidden_states.device) + freqs = self.rotary_freqs + + # patchify and add the action adapter's features onto the patch tokens + hidden_states = self.patch_embedding(hidden_states) + if action_hidden_states is not None: + action_features = self.act_control_adapter(action_hidden_states) + frames = hidden_states.shape[2] + action_frames = action_features.shape[2] + if frames > action_frames: + offset = frames - action_frames + hidden_states = torch.cat( + [hidden_states[:, :, :offset], hidden_states[:, :, offset:] + action_features * action_scale], + dim=2, + ) + else: + hidden_states = hidden_states + action_features * action_scale + + grid_sizes = torch.tensor(hidden_states.shape[2:], dtype=torch.long).unsqueeze(0) + batch_size = hidden_states.shape[0] + frame_seqlen = int(math.prod(hidden_states.shape[3:])) + hidden_states = hidden_states.flatten(2).transpose(1, 2) + + temb, temb_proj = self.condition_embedder(timestep, hidden_states.dtype) + + # zero-pad the text embeddings to text_len, matching the reference + seq = encoder_hidden_states.shape[1] + if seq < self.config.text_len: + encoder_hidden_states = torch.cat( + [ + encoder_hidden_states, + encoder_hidden_states.new_zeros( + batch_size, self.config.text_len - seq, encoder_hidden_states.shape[2] + ), + ], + dim=1, + ) + encoder_hidden_states = self.condition_embedder.text_embedder(encoder_hidden_states) + + # reference tokens ride along on the first block of a stream and are pinned into the cache + ref_token_len = kv_cache.ref_token_len + query_ref_token_len = 0 + ref_rotary_emb = None + if reference_hidden_states is not None and current_start == 0: + batch, num_slots, channels, ref_t, ref_h, ref_w = reference_hidden_states.shape + ref_features = self.patch_embedding( + reference_hidden_states.reshape(batch * num_slots, channels, ref_t, ref_h, ref_w).to( + hidden_states.dtype + ) + ) + patch_t, patch_h, patch_w = ref_features.shape[2:] + tokens_per_slot = patch_t * patch_h * patch_w + ref_tokens = ref_features.flatten(2).transpose(1, 2).reshape(batch, num_slots, tokens_per_slot, -1) + if reference_mask is None: + reference_mask = reference_hidden_states.new_ones(batch, num_slots) + ref_tokens = ref_tokens * reference_mask[:, :, None, None].to(ref_tokens.dtype) + ref_tokens = ref_tokens.reshape(batch, num_slots * tokens_per_slot, -1) + + query_ref_token_len = num_slots * tokens_per_slot + if query_ref_token_len != ref_token_len: + raise ValueError( + f"The KV cache was allocated for {ref_token_len} reference tokens but " + f"`reference_hidden_states` produced {query_ref_token_len}." + ) + self._ref_grid = (int(patch_t), int(patch_h), int(patch_w)) + self._ref_num_slots = int(num_slots) + self._ref_tokens_per_slot = int(tokens_per_slot) + + # token-level modulation: video tokens keep their frame's modulation, reference tokens get timestep 0 + temb_proj = temb_proj.repeat_interleave(frame_seqlen, dim=1) + ref_timestep = torch.zeros((batch_size, 1), dtype=torch.long, device=hidden_states.device) + _, ref_temb_proj = self.condition_embedder(ref_timestep, hidden_states.dtype) + temb_proj = torch.cat([ref_temb_proj.expand(-1, query_ref_token_len, -1, -1), temb_proj], dim=1) + hidden_states = torch.cat([ref_tokens, hidden_states], dim=1) + + if ref_token_len > 0: + ref_rotary_emb = reference_rope_freqs( + freqs, self._ref_num_slots, self._ref_tokens_per_slot, self._ref_grid, hidden_states.device + ) + + for layer_idx, block in enumerate(self.blocks): + hidden_states = block( + hidden_states, + encoder_hidden_states, + temb_proj, + freqs, + grid_sizes, + kv_cache.get(layer_idx), + current_start, + query_ref_token_len, + ref_token_len, + ref_rotary_emb, + ) + + if query_ref_token_len > 0: + hidden_states = hidden_states[:, query_ref_token_len:] + + # head: frame-level modulation with the un-projected time embedding + shift, scale = ( + self.scale_shift_table.unsqueeze(1).float() + temb.unflatten(0, timestep.shape).unsqueeze(2).float() + ).chunk(2, dim=2) + num_frames = timestep.shape[1] + hidden_states = ( + self.norm_out(hidden_states.float()).unflatten(1, (num_frames, frame_seqlen)) * (1 + scale) + shift + ).type_as(hidden_states) + hidden_states = self.proj_out(hidden_states) + + # unpatchify: (B, F, frame_seqlen, prod(patch) * C) -> (B, C, F*pt, H*ph, W*pw) + p_t, p_h, p_w = self.config.patch_size + _, latent_h, latent_w = grid_sizes[0].tolist() + hidden_states = hidden_states.reshape( + batch_size, num_frames, latent_h, latent_w, p_t, p_h, p_w, self.config.out_channels + ) + hidden_states = torch.einsum("bfhwpqrc->bcfphqwr", hidden_states) + hidden_states = hidden_states.reshape( + batch_size, self.config.out_channels, num_frames * p_t, latent_h * p_h, latent_w * p_w + ) + + if not return_dict: + return (hidden_states,) + return Transformer2DModelOutput(sample=hidden_states) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 572bab22745b..cddbab1f037b 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -102,6 +102,11 @@ "QwenImageLayeredModularPipeline", "QwenImageLayeredAutoBlocks", ] + _import_structure["abot_world"] = [ + "ABotWorldBlocks", + "ABotWorldModularPipeline", + "ABotWorldStreamingBlocks", + ] _import_structure["anima"] = [ "AnimaAutoBlocks", "AnimaModularPipeline", @@ -151,6 +156,7 @@ except OptionalDependencyNotAvailable: from ..utils.dummy_pt_objects import * # noqa F403 else: + from .abot_world import ABotWorldBlocks, ABotWorldModularPipeline, ABotWorldStreamingBlocks from .anima import AnimaAutoBlocks, AnimaModularPipeline from .components_manager import ComponentsManager from .cosmos import ( diff --git a/src/diffusers/modular_pipelines/abot_world/__init__.py b/src/diffusers/modular_pipelines/abot_world/__init__.py new file mode 100644 index 000000000000..35040557f39b --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/__init__.py @@ -0,0 +1,47 @@ +from typing import TYPE_CHECKING + +from ...utils import ( + DIFFUSERS_SLOW_IMPORT, + OptionalDependencyNotAvailable, + _LazyModule, + get_objects_from_module, + is_torch_available, + is_transformers_available, +) + + +_dummy_objects = {} +_import_structure = {} + +try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() +except OptionalDependencyNotAvailable: + from ...utils import dummy_torch_and_transformers_objects # noqa F403 + + _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) +else: + _import_structure["modular_blocks_abot_world"] = ["ABotWorldBlocks", "ABotWorldStreamingBlocks"] + _import_structure["modular_pipeline"] = ["ABotWorldModularPipeline"] + +if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: + try: + if not (is_transformers_available() and is_torch_available()): + raise OptionalDependencyNotAvailable() + except OptionalDependencyNotAvailable: + from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 + else: + from .modular_blocks_abot_world import ABotWorldBlocks, ABotWorldStreamingBlocks + from .modular_pipeline import ABotWorldModularPipeline +else: + import sys + + sys.modules[__name__] = _LazyModule( + __name__, + globals()["__file__"], + _import_structure, + module_spec=__spec__, + ) + + for name, value in _dummy_objects.items(): + setattr(sys.modules[__name__], name, value) diff --git a/src/diffusers/modular_pipelines/abot_world/before_denoise.py b/src/diffusers/modular_pipelines/abot_world/before_denoise.py new file mode 100644 index 000000000000..77a0164a837a --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/before_denoise.py @@ -0,0 +1,121 @@ +# 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. + +import numpy as np +import torch + +from ...models import ABotWorldTransformer3DModel +from ...models.autoencoders.autoencoder_kl_wan import WanDecodeCache +from ...models.autoencoders.autoencoder_tiny_video import TinyVideoDecodeCache +from ...models.transformers.transformer_abot_world import ABotWorldKVCache +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ABotWorldPrepareStep(ModularPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Prepare step for the causal rollout: sets the scheduler's full shifted flow-match grid and warps the " + "distilled `denoising_timesteps` through it, and allocates the transformer's rolling K/V cache with " + "the reference tokens pinned at its head." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "denoising_timesteps", + type_hint=list[int], + default=[1000, 750, 500, 250], + description="The distilled student's denoising timesteps, before shift-warping", + ), + InputParam("height", type_hint=int, default=704, description="Height of the generated video in pixels"), + InputParam("width", type_hint=int, default=1280, description="Width of the generated video in pixels"), + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="Normalized VAE latents of the reference views `[B, K, C, 1, h, w]`", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "denoise_timesteps", + type_hint=torch.Tensor, + description="The warped denoising timesteps the rollout loop iterates", + ), + OutputParam("kv_cache", type_hint=ABotWorldKVCache, description="The rollout's rolling K/V cache"), + OutputParam( + "decode_cache", + type_hint=WanDecodeCache | TinyVideoDecodeCache, + description="The decoder's cache carried across the rollout loop; starts as `None` and is created by " + "the decode step on the first block", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + # the full 1000-point flow-match grid the reference warps its step list through: the scheduler + # applies its configured shift to sigmas = linspace(1, 0, 1001)[:-1] + components.scheduler.set_timesteps(sigmas=np.linspace(1.0, 0.0, 1001)[:-1].tolist()) + timesteps = components.scheduler.timesteps.float() + step_list = torch.tensor(block_state.denoising_timesteps, dtype=torch.long) + block_state.denoise_timesteps = torch.cat([timesteps, timesteps.new_zeros(1)])[1000 - step_list] + + config = components.transformer.config + frame_seqlen = (block_state.height // 16 // config.patch_size[1]) * ( + block_state.width // 16 // config.patch_size[2] + ) + num_slots, _, ref_t, ref_h, ref_w = block_state.reference_latents.shape[1:] + ref_token_len = ( + num_slots + * (ref_t // config.patch_size[0]) + * (ref_h // config.patch_size[1]) + * (ref_w // config.patch_size[2]) + ) + block_state.decode_cache = None + block_state.kv_cache = ABotWorldKVCache( + num_layers=config.num_layers, + batch_size=block_state.reference_latents.shape[0], + num_tokens=ref_token_len + config.local_attn_size * frame_seqlen, + ref_token_len=ref_token_len, + num_heads=config.num_attention_heads, + head_dim=config.attention_head_dim, + device=device, + dtype=components.transformer.dtype, + ) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/abot_world/denoise.py b/src/diffusers/modular_pipelines/abot_world/denoise.py new file mode 100644 index 000000000000..ceef7e514c44 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -0,0 +1,679 @@ +# 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. + +from typing import Callable + +import numpy as np +import torch +from tqdm import tqdm + +from ...configuration_utils import FrozenDict +from ...models import ABotWorldTransformer3DModel, AutoencoderKLWan, AutoencoderTinyVideo +from ...models.autoencoders.autoencoder_kl_wan import WanDecodeCache +from ...models.autoencoders.autoencoder_tiny_video import TinyVideoDecodeCache +from ...models.transformers.transformer_abot_world import ABotWorldKVCache +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +from ...video_processor import VideoProcessor +from ..modular_pipeline import IterativePipelineBlocks, ModularLoopPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ABotWorldSetActionStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the rollout loop that broadcasts this block's `[W, A, S, D, I, J, K, L]` action vector into " + "constant pixel-resolution planes (each key repeated over 4 channels), which the transformer's action " + "adapter encodes and adds onto the patch tokens. `action` is a loop variable: the rollout takes it from " + "the `actions` list or callable, and a live driver passes it to " + "`loop_step(components, state, k=k, action=action)`." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("height", type_hint=int, default=704, description="Height of the generated video in pixels"), + InputParam("width", type_hint=int, default=1280, description="Width of the generated video in pixels"), + InputParam( + "num_frames_per_block", + type_hint=int, + default=3, + description="Latent frames generated per block (the model was trained with 3)", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "action_planes", + type_hint=torch.Tensor, + description="This block's broadcast action planes `[B, 32, F, height, width]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + device = components._execution_device + + action = torch.as_tensor(action, dtype=torch.float32) + if action.shape != (8,): + raise ValueError( + f"`action` must be an 8-element `[W, A, S, D, I, J, K, L]` vector, got shape {action.shape}" + ) + action = action.to(device=device, dtype=components.transformer.dtype) + block_state.action_planes = ( + action.view(1, 8, 1, 1, 1) + .repeat_interleave(4, dim=1) + .repeat(1, 1, block_state.num_frames_per_block, block_state.height, block_state.width) + ) + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldPrepareNoiseStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the rollout loop that draws this block's initial noise and computes its token offset " + "`current_start` in the rollout. On the first block (`current_start == 0`) the clean starting-frame " + "latent is pinned as frame 0." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam.template("generator"), + InputParam( + "first_frame_latents", + required=True, + type_hint=torch.Tensor, + description="Normalized VAE latent of the starting frame `[B, C, 1, h, w]`", + ), + InputParam("height", type_hint=int, default=704, description="Height of the generated video in pixels"), + InputParam("width", type_hint=int, default=1280, description="Width of the generated video in pixels"), + InputParam( + "num_frames_per_block", + type_hint=int, + default=3, + description="Latent frames generated per block (the model was trained with 3)", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "latents", type_hint=torch.Tensor, description="This block's working latents `[B, C, F, h, w]`" + ), + OutputParam( + "current_start", + type_hint=int, + description="Token offset of this block in the rollout: `k * F * tokens_per_frame`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + device = components._execution_device + + config = components.transformer.config + latent_height = block_state.height // 16 + latent_width = block_state.width // 16 + frame_seqlen = (latent_height // config.patch_size[1]) * (latent_width // config.patch_size[2]) + num_frames = block_state.num_frames_per_block + batch_size = block_state.first_frame_latents.shape[0] + + # drawn in the reference's [B, F, C, h, w] layout so a seeded run consumes the RNG identically + noise = randn_tensor( + (batch_size, num_frames, config.in_channels, latent_height, latent_width), + generator=block_state.generator, + device=device, + dtype=components.transformer.dtype, + ) + block_state.latents = noise.permute(0, 2, 1, 3, 4) + block_state.current_start = k * num_frames * frame_seqlen + if block_state.current_start == 0: + block_state.latents[:, :, :1] = block_state.first_frame_latents.to(block_state.latents.dtype) + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldLoopDenoiser(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the block's denoising loop: one transformer forward at timestep `t` (per-frame; frame 0 of " + "the first block is held at timestep 0), velocity converted to x0 with the warped sigma grid, then — " + "except on the last step — re-noised to the next timestep with fresh noise. On the first block the " + "clean starting-frame latent is re-pinned after every step. This block should be used to compose the " + "`sub_blocks` attribute of an `IterativePipelineBlocks` object (e.g. `ABotWorldDenoiseLoopWrapper`)." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ComponentSpec("scheduler", FlowMatchEulerDiscreteScheduler), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", required=True, type_hint=torch.Tensor, description="This block's working latents"), + InputParam( + "action_planes", + required=True, + type_hint=torch.Tensor, + description="This block's broadcast action planes", + ), + InputParam.template("prompt_embeds"), + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="Normalized VAE latents of the reference views `[B, K, C, 1, h, w]`", + ), + InputParam( + "reference_mask", + required=True, + type_hint=torch.Tensor, + description="Per-slot validity mask `[B, K]` for the reference views", + ), + InputParam( + "first_frame_latents", + required=True, + type_hint=torch.Tensor, + description="Normalized VAE latent of the starting frame `[B, C, 1, h, w]`", + ), + InputParam( + "kv_cache", + required=True, + type_hint=ABotWorldKVCache, + description="The rollout's rolling K/V cache", + ), + InputParam( + "current_start", + required=True, + type_hint=int, + description="Token offset of this block in the rollout", + ), + InputParam( + "denoise_timesteps", + required=True, + type_hint=torch.Tensor, + description="The warped denoising timesteps the loop iterates", + ), + InputParam.template("generator"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("latents", type_hint=torch.Tensor, description="The (partially) denoised block latents"), + ] + + def _lookup_sigma(self, scheduler, timestep: torch.Tensor) -> torch.Tensor: + """Per-element sigma via nearest-timestep lookup on the scheduler's warped grid, like the reference + wrapper's flow -> x0 conversion (timestep 0 of the pinned first frame is off-grid, hence nearest).""" + timesteps = scheduler.timesteps.double().to(timestep.device) + index = torch.argmin((timesteps.unsqueeze(0) - timestep.double().unsqueeze(1)).abs(), dim=1) + return scheduler.sigmas.double().to(timestep.device)[index] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size, _, num_frames = block_state.latents.shape[:3] + timestep = t.to(device).expand(batch_size, num_frames).clone() + if block_state.current_start == 0: + timestep[:, 0] = 0 + + velocity = components.transformer( + hidden_states=block_state.latents.to(components.transformer.dtype), + timestep=timestep, + encoder_hidden_states=block_state.prompt_embeds.to(components.transformer.dtype), + action_hidden_states=block_state.action_planes, + reference_hidden_states=block_state.reference_latents.to(components.transformer.dtype), + reference_mask=block_state.reference_mask, + kv_cache=block_state.kv_cache, + current_start=block_state.current_start, + return_dict=False, + )[0] + + # velocity -> x0 in double precision with per-frame sigmas (frame 0 of the first block sits at + # timestep 0), matching the reference wrapper's `_convert_flow_to_x0` + sigma = self._lookup_sigma(components.scheduler, timestep.flatten()).reshape(batch_size, 1, num_frames, 1, 1) + x0 = (block_state.latents.double() - sigma * velocity.double()).to(block_state.latents.dtype) + + if i < len(block_state.denoise_timesteps) - 1: + noise = randn_tensor( + (batch_size, num_frames, x0.shape[1], x0.shape[3], x0.shape[4]), + generator=block_state.generator, + device=device, + dtype=x0.dtype, + ).permute(0, 2, 1, 3, 4) + next_t = block_state.denoise_timesteps[i + 1].to(device).unsqueeze(0) + block_state.latents = components.scheduler.scale_noise(x0, next_t, noise) + else: + block_state.latents = x0 + + if block_state.current_start == 0: + block_state.latents[:, :, :1] = block_state.first_frame_latents.to(block_state.latents.dtype) + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldDenoiseLoopWrapper(IterativePipelineBlocks): + model_name = "abot-world" + + @property + def loop_variables(self) -> list[str]: + return ["i", "t"] + + @property + def description(self) -> str: + return ( + "Pipeline block that denoises one rollout block over the distilled `denoise_timesteps`. It runs inside " + "the rollout loop and accepts its loop variables (`k`, `action`)." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "denoise_timesteps", + required=True, + type_hint=torch.Tensor, + description="The warped denoising timesteps the loop iterates", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.denoise_timesteps): + components, state = self.loop_step(components, state, i=i, t=t) + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + for i, t in enumerate(block_state.denoise_timesteps): + components, state = yield from self.stream_step(components, state, i=i, t=t) + return components, state + + +class ABotWorldDenoiseStep(ABotWorldDenoiseLoopWrapper): + block_classes = [ABotWorldLoopDenoiser] + block_names = ["denoiser"] + + @property + def description(self) -> str: + return ( + "Denoise step for one rollout block: the distilled few-step loop. \n" + "Its loop logic is defined in `ABotWorldDenoiseLoopWrapper.__call__` method \n" + "At each iteration, it runs blocks defined in `sub_blocks` sequentially:\n" + " - `ABotWorldLoopDenoiser`\n" + ) + + +class ABotWorldCacheUpdateStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the rollout loop that runs one extra transformer forward on the finished block at the " + "context noise level (timestep 0), purely to write the clean block into the K/V cache that future " + "blocks attend over. The output is discarded." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", required=True, type_hint=torch.Tensor, description="The denoised block latents"), + InputParam( + "action_planes", + required=True, + type_hint=torch.Tensor, + description="This block's broadcast action planes", + ), + InputParam.template("prompt_embeds"), + InputParam( + "reference_latents", + required=True, + type_hint=torch.Tensor, + description="Normalized VAE latents of the reference views `[B, K, C, 1, h, w]`", + ), + InputParam( + "reference_mask", + required=True, + type_hint=torch.Tensor, + description="Per-slot validity mask `[B, K]` for the reference views", + ), + InputParam( + "kv_cache", + required=True, + type_hint=ABotWorldKVCache, + description="The rollout's rolling K/V cache", + ), + InputParam( + "current_start", + required=True, + type_hint=int, + description="Token offset of this block in the rollout", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + device = components._execution_device + + batch_size, _, num_frames = block_state.latents.shape[:3] + timestep = torch.zeros((batch_size, num_frames), dtype=torch.long, device=device) + + components.transformer( + hidden_states=block_state.latents.to(components.transformer.dtype), + timestep=timestep, + encoder_hidden_states=block_state.prompt_embeds.to(components.transformer.dtype), + action_hidden_states=block_state.action_planes, + reference_hidden_states=block_state.reference_latents.to(components.transformer.dtype), + reference_mask=block_state.reference_mask, + kv_cache=block_state.kv_cache, + current_start=block_state.current_start, + return_dict=False, + ) + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldDecodeStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the rollout loop that decodes the finished block to pixels with the full Wan VAE. The causal " + "VAE continues from the previous blocks through a loop-carried decode cache, so only the new block's " + "latents are decoded and the frames match a single decode of the whole rollout." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", required=True, type_hint=torch.Tensor, description="The denoised block latents"), + InputParam( + "decode_cache", + type_hint=WanDecodeCache, + description="The VAE's causal-conv cache carried over from the previous blocks; `None` for the first block", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("frames", type_hint=np.ndarray, description="This block's decoded frames `[T, H, W, 3]`"), + OutputParam( + "decode_cache", + type_hint=WanDecodeCache, + description="The VAE's causal-conv cache after this block, carried to the next block", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + vae = components.vae + + latents = block_state.latents.to(vae.dtype) + latents_mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=vae.dtype).view( + 1, -1, 1, 1, 1 + ) + latents_std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=vae.dtype).view(1, -1, 1, 1, 1) + + cache = block_state.decode_cache if block_state.decode_cache is not None else WanDecodeCache() + video = vae.decode(latents * latents_std + latents_mean, return_dict=False, cache=cache)[0] + block_state.frames = components.video_processor.postprocess_video(video, output_type="np")[0] + block_state.decode_cache = cache + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldTinyDecodeStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the rollout loop that decodes the finished block to pixels with the tiny VAE (`taew2_2`), " + "fast enough for real-time play. The decoder continues from the previous blocks through a loop-carried " + "decode cache, so only the new block's latents are decoded." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("tiny_vae", AutoencoderTinyVideo), + ComponentSpec( + "video_processor", + VideoProcessor, + config=FrozenDict({"vae_scale_factor": 16}), + default_creation_method="from_config", + ), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("latents", required=True, type_hint=torch.Tensor, description="The denoised block latents"), + InputParam( + "decode_cache", + type_hint=TinyVideoDecodeCache, + description="The tiny VAE's memory carried over from the previous blocks; `None` for the first block", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam("frames", type_hint=np.ndarray, description="This block's decoded frames `[T, H, W, 3]`"), + OutputParam( + "decode_cache", + type_hint=TinyVideoDecodeCache, + description="The tiny VAE's memory after this block, carried to the next block", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): + block_state = self.get_block_state(state) + tiny_vae = components.tiny_vae + + # the tiny VAE works in the normalized latent space the transformer produces + cache = block_state.decode_cache if block_state.decode_cache is not None else TinyVideoDecodeCache() + video = tiny_vae.decode(block_state.latents.to(tiny_vae.dtype), return_dict=False, cache=cache)[0] + block_state.frames = components.video_processor.postprocess_video(video, output_type="np")[0] + block_state.decode_cache = cache + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldRolloutWrapper(IterativePipelineBlocks): + model_name = "abot-world" + + @property + def loop_variables(self) -> list[str]: + return ["k", "action"] + + @property + def description(self) -> str: + return ( + "Pipeline block that rolls the world out block by block: at each block it encodes the block's action, " + "draws noise, runs the distilled denoising loop against the rolling K/V cache, writes the finished block " + "back into the cache and decodes it to pixels. `actions` is either a scripted list (one block per " + "vector) or, with `stream()`, a callable polled once per block. Drive it through " + "`loop_step(components, state, k=k, action=action)` to own the iteration yourself." + ) + + @property + def loop_inputs(self) -> list[InputParam]: + return [ + InputParam( + "actions", + required=True, + type_hint=list[list[int]] | Callable, + description=( + "Per-block `[W, A, S, D, I, J, K, L]` 0/1 action vectors (W/A/S/D move, I/J/K/L turn the " + "camera): a list generates one block per vector; with `stream()`, a callable " + "`(block_index) -> vector or None` is polled once per block instead and the rollout runs until " + "it returns `None`." + ), + ), + ] + + @property + def loop_intermediate_outputs(self) -> list[OutputParam]: + # the loop logic collects each block's decoded frames + return [OutputParam("videos", type_hint=list[np.ndarray], description="The generated videos")] + + @torch.no_grad() + def __call__(self, components, state: PipelineState): + block_state = self.get_block_state(state) + if callable(block_state.actions): + raise ValueError( + "A callable `actions` needs `stream()`: with `__call__` no frames flow back between polls, so there " + "is nothing to react to. Use `pipe.stream(...)`, or pass a scripted list instead." + ) + + frames = [] + with tqdm(total=len(block_state.actions), desc="Rollout") as progress_bar: + for k, action in enumerate(block_state.actions): + components, state = self.loop_step(components, state, k=k, action=action) + frames.append(state.get("frames")) + progress_bar.update() + block_state.videos = [np.concatenate(frames, axis=0)] + self.set_block_state(state, block_state) + + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + if callable(block_state.actions): + next_action = block_state.actions + else: + scripted = block_state.actions + + def next_action(k): + return scripted[k] if k < len(scripted) else None + + frames, k = [], 0 + while (action := next_action(k)) is not None: + components, state = yield from self.stream_step(components, state, k=k, action=action) + frames.append(state.get("frames")) + k += 1 + block_state.videos = [np.concatenate(frames, axis=0)] + self.set_block_state(state, block_state) + + return components, state + + +class ABotWorldRolloutStep(ABotWorldRolloutWrapper): + block_classes = [ + ABotWorldSetActionStep, + ABotWorldPrepareNoiseStep, + ABotWorldDenoiseStep, + ABotWorldCacheUpdateStep, + ABotWorldDecodeStep, + ] + block_names = ["set_action", "prepare_noise", "denoise", "cache_update", "decode"] + + @property + def description(self) -> str: + return ( + "Rollout step that generates and decodes the world block by block.\n" + "At each block: set_action -> prepare_noise -> denoise (a nested distilled denoising loop) -> " + "cache_update -> decode." + ) + + +class ABotWorldStreamingRolloutStep(ABotWorldRolloutWrapper): + block_classes = [ + ABotWorldSetActionStep, + ABotWorldPrepareNoiseStep, + ABotWorldDenoiseStep, + ABotWorldCacheUpdateStep, + ABotWorldTinyDecodeStep, + ] + block_names = ["set_action", "prepare_noise", "denoise", "cache_update", "decode"] + + @property + def description(self) -> str: + return ( + "Streaming rollout step that generates and decodes the world block by block.\n" + "At each block: set_action -> prepare_noise -> denoise (a nested distilled denoising loop) -> " + "cache_update -> decode." + ) diff --git a/src/diffusers/modular_pipelines/abot_world/encoders.py b/src/diffusers/modular_pipelines/abot_world/encoders.py new file mode 100644 index 000000000000..0d0fc5b1cf58 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/encoders.py @@ -0,0 +1,224 @@ +# 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. + +import numpy as np +import PIL.Image +import torch +from PIL import ImageOps +from transformers import AutoTokenizer, UMT5EncoderModel + +from ...models import AutoencoderKLWan +from ...utils import logging +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +def encode_image_to_latent(vae: AutoencoderKLWan, image: PIL.Image.Image, device, dtype) -> torch.Tensor: + """Preprocess a PIL image to `[-1, 1]` and VAE-encode it to a normalized latent `[1, C, 1, h, w]`.""" + pixels = torch.from_numpy(np.array(image, dtype=np.float32)) + pixels = pixels.to(device=device, dtype=dtype) * (2 / 255) - 1 + pixels = pixels.permute(2, 0, 1)[None, :, None] # [1, C, 1, H, W] + + latent = vae.encode(pixels).latent_dist.mode().float() + latents_mean = torch.tensor(vae.config.latents_mean, device=device).view(1, -1, 1, 1, 1) + latents_std = torch.tensor(vae.config.latents_std, device=device).view(1, -1, 1, 1, 1) + return (latent - latents_mean) / latents_std + + +class ABotWorldTextEncoderStep(ModularPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return "Text encoder step that encodes the prompt with umt5-xxl; masked positions are zeroed." + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("text_encoder", UMT5EncoderModel), + ComponentSpec("tokenizer", AutoTokenizer), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("prompt", required=True, type_hint=str, description="The text prompt describing the world"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam.template("prompt_embeds"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + text_inputs = components.tokenizer( + [block_state.prompt], + padding="max_length", + max_length=512, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + input_ids = text_inputs.input_ids.to(device) + mask = text_inputs.attention_mask.to(device) + seq_lens = mask.gt(0).sum(dim=1).long() + + prompt_embeds = components.text_encoder(input_ids, mask).last_hidden_state + prompt_embeds = prompt_embeds.to(components.text_encoder.dtype) + prompt_embeds = torch.stack( + [torch.cat([u[:v], u.new_zeros(u.size(0) - v, u.size(1))]) for u, v in zip(prompt_embeds, seq_lens)] + ) + block_state.prompt_embeds = prompt_embeds + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldImageEncoderStep(ModularPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Image encoder step that fits the input image to the target resolution (cover + center-crop) and " + "VAE-encodes it into `first_frame_latents`. The rollout loop pins this clean latent as frame 0 of the " + "first block, which is how the generated world starts from the input image." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam("image", required=True, type_hint=PIL.Image.Image, description="The starting frame"), + InputParam("height", type_hint=int, default=704, description="Height of the generated video in pixels"), + InputParam("width", type_hint=int, default=1280, description="Width of the generated video in pixels"), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "first_frame_latents", + type_hint=torch.Tensor, + description="Normalized VAE latent of the starting frame `[B, C, 1, h, w]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + image = ImageOps.fit( + block_state.image.convert("RGB"), + (block_state.width, block_state.height), + method=PIL.Image.LANCZOS, + centering=(0.5, 0.5), + ) + block_state.first_frame_latents = encode_image_to_latent(components.vae, image, device, components.vae.dtype) + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldRefImagesEncoderStep(ModularPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Reference encoder step that VAE-encodes the character reference views (e.g. head/left/right/front/back " + "at 512x512) into `reference_latents`. The transformer pins these tokens at the head of its K/V cache, " + "so every generated frame attends to them — this is what keeps the character consistent over an " + "unbounded rollout. Without `reference_images` (a plain scene rollout), zero latents with a zero " + "`reference_mask` are emitted, matching the reference implementation's ref-less mode." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "reference_images", + type_hint=list[PIL.Image.Image], + description=( + "The character reference views; each is resized to `reference_resolution`. Omit for a plain " + "scene rollout without a reference character." + ), + ), + InputParam( + "reference_resolution", + type_hint=int, + default=512, + description="Side length the reference views are resized to before encoding", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "reference_latents", + type_hint=torch.Tensor, + description="Normalized VAE latents of the reference views `[B, K, C, 1, h, w]`", + ), + OutputParam( + "reference_mask", + type_hint=torch.Tensor, + description="Per-slot validity mask `[B, K]`: ones for encoded views, zeros in the ref-less mode", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + if block_state.reference_images is None: + # ref-less rollout: zero latents in the standard 5 slots, masked out entirely + latent_size = block_state.reference_resolution // 16 + block_state.reference_latents = torch.zeros( + 1, 5, components.vae.config.z_dim, 1, latent_size, latent_size, device=device + ) + block_state.reference_mask = torch.zeros(1, 5, device=device) + else: + size = (block_state.reference_resolution, block_state.reference_resolution) + latents = [ + encode_image_to_latent(components.vae, img.convert("RGB").resize(size), device, components.vae.dtype) + for img in block_state.reference_images + ] + block_state.reference_latents = torch.stack(latents, dim=1) # [B, K, C, 1, h, w] + block_state.reference_mask = torch.ones(1, len(latents), device=device) + + self.set_block_state(state, block_state) + return components, state diff --git a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py new file mode 100644 index 000000000000..c1d4bb17527c --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -0,0 +1,354 @@ +# 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. + +from ...utils import logging +from ..modular_pipeline import SequentialPipelineBlocks +from ..modular_pipeline_utils import InsertableDict +from .before_denoise import ABotWorldPrepareStep +from .denoise import ABotWorldRolloutStep, ABotWorldStreamingRolloutStep +from .encoders import ABotWorldImageEncoderStep, ABotWorldRefImagesEncoderStep, ABotWorldTextEncoderStep + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +ABotWorldCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", ABotWorldPrepareStep()), + ("rollout", ABotWorldRolloutStep()), + ] +) + + +# auto_docstring +class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step that prepares the denoising schedule and the rolling K/V cache, then rolls the world out block by + block conditioned on the per-block actions. + + Components: + transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + height (`int`, *optional*, defaults to 704): + Height of the generated video in pixels + width (`int`, *optional*, defaults to 1280): + Width of the generated video in pixels + reference_latents (`Tensor`): + Normalized VAE latents of the reference views `[B, K, C, 1, h, w]` + actions (`list | Callable`): + Per-block `[W, A, S, D, I, J, K, L]` 0/1 action vectors (W/A/S/D move, I/J/K/L turn the camera): a list + generates one block per vector; with `stream()`, a callable `(block_index) -> vector or None` is polled + once per block instead and the rollout runs until it returns `None`. + num_frames_per_block (`int`, *optional*, defaults to 3): + Latent frames generated per block (the model was trained with 3) + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + first_frame_latents (`Tensor`): + Normalized VAE latent of the starting frame `[B, C, 1, h, w]` + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + reference_mask (`Tensor`): + Per-slot validity mask `[B, K]` for the reference views + + Outputs: + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + decode_cache (`WanDecodeCache | TinyVideoDecodeCache`): + The decoder's cache carried across the rollout loop; starts as `None` and is created by the decode step + on the first block + action_planes (`Tensor`): + This block's broadcast action planes `[B, 32, F, height, width]` + latents (`Tensor`): + This block's working latents `[B, C, F, h, w]` + current_start (`int`): + Token offset of this block in the rollout: `k * F * tokens_per_frame` + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + videos (`list`): + The generated videos + """ + + model_name = "abot-world" + block_classes = ABotWorldCoreDenoiseBlocks.values() + block_names = ABotWorldCoreDenoiseBlocks.keys() + + @property + def description(self): + return ( + "Core denoise step that prepares the denoising schedule and the rolling K/V cache, then rolls the world " + "out block by block conditioned on the per-block actions." + ) + + +BLOCKS = InsertableDict( + [ + ("text_encoder", ABotWorldTextEncoderStep()), + ("image_encoder", ABotWorldImageEncoderStep()), + ("ref_encoder", ABotWorldRefImagesEncoderStep()), + ("denoise", ABotWorldCoreDenoiseStep()), + ] +) + + +ABotWorldStreamingCoreDenoiseBlocks = InsertableDict( + [ + ("prepare", ABotWorldPrepareStep()), + ("rollout", ABotWorldStreamingRolloutStep()), + ] +) + + +# auto_docstring +class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): + """ + Core denoise step for the streaming workflow: prepares the denoising schedule and the rolling K/V cache, then rolls + the world out block by block, decoding each block to pixels inside the loop. + + Components: + transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) tiny_vae + (`AutoencoderTinyVideo`) video_processor (`VideoProcessor`) + + Inputs: + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + height (`int`, *optional*, defaults to 704): + Height of the generated video in pixels + width (`int`, *optional*, defaults to 1280): + Width of the generated video in pixels + reference_latents (`Tensor`): + Normalized VAE latents of the reference views `[B, K, C, 1, h, w]` + actions (`list | Callable`): + Per-block `[W, A, S, D, I, J, K, L]` 0/1 action vectors (W/A/S/D move, I/J/K/L turn the camera): a list + generates one block per vector; with `stream()`, a callable `(block_index) -> vector or None` is polled + once per block instead and the rollout runs until it returns `None`. + num_frames_per_block (`int`, *optional*, defaults to 3): + Latent frames generated per block (the model was trained with 3) + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + first_frame_latents (`Tensor`): + Normalized VAE latent of the starting frame `[B, C, 1, h, w]` + prompt_embeds (`Tensor`): + text embeddings used to guide the image generation. Can be generated from text_encoder step. + reference_mask (`Tensor`): + Per-slot validity mask `[B, K]` for the reference views + + Outputs: + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + decode_cache (`WanDecodeCache | TinyVideoDecodeCache`): + The decoder's cache carried across the rollout loop; starts as `None` and is created by the decode step + on the first block + action_planes (`Tensor`): + This block's broadcast action planes `[B, 32, F, height, width]` + latents (`Tensor`): + This block's working latents `[B, C, F, h, w]` + current_start (`int`): + Token offset of this block in the rollout: `k * F * tokens_per_frame` + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + videos (`list`): + The generated videos + """ + + model_name = "abot-world" + block_classes = ABotWorldStreamingCoreDenoiseBlocks.values() + block_names = ABotWorldStreamingCoreDenoiseBlocks.keys() + + @property + def description(self): + return ( + "Core denoise step for the streaming workflow: prepares the denoising schedule and the rolling K/V " + "cache, then rolls the world out block by block, decoding each block to pixels inside the loop." + ) + + +STREAMING_BLOCKS = InsertableDict( + [ + ("text_encoder", ABotWorldTextEncoderStep()), + ("image_encoder", ABotWorldImageEncoderStep()), + ("ref_encoder", ABotWorldRefImagesEncoderStep()), + ("denoise", ABotWorldStreamingCoreDenoiseStep()), + ] +) + + +# auto_docstring +class ABotWorldStreamingBlocks(SequentialPipelineBlocks): + """ + Streaming/interactive ABot-World world generation: like the default blockset, but each block is decoded to pixels + inside the rollout loop, so `pipe.stream(...)` yields ready frames per ~1 s block and a live driver gets frames + back from every `loop_step` call. Interactive drivers own the loop via the rollout block's `loop_step`, writing the + current `action` into the state between calls. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer + (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) tiny_vae + (`AutoencoderTinyVideo`) video_processor (`VideoProcessor`) + + Inputs: + prompt (`str`): + The text prompt describing the world + image (`Image`): + The starting frame + height (`int`, *optional*, defaults to 704): + Height of the generated video in pixels + width (`int`, *optional*, defaults to 1280): + Width of the generated video in pixels + reference_images (`list`, *optional*): + The character reference views; each is resized to `reference_resolution`. Omit for a plain scene rollout + without a reference character. + reference_resolution (`int`, *optional*, defaults to 512): + Side length the reference views are resized to before encoding + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + actions (`list | Callable`): + Per-block `[W, A, S, D, I, J, K, L]` 0/1 action vectors (W/A/S/D move, I/J/K/L turn the camera): a list + generates one block per vector; with `stream()`, a callable `(block_index) -> vector or None` is polled + once per block instead and the rollout runs until it returns `None`. + num_frames_per_block (`int`, *optional*, defaults to 3): + Latent frames generated per block (the model was trained with 3) + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + first_frame_latents (`Tensor`): + Normalized VAE latent of the starting frame `[B, C, 1, h, w]` + reference_latents (`Tensor`): + Normalized VAE latents of the reference views `[B, K, C, 1, h, w]` + reference_mask (`Tensor`): + Per-slot validity mask `[B, K]`: ones for encoded views, zeros in the ref-less mode + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + decode_cache (`WanDecodeCache | TinyVideoDecodeCache`): + The decoder's cache carried across the rollout loop; starts as `None` and is created by the decode step + on the first block + action_planes (`Tensor`): + This block's broadcast action planes `[B, 32, F, height, width]` + latents (`Tensor`): + This block's working latents `[B, C, F, h, w]` + current_start (`int`): + Token offset of this block in the rollout: `k * F * tokens_per_frame` + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + videos (`list`): + The generated videos + """ + + model_name = "abot-world" + block_classes = STREAMING_BLOCKS.values() + block_names = STREAMING_BLOCKS.keys() + + @property + def description(self): + return ( + "Streaming/interactive ABot-World world generation: like the default blockset, but each block is " + "decoded to pixels inside the rollout loop, so `pipe.stream(...)` yields ready frames per ~1 s block " + "and a live driver gets frames back from every `loop_step` call. Interactive drivers own the loop via " + "the rollout block's `loop_step`, writing the current `action` into the state between calls." + ) + + +# auto_docstring +class ABotWorldBlocks(SequentialPipelineBlocks): + """ + Action-conditioned world generation with ABot-World: starting from an input image and character reference views, + the model rolls the world out block by block (3 latent frames each), steered by a per-block `[W, A, S, D, I, J, K, + L]` action vector. Scripted rollouts pass the full action list; streaming consumers use `pipe.stream(...)` for a + live state after every denoise step and block; interactive drivers own the loop via the rollout block's + `loop_step`, writing new actions into the state between calls. + + Components: + text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer + (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) video_processor + (`VideoProcessor`) + + Inputs: + prompt (`str`): + The text prompt describing the world + image (`Image`): + The starting frame + height (`int`, *optional*, defaults to 704): + Height of the generated video in pixels + width (`int`, *optional*, defaults to 1280): + Width of the generated video in pixels + reference_images (`list`, *optional*): + The character reference views; each is resized to `reference_resolution`. Omit for a plain scene rollout + without a reference character. + reference_resolution (`int`, *optional*, defaults to 512): + Side length the reference views are resized to before encoding + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + actions (`list | Callable`): + Per-block `[W, A, S, D, I, J, K, L]` 0/1 action vectors (W/A/S/D move, I/J/K/L turn the camera): a list + generates one block per vector; with `stream()`, a callable `(block_index) -> vector or None` is polled + once per block instead and the rollout runs until it returns `None`. + num_frames_per_block (`int`, *optional*, defaults to 3): + Latent frames generated per block (the model was trained with 3) + generator (`Generator`, *optional*): + Torch generator for deterministic generation. + + Outputs: + prompt_embeds (`Tensor`): + The prompt embeddings. + first_frame_latents (`Tensor`): + Normalized VAE latent of the starting frame `[B, C, 1, h, w]` + reference_latents (`Tensor`): + Normalized VAE latents of the reference views `[B, K, C, 1, h, w]` + reference_mask (`Tensor`): + Per-slot validity mask `[B, K]`: ones for encoded views, zeros in the ref-less mode + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + decode_cache (`WanDecodeCache | TinyVideoDecodeCache`): + The decoder's cache carried across the rollout loop; starts as `None` and is created by the decode step + on the first block + action_planes (`Tensor`): + This block's broadcast action planes `[B, 32, F, height, width]` + latents (`Tensor`): + This block's working latents `[B, C, F, h, w]` + current_start (`int`): + Token offset of this block in the rollout: `k * F * tokens_per_frame` + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + videos (`list`): + The generated videos + """ + + model_name = "abot-world" + block_classes = BLOCKS.values() + block_names = BLOCKS.keys() + + @property + def description(self): + return ( + "Action-conditioned world generation with ABot-World: starting from an input image and character " + "reference views, the model rolls the world out block by block (3 latent frames each), steered by a " + "per-block `[W, A, S, D, I, J, K, L]` action vector. Scripted rollouts pass the full action list; " + "streaming consumers use `pipe.stream(...)` for a live state after every denoise step and block; " + "interactive drivers own the loop via the rollout block's `loop_step`, writing new actions into the " + "state between calls." + ) diff --git a/src/diffusers/modular_pipelines/abot_world/modular_pipeline.py b/src/diffusers/modular_pipelines/abot_world/modular_pipeline.py new file mode 100644 index 000000000000..cfe91f2ea29e --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/modular_pipeline.py @@ -0,0 +1,28 @@ +# 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. + +from ...utils import logging +from ..modular_pipeline import ModularPipeline + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ABotWorldModularPipeline(ModularPipeline): + """ + A ModularPipeline for ABot-World action-conditioned world generation. + + """ + + default_blocks_name = "ABotWorldBlocks" diff --git a/src/diffusers/modular_pipelines/modular_pipeline.py b/src/diffusers/modular_pipelines/modular_pipeline.py index 5d14dfd4e10b..859a681479f4 100644 --- a/src/diffusers/modular_pipelines/modular_pipeline.py +++ b/src/diffusers/modular_pipelines/modular_pipeline.py @@ -145,6 +145,7 @@ def _helios_pyramid_map_fn(config_dict=None): ("qwenimage-edit", _create_default_map_fn("QwenImageEditModularPipeline")), ("qwenimage-edit-plus", _create_default_map_fn("QwenImageEditPlusModularPipeline")), ("qwenimage-layered", _create_default_map_fn("QwenImageLayeredModularPipeline")), + ("abot-world", _create_default_map_fn("ABotWorldModularPipeline")), ("anima", _create_default_map_fn("AnimaModularPipeline")), ("z-image", _create_default_map_fn("ZImageModularPipeline")), ("cosmos3-omni", _create_default_map_fn("Cosmos3OmniModularPipeline")), diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index d4691a6a3a76..44aaafe04300 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -405,6 +405,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class ABotWorldTransformer3DModel(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AceStepTransformer1DModel(metaclass=DummyObject): _backends = ["torch"] @@ -915,6 +930,21 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) +class AutoencoderTinyVideo(metaclass=DummyObject): + _backends = ["torch"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch"]) + + class AutoencoderVidTok(metaclass=DummyObject): _backends = ["torch"] diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 376596d632ea..b0c680eb04e8 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2,6 +2,36 @@ from ..utils import DummyObject, requires_backends +class ABotWorldBlocks(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + +class ABotWorldModularPipeline(metaclass=DummyObject): + _backends = ["torch", "transformers"] + + def __init__(self, *args, **kwargs): + requires_backends(self, ["torch", "transformers"]) + + @classmethod + def from_config(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + @classmethod + def from_pretrained(cls, *args, **kwargs): + requires_backends(cls, ["torch", "transformers"]) + + class AnimaAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"]