From e7fdfdb2bac09d67f4f5f4ece4c6808f62ac595f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 20 Aug 2026 02:54:38 +0000 Subject: [PATCH 1/6] Add ABot-World: real-time action-conditioned world model (Wan2.2-TI2V-5B) Integrates https://github.com/amap-cvlab/ABot-World (acvlab/ABot-World-0-5B-LF, Apache-2.0) as a modular pipeline. - ABotWorldTransformer3DModel: block-causal rollout over a rolling K/V cache (window eviction, pinned reference-token prefix, relative RoPE with periodic rebase), per-frame timesteps, keyboard-action adapter added onto the patch tokens, per-stream cross-attention cache. Bit-exact vs the reference (CPU/fp32) incl. eviction, rebase, and the real 5B weights. - scripts/convert_abot_world_to_diffusers.py: pure key renames, no surgery. VAE and umt5 are byte-identical to Wan-AI/Wan2.2-TI2V-5B and are reused from Wan-AI/Wan2.2-TI2V-5B-Diffusers; FlowMatchEulerDiscreteScheduler(shift=5.0) covers the warped DMD grid (scale_noise == the reference re-noise). - modular_pipelines/abot_world: text/image/reference encoders -> core denoise (prepare + rollout IterativePipelineBlocks over blocks k, with a nested distilled denoise loop over (i, t) and a KV-cache context update) -> decode. Streams via pipe.stream() (events per denoise step and per ~1s block) and drives interactively via loop_step, writing new actions into the state between calls. - Converted checkpoint + runnable example: YiYiXu/ABot-World-0-5B-LF-Diffusers Co-Authored-By: Claude Fable 5 --- scripts/convert_abot_world_to_diffusers.py | 57 ++ src/diffusers/__init__.py | 6 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/transformers/__init__.py | 1 + .../transformers/transformer_abot_world.py | 699 ++++++++++++++++++ src/diffusers/modular_pipelines/__init__.py | 5 + .../modular_pipelines/abot_world/__init__.py | 47 ++ .../abot_world/before_denoise.py | 126 ++++ .../modular_pipelines/abot_world/decoders.py | 95 +++ .../modular_pipelines/abot_world/denoise.py | 506 +++++++++++++ .../modular_pipelines/abot_world/encoders.py | 207 ++++++ .../abot_world/modular_blocks_abot_world.py | 182 +++++ .../abot_world/modular_pipeline.py | 28 + .../modular_pipelines/modular_pipeline.py | 1 + src/diffusers/utils/dummy_pt_objects.py | 15 + .../dummy_torch_and_transformers_objects.py | 30 + 16 files changed, 2007 insertions(+) create mode 100644 scripts/convert_abot_world_to_diffusers.py create mode 100644 src/diffusers/models/transformers/transformer_abot_world.py create mode 100644 src/diffusers/modular_pipelines/abot_world/__init__.py create mode 100644 src/diffusers/modular_pipelines/abot_world/before_denoise.py create mode 100644 src/diffusers/modular_pipelines/abot_world/decoders.py create mode 100644 src/diffusers/modular_pipelines/abot_world/denoise.py create mode 100644 src/diffusers/modular_pipelines/abot_world/encoders.py create mode 100644 src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py create mode 100644 src/diffusers/modular_pipelines/abot_world/modular_pipeline.py 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/src/diffusers/__init__.py b/src/diffusers/__init__.py index 88ede2fd009a..5565ba2a7673 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -224,6 +224,7 @@ ] _import_structure["models"].extend( [ + "ABotWorldTransformer3DModel", "AceStepTransformer1DModel", "AllegroTransformer3DModel", "AnimaTextConditioner", @@ -513,6 +514,8 @@ else: _import_structure["modular_pipelines"].extend( [ + "ABotWorldBlocks", + "ABotWorldModularPipeline", "AnimaAutoBlocks", "AnimaModularPipeline", "Cosmos3DistilledBlocks", @@ -1098,6 +1101,7 @@ VaeImageProcessorLDM3D, ) from .models import ( + ABotWorldTransformer3DModel, AceStepTransformer1DModel, AllegroTransformer3DModel, AnimaTextConditioner, @@ -1366,6 +1370,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..ec7f66df9c0e 100755 --- a/src/diffusers/models/__init__.py +++ b/src/diffusers/models/__init__.py @@ -103,6 +103,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"] @@ -234,6 +235,7 @@ from .embeddings import ImageProjection from .modeling_utils import ModelMixin from .transformers import ( + ABotWorldTransformer3DModel, AceStepTransformer1DModel, AllegroTransformer3DModel, AnyFlowFARTransformer3DModel, 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..5ceac2b1406d 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -102,6 +102,10 @@ "QwenImageLayeredModularPipeline", "QwenImageLayeredAutoBlocks", ] + _import_structure["abot_world"] = [ + "ABotWorldBlocks", + "ABotWorldModularPipeline", + ] _import_structure["anima"] = [ "AnimaAutoBlocks", "AnimaModularPipeline", @@ -151,6 +155,7 @@ except OptionalDependencyNotAvailable: from ..utils.dummy_pt_objects import * # noqa F403 else: + from .abot_world import ABotWorldBlocks, ABotWorldModularPipeline 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..8096eb725cfa --- /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"] + _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 + 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..5d5c6a29412a --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/before_denoise.py @@ -0,0 +1,126 @@ +# 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.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, validates the per-block actions, 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( + "actions", + required=True, + type_hint=list[list[int]], + description=( + "Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block " + "(W/A/S/D move, I/J/K/L turn the camera). The rollout generates `len(actions)` blocks." + ), + ), + 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("actions", type_hint=torch.Tensor, description="The actions as a `[num_blocks, 8]` tensor"), + 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"), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + block_state.actions = torch.tensor(block_state.actions, dtype=torch.float32) + if block_state.actions.ndim != 2 or block_state.actions.shape[1] != 8: + raise ValueError(f"`actions` must be a list of 8-element vectors, got shape {block_state.actions.shape}") + + # 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.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/decoders.py b/src/diffusers/modular_pipelines/abot_world/decoders.py new file mode 100644 index 000000000000..e5d05ded2615 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/decoders.py @@ -0,0 +1,95 @@ +# 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 Union + +import numpy as np +import PIL +import torch + +from ...configuration_utils import FrozenDict +from ...models import AutoencoderKLWan +from ...utils import logging +from ...video_processor import VideoProcessor +from ..modular_pipeline import ModularPipelineBlocks, PipelineState +from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam + + +logger = logging.get_logger(__name__) # pylint: disable=invalid-name + + +class ABotWorldDecodeStep(ModularPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return "Step that de-normalizes the rollout's latents and VAE-decodes them into the output video." + + @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("output_type", default="np"), + InputParam( + "video_latents", + required=True, + type_hint=torch.Tensor, + description="The rollout's accumulated latents `[B, C, T, h, w]`", + ), + ] + + @property + def intermediate_outputs(self) -> list[OutputParam]: + return [ + OutputParam( + "videos", + type_hint=Union[list[list[PIL.Image.Image]], list[torch.Tensor], list[np.ndarray]], + description="The generated videos", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + vae = components.vae + + if block_state.output_type == "latent": + block_state.videos = block_state.video_latents + else: + latents = block_state.video_latents.to(vae.dtype) + latents_mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype).view( + 1, -1, 1, 1, 1 + ) + latents_std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype).view( + 1, -1, 1, 1, 1 + ) + latents = latents * latents_std + latents_mean + video = vae.decode(latents, return_dict=False)[0] + block_state.videos = components.video_processor.postprocess_video( + video, output_type=block_state.output_type + ) + + 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..6d509915be06 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -0,0 +1,506 @@ +# 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 torch +from tqdm import tqdm + +from ...models import ABotWorldTransformer3DModel +from ...models.transformers.transformer_abot_world import ABotWorldKVCache +from ...schedulers import FlowMatchEulerDiscreteScheduler +from ...utils import logging +from ...utils.torch_utils import randn_tensor +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. An interactive driver overwrites `actions` in the " + "state between `loop_step` calls to steer the world live." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "actions", + required=True, + type_hint=torch.Tensor, + description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", + ), + 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): + block_state = self.get_block_state(state) + device = components._execution_device + + action = block_state.actions[k].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): + 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( + "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), + 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 reads the current block index `k` from the loop scope." + ) + + @property + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself, on top of what the sub-blocks declare + loop_inputs = [ + InputParam( + "denoise_timesteps", + required=True, + type_hint=torch.Tensor, + description="The warped denoising timesteps the loop iterates", + ), + ] + return [param for param in loop_inputs if param.name not in names] + inputs + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int): + 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): + 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( + "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): + 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), + 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 ABotWorldRolloutWrapper(IterativePipelineBlocks): + model_name = "abot-world" + + @property + def loop_variables(self) -> list[str]: + return ["k"] + + @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, and writes the finished " + "block back into the cache. Drive it through `loop_step(components, state, k=k)` to own the iteration — " + "write new `actions` into the state between calls for live interaction." + ) + + @property + def inputs(self) -> list[InputParam]: + inputs = super().inputs + names = {param.name for param in inputs} + # `actions` is also consumed by the loop logic itself (the rollout length) + loop_inputs = [ + InputParam( + "actions", + required=True, + type_hint=torch.Tensor, + description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", + ), + ] + return [param for param in loop_inputs if param.name not in names] + inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + # produced by the loop logic itself, which collects each finished block + return super().intermediate_outputs + [ + OutputParam( + "video_latents", + type_hint=torch.Tensor, + description="The rollout's accumulated latents `[B, C, num_blocks * F, h, w]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState): + block_state = self.get_block_state(state) + + video_latents = [] + with tqdm(total=block_state.actions.shape[0], desc="Rollout") as progress_bar: + for k in range(block_state.actions.shape[0]): + components, state = self.loop_step(components, state, k=k) + video_latents.append(state.get("latents")) + progress_bar.update() + state.set("video_latents", torch.cat(video_latents, dim=2)) + + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + + video_latents = [] + for k in range(block_state.actions.shape[0]): + components, state = yield from self.stream_step(components, state, k=k) + video_latents.append(state.get("latents")) + state.set("video_latents", torch.cat(video_latents, dim=2)) + + return components, state + + +class ABotWorldRolloutStep(ABotWorldRolloutWrapper): + block_classes = [ + ABotWorldSetActionStep, + ABotWorldPrepareNoiseStep, + ABotWorldDenoiseStep, + ABotWorldCacheUpdateStep, + ] + block_names = ["set_action", "prepare_noise", "denoise", "cache_update"] + + @property + def description(self) -> str: + return ( + "Rollout step that generates the world block by block.\n" + "At each block: set_action -> prepare_noise -> denoise (a nested distilled denoising loop) -> " + "cache_update." + ) 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..bdd6227abdc9 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/encoders.py @@ -0,0 +1,207 @@ +# 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." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("vae", AutoencoderKLWan), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "reference_images", + required=True, + type_hint=list[PIL.Image.Image], + description="The character reference views; each is resized to `reference_resolution`", + ), + 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]`", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState) -> PipelineState: + block_state = self.get_block_state(state) + device = components._execution_device + + 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] + + 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..92215ce39f97 --- /dev/null +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -0,0 +1,182 @@ +# 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 .decoders import ABotWorldDecodeStep +from .denoise import ABotWorldRolloutStep +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`) + + Inputs: + actions (`list`): + Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L + turn the camera). The rollout generates `len(actions)` blocks. + 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]` + 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. + + Outputs: + actions (`Tensor`): + The actions as a `[num_blocks, 8]` tensor + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + 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` + video_latents (`Tensor`): + The rollout's accumulated latents `[B, C, num_blocks * F, h, w]` + """ + + 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()), + ("decode", ABotWorldDecodeStep()), + ] +) + + +# 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`): + The character reference views; each is resized to `reference_resolution` + reference_resolution (`int`, *optional*, defaults to 512): + Side length the reference views are resized to before encoding + actions (`list`): + Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L + turn the camera). The rollout generates `len(actions)` blocks. + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + 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. + output_type (`None`, *optional*, defaults to np): + TODO: Add description. + + 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]` + actions (`Tensor`): + The actions as a `[num_blocks, 8]` tensor + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + 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` + video_latents (`Tensor`): + The rollout's accumulated latents `[B, C, num_blocks * F, h, w]` + videos (`list | list | 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 d3e703d6bd85..b10a8b891bf9 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..22bed2f6472f 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -420,6 +420,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 AllegroTransformer3DModel(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..14cc3d937e78 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -17,6 +17,36 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) +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 AnimaModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] From 8e4198e7097972a2ab12213f9fbcb39e7ec7f842 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 20 Aug 2026 06:21:05 +0000 Subject: [PATCH 2/6] Add the ABot-World streaming/interactive blockset MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - ABotWorldStreamingBlocks: like the default blockset, but each block is decoded to pixels inside the rollout loop (trailing-window decode with the previous block as temporal context; independent per-block decode of the causal VAE produces artifacts), so stream() yields ready frames per ~1 s block - the rollout takes a scripted `actions` list or an `action_source` callable polled once per block (return None to stop) — a live driver is just a stream() consumer whose callable returns the player's current action; with __call__ the callable combination raises, since no frames flow back between polls - set_action reads the single current `action` from state (no actions[k] indexing); `actions` becomes optional in the prepare step Co-Authored-By: Claude Fable 5 --- src/diffusers/modular_pipelines/__init__.py | 3 +- .../modular_pipelines/abot_world/__init__.py | 4 +- .../abot_world/before_denoise.py | 13 +- .../modular_pipelines/abot_world/denoise.py | 269 +++++++++++++++++- .../abot_world/modular_blocks_abot_world.py | 181 +++++++++++- src/diffusers/utils/dummy_pt_objects.py | 4 +- .../dummy_torch_and_transformers_objects.py | 6 +- 7 files changed, 461 insertions(+), 19 deletions(-) diff --git a/src/diffusers/modular_pipelines/__init__.py b/src/diffusers/modular_pipelines/__init__.py index 5ceac2b1406d..cddbab1f037b 100644 --- a/src/diffusers/modular_pipelines/__init__.py +++ b/src/diffusers/modular_pipelines/__init__.py @@ -105,6 +105,7 @@ _import_structure["abot_world"] = [ "ABotWorldBlocks", "ABotWorldModularPipeline", + "ABotWorldStreamingBlocks", ] _import_structure["anima"] = [ "AnimaAutoBlocks", @@ -155,7 +156,7 @@ except OptionalDependencyNotAvailable: from ..utils.dummy_pt_objects import * # noqa F403 else: - from .abot_world import ABotWorldBlocks, ABotWorldModularPipeline + 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 index 8096eb725cfa..35040557f39b 100644 --- a/src/diffusers/modular_pipelines/abot_world/__init__.py +++ b/src/diffusers/modular_pipelines/abot_world/__init__.py @@ -21,7 +21,7 @@ _dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects)) else: - _import_structure["modular_blocks_abot_world"] = ["ABotWorldBlocks"] + _import_structure["modular_blocks_abot_world"] = ["ABotWorldBlocks", "ABotWorldStreamingBlocks"] _import_structure["modular_pipeline"] = ["ABotWorldModularPipeline"] if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT: @@ -31,7 +31,7 @@ except OptionalDependencyNotAvailable: from ...utils.dummy_torch_and_transformers_objects import * # noqa F403 else: - from .modular_blocks_abot_world import ABotWorldBlocks + from .modular_blocks_abot_world import ABotWorldBlocks, ABotWorldStreamingBlocks from .modular_pipeline import ABotWorldModularPipeline else: import sys diff --git a/src/diffusers/modular_pipelines/abot_world/before_denoise.py b/src/diffusers/modular_pipelines/abot_world/before_denoise.py index 5d5c6a29412a..3e354a810129 100644 --- a/src/diffusers/modular_pipelines/abot_world/before_denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/before_denoise.py @@ -49,11 +49,11 @@ def inputs(self) -> list[InputParam]: return [ InputParam( "actions", - required=True, type_hint=list[list[int]], description=( "Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block " - "(W/A/S/D move, I/J/K/L turn the camera). The rollout generates `len(actions)` blocks." + "(W/A/S/D move, I/J/K/L turn the camera); the scripted rollout generates `len(actions)` " + "blocks. Omit when driving the rollout interactively through `loop_step`." ), ), InputParam( @@ -89,9 +89,12 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - block_state.actions = torch.tensor(block_state.actions, dtype=torch.float32) - if block_state.actions.ndim != 2 or block_state.actions.shape[1] != 8: - raise ValueError(f"`actions` must be a list of 8-element vectors, got shape {block_state.actions.shape}") + if block_state.actions is not None: + block_state.actions = torch.tensor(block_state.actions, dtype=torch.float32) + if block_state.actions.ndim != 2 or block_state.actions.shape[1] != 8: + raise ValueError( + f"`actions` must be a list of 8-element vectors, got shape {block_state.actions.shape}" + ) # 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] diff --git a/src/diffusers/modular_pipelines/abot_world/denoise.py b/src/diffusers/modular_pipelines/abot_world/denoise.py index 6d509915be06..5c29e1dfe3a2 100644 --- a/src/diffusers/modular_pipelines/abot_world/denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -12,14 +12,19 @@ # 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 ...models import ABotWorldTransformer3DModel +from ...configuration_utils import FrozenDict +from ...models import ABotWorldTransformer3DModel, AutoencoderKLWan 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 @@ -464,6 +469,11 @@ def intermediate_outputs(self) -> list[OutputParam]: @torch.no_grad() def __call__(self, components, state: PipelineState): block_state = self.get_block_state(state) + if block_state.actions is None: + raise ValueError( + "A scripted rollout requires `actions` (one vector per block); to drive the rollout " + "interactively, call `loop_step` yourself and write `action` into the state between calls." + ) video_latents = [] with tqdm(total=block_state.actions.shape[0], desc="Rollout") as progress_bar: @@ -478,6 +488,11 @@ def __call__(self, components, state: PipelineState): @torch.no_grad() def stream(self, components, state: PipelineState): block_state = self.get_block_state(state) + if block_state.actions is None: + raise ValueError( + "A scripted rollout requires `actions` (one vector per block); to drive the rollout " + "interactively, call `loop_step` yourself and write `action` into the state between calls." + ) video_latents = [] for k in range(block_state.actions.shape[0]): @@ -504,3 +519,255 @@ def description(self) -> str: "At each block: set_action -> prepare_noise -> denoise (a nested distilled denoising loop) -> " "cache_update." ) + + +class ABotWorldCurrentActionStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the streaming rollout loop that broadcasts the *current* `[W, A, S, D, I, J, K, L]` action " + "vector (the `action` state value) into the conditioning planes. The scripted `__call__`/`stream` set " + "`action` from the `actions` list each iteration; a live driver writes it into the state between " + "`loop_step` calls." + ) + + @property + def expected_components(self) -> list[ComponentSpec]: + return [ + ComponentSpec("transformer", ABotWorldTransformer3DModel), + ] + + @property + def inputs(self) -> list[InputParam]: + return [ + InputParam( + "action", + required=True, + type_hint=torch.Tensor, + description="The current block's `[W, A, S, D, I, J, K, L]` 0/1 action vector", + ), + 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): + block_state = self.get_block_state(state) + device = components._execution_device + + action = torch.as_tensor(block_state.action, dtype=torch.float32) + 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 ABotWorldStreamingDecodeStep(ModularLoopPipelineBlocks): + model_name = "abot-world" + + @property + def description(self) -> str: + return ( + "Step within the streaming rollout loop that decodes the finished block to pixels. The causal VAE needs " + "temporal context, so the previous block's latents are decoded alongside and only the new block's " + "frames are kept." + ) + + @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( + "previous_latents", + type_hint=torch.Tensor, + description="The previous block's latents, decoded as temporal context; `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( + "previous_latents", + type_hint=torch.Tensor, + description="This block's latents, kept as the next block's decode context", + ), + ] + + @torch.no_grad() + def __call__(self, components, state: PipelineState, k: int): + block_state = self.get_block_state(state) + vae = components.vae + + latents = block_state.latents + if block_state.previous_latents is None: + decode_input = latents + new_frames = None # keep everything + else: + decode_input = torch.cat([block_state.previous_latents, latents], dim=2) + new_frames = latents.shape[2] * 4 + + decode_input = decode_input.to(vae.dtype) + latents_mean = torch.tensor(vae.config.latents_mean, device=decode_input.device, dtype=vae.dtype).view( + 1, -1, 1, 1, 1 + ) + latents_std = torch.tensor(vae.config.latents_std, device=decode_input.device, dtype=vae.dtype).view( + 1, -1, 1, 1, 1 + ) + video = vae.decode(decode_input * latents_std + latents_mean, return_dict=False)[0] + video = components.video_processor.postprocess_video(video, output_type="np")[0] + + block_state.frames = video if new_frames is None else video[-new_frames:] + block_state.previous_latents = latents + + self.set_block_state(state, block_state) + return components, state + + +class ABotWorldStreamingRolloutWrapper(IterativePipelineBlocks): + model_name = "abot-world" + + @property + def loop_variables(self) -> list[str]: + return ["k"] + + @property + def description(self) -> str: + return ( + "Pipeline block that rolls the world out block by block and decodes each block to pixels inside the " + "loop. Scripted runs iterate the `actions` list (each iteration writes the block's `action` into the " + "state, exactly as a live driver would); interactive drivers own the loop via " + "`loop_step(components, state, k=k)` and write `action` between calls." + ) + + @property + def inputs(self) -> list[InputParam]: + # `action` and `previous_latents` are loop-carried — supplied per iteration by the loop logic (or a live + # driver) and by the decode step of the previous iteration — never user-provided. + inputs = [param for param in super().inputs if param.name not in ("action", "previous_latents")] + names = {param.name for param in inputs} + # inputs consumed by the loop logic itself + loop_inputs = [ + InputParam( + "actions", + type_hint=torch.Tensor, + description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", + ), + InputParam( + "action_source", + type_hint=Callable, + description=( + "Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` " + "polled once per block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, " + "or `None` to stop. The rollout is unbounded while it returns actions." + ), + ), + ] + return [param for param in loop_inputs if param.name not in names] + inputs + + @property + def intermediate_outputs(self) -> list[OutputParam]: + # produced by the loop logic itself, which collects each block's decoded frames + return super().intermediate_outputs + [ + OutputParam("videos", type_hint=list[np.ndarray], description="The generated videos"), + ] + + def _next_action(self, block_state, k): + if block_state.actions is not None: + return block_state.actions[k] if k < block_state.actions.shape[0] else None + if block_state.action_source is not None: + action = block_state.action_source(k) + return None if action is None else torch.as_tensor(action, dtype=torch.float32) + raise ValueError( + "The streaming rollout needs `actions` (a scripted list) or `action_source` (a callable polled per " + "block); to own the loop yourself instead, call `loop_step` and write `action` into the state " + "between calls." + ) + + @torch.no_grad() + def __call__(self, components, state: PipelineState): + block_state = self.get_block_state(state) + if block_state.actions is None and block_state.action_source is not None: + raise ValueError( + "`action_source` needs `stream()`: with `__call__` no frames flow back between polls, so there " + "is nothing to react to. Use `pipe.stream(...)`, or pass a scripted `actions` list instead." + ) + + frames, k = [], 0 + while (action := self._next_action(block_state, k)) is not None: + state.set("action", action) + components, state = self.loop_step(components, state, k=k) + frames.append(state.get("frames")) + k += 1 + state.set("videos", [np.concatenate(frames, axis=0)]) + + return components, state + + @torch.no_grad() + def stream(self, components, state: PipelineState): + block_state = self.get_block_state(state) + + frames, k = [], 0 + while (action := self._next_action(block_state, k)) is not None: + state.set("action", action) + components, state = yield from self.stream_step(components, state, k=k) + frames.append(state.get("frames")) + k += 1 + state.set("videos", [np.concatenate(frames, axis=0)]) + + return components, state + + +class ABotWorldStreamingRolloutStep(ABotWorldStreamingRolloutWrapper): + block_classes = [ + ABotWorldCurrentActionStep, + ABotWorldPrepareNoiseStep, + ABotWorldDenoiseStep, + ABotWorldCacheUpdateStep, + ABotWorldStreamingDecodeStep, + ] + 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/modular_blocks_abot_world.py b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py index 92215ce39f97..2a5b0f9629b0 100644 --- a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -17,7 +17,7 @@ from ..modular_pipeline_utils import InsertableDict from .before_denoise import ABotWorldPrepareStep from .decoders import ABotWorldDecodeStep -from .denoise import ABotWorldRolloutStep +from .denoise import ABotWorldRolloutStep, ABotWorldStreamingRolloutStep from .encoders import ABotWorldImageEncoderStep, ABotWorldRefImagesEncoderStep, ABotWorldTextEncoderStep @@ -42,9 +42,10 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) Inputs: - actions (`list`): + actions (`list`, *optional*): Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera). The rollout generates `len(actions)` blocks. + turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout + interactively through `loop_step`. 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): @@ -102,6 +103,175 @@ def description(self): ) +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`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + + Inputs: + actions (`list`, *optional*): + Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L + turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout + interactively through `loop_step`. + 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]` + action_source (`Callable`, *optional*): + Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` polled once per + block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, or `None` to stop. The + rollout is unbounded while it returns actions. + 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. + + Outputs: + actions (`Tensor`): + The actions as a `[num_blocks, 8]` tensor + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + 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]` + previous_latents (`Tensor`): + This block's latents, kept as the next block's decode context + 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`) 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`): + The character reference views; each is resized to `reference_resolution` + reference_resolution (`int`, *optional*, defaults to 512): + Side length the reference views are resized to before encoding + actions (`list`, *optional*): + Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L + turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout + interactively through `loop_step`. + denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): + The distilled student's denoising timesteps, before shift-warping + action_source (`Callable`, *optional*): + Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` polled once per + block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, or `None` to stop. The + rollout is unbounded while it returns actions. + 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]` + actions (`Tensor`): + The actions as a `[num_blocks, 8]` tensor + denoise_timesteps (`Tensor`): + The warped denoising timesteps the rollout loop iterates + kv_cache (`ABotWorldKVCache`): + The rollout's rolling K/V cache + 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]` + previous_latents (`Tensor`): + This block's latents, kept as the next block's decode context + 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): """ @@ -129,9 +299,10 @@ class ABotWorldBlocks(SequentialPipelineBlocks): The character reference views; each is resized to `reference_resolution` reference_resolution (`int`, *optional*, defaults to 512): Side length the reference views are resized to before encoding - actions (`list`): + actions (`list`, *optional*): Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera). The rollout generates `len(actions)` blocks. + turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout + interactively through `loop_step`. denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): The distilled student's denoising timesteps, before shift-warping num_frames_per_block (`int`, *optional*, defaults to 3): diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index 22bed2f6472f..c8a176e9b0f3 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -405,7 +405,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class AceStepTransformer1DModel(metaclass=DummyObject): +class ABotWorldTransformer3DModel(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): @@ -420,7 +420,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch"]) -class ABotWorldTransformer3DModel(metaclass=DummyObject): +class AceStepTransformer1DModel(metaclass=DummyObject): _backends = ["torch"] def __init__(self, *args, **kwargs): diff --git a/src/diffusers/utils/dummy_torch_and_transformers_objects.py b/src/diffusers/utils/dummy_torch_and_transformers_objects.py index 14cc3d937e78..b0c680eb04e8 100644 --- a/src/diffusers/utils/dummy_torch_and_transformers_objects.py +++ b/src/diffusers/utils/dummy_torch_and_transformers_objects.py @@ -2,7 +2,7 @@ from ..utils import DummyObject, requires_backends -class AnimaAutoBlocks(metaclass=DummyObject): +class ABotWorldBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -17,7 +17,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class ABotWorldBlocks(metaclass=DummyObject): +class ABotWorldModularPipeline(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): @@ -32,7 +32,7 @@ def from_pretrained(cls, *args, **kwargs): requires_backends(cls, ["torch", "transformers"]) -class ABotWorldModularPipeline(metaclass=DummyObject): +class AnimaAutoBlocks(metaclass=DummyObject): _backends = ["torch", "transformers"] def __init__(self, *args, **kwargs): From ffee5d2d9ccf86db5c21190e223ea49435fc5145 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Thu, 20 Aug 2026 21:22:47 +0000 Subject: [PATCH 3/6] ABot-World: make reference_images optional (ref-less scene rollout) Without reference images the ref encoder emits zero reference latents with a zero reference_mask, matching the reference implementation's ref-less mode; the mask is plumbed through both transformer calls in the streaming rollout. Co-Authored-By: Claude Fable 5 --- .../modular_pipelines/abot_world/denoise.py | 14 ++++++++ .../modular_pipelines/abot_world/encoders.py | 35 ++++++++++++++----- .../abot_world/modular_blocks_abot_world.py | 18 +++++++--- 3 files changed, 54 insertions(+), 13 deletions(-) diff --git a/src/diffusers/modular_pipelines/abot_world/denoise.py b/src/diffusers/modular_pipelines/abot_world/denoise.py index 5c29e1dfe3a2..a85a48ec59aa 100644 --- a/src/diffusers/modular_pipelines/abot_world/denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -210,6 +210,12 @@ def inputs(self) -> list[InputParam]: 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, @@ -266,6 +272,7 @@ def __call__(self, components, state: PipelineState, i: int, t: torch.Tensor): 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, @@ -387,6 +394,12 @@ def inputs(self) -> list[InputParam]: 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, @@ -415,6 +428,7 @@ def __call__(self, components, state: PipelineState, k: int): 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, diff --git a/src/diffusers/modular_pipelines/abot_world/encoders.py b/src/diffusers/modular_pipelines/abot_world/encoders.py index bdd6227abdc9..0d0fc5b1cf58 100644 --- a/src/diffusers/modular_pipelines/abot_world/encoders.py +++ b/src/diffusers/modular_pipelines/abot_world/encoders.py @@ -155,7 +155,8 @@ def description(self) -> str: "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." + "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 @@ -169,9 +170,11 @@ def inputs(self) -> list[InputParam]: return [ InputParam( "reference_images", - required=True, type_hint=list[PIL.Image.Image], - description="The character reference views; each is resized to `reference_resolution`", + description=( + "The character reference views; each is resized to `reference_resolution`. Omit for a plain " + "scene rollout without a reference character." + ), ), InputParam( "reference_resolution", @@ -189,6 +192,11 @@ def intermediate_outputs(self) -> list[OutputParam]: 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() @@ -196,12 +204,21 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - 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] + 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 index 2a5b0f9629b0..f1de60d97cf3 100644 --- a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -62,6 +62,8 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): 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: actions (`Tensor`): @@ -146,6 +148,8 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): 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: actions (`Tensor`): @@ -212,8 +216,9 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): Height of the generated video in pixels width (`int`, *optional*, defaults to 1280): Width of the generated video in pixels - reference_images (`list`): - The character reference views; each is resized to `reference_resolution` + 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 actions (`list`, *optional*): @@ -238,6 +243,8 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): 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 actions (`Tensor`): The actions as a `[num_blocks, 8]` tensor denoise_timesteps (`Tensor`): @@ -295,8 +302,9 @@ class ABotWorldBlocks(SequentialPipelineBlocks): Height of the generated video in pixels width (`int`, *optional*, defaults to 1280): Width of the generated video in pixels - reference_images (`list`): - The character reference views; each is resized to `reference_resolution` + 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 actions (`list`, *optional*): @@ -319,6 +327,8 @@ class ABotWorldBlocks(SequentialPipelineBlocks): 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 actions (`Tensor`): The actions as a `[num_blocks, 8]` tensor denoise_timesteps (`Tensor`): From 5a29766dd898749056465fcde3746e71fdaaa6d7 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 21 Aug 2026 01:10:36 +0000 Subject: [PATCH 4/6] AutoencoderKLWan: WanDecodeCache for chunk-by-chunk decode; use it in the ABot-World streaming rollout decode(z, cache=cache) keeps the causal-conv feature cache across calls so a video can be decoded chunk by chunk with results identical to a single decode. The streaming decode step carries the cache through the rollout loop instead of re-decoding a trailing window of the previous block (bit-exact vs. a full decode; 3.3 -> 2.3 s per block on H100 at 704x1280). Co-Authored-By: Claude Fable 5 --- .../models/autoencoders/autoencoder_kl_wan.py | 44 ++++++++++++++--- .../modular_pipelines/abot_world/denoise.py | 47 ++++++++----------- .../abot_world/modular_blocks_abot_world.py | 8 ++-- 3 files changed, 60 insertions(+), 39 deletions(-) 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/modular_pipelines/abot_world/denoise.py b/src/diffusers/modular_pipelines/abot_world/denoise.py index a85a48ec59aa..078fa1069d8e 100644 --- a/src/diffusers/modular_pipelines/abot_world/denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -20,6 +20,7 @@ from ...configuration_utils import FrozenDict from ...models import ABotWorldTransformer3DModel, AutoencoderKLWan +from ...models.autoencoders.autoencoder_kl_wan import WanDecodeCache from ...models.transformers.transformer_abot_world import ABotWorldKVCache from ...schedulers import FlowMatchEulerDiscreteScheduler from ...utils import logging @@ -605,9 +606,9 @@ class ABotWorldStreamingDecodeStep(ModularLoopPipelineBlocks): @property def description(self) -> str: return ( - "Step within the streaming rollout loop that decodes the finished block to pixels. The causal VAE needs " - "temporal context, so the previous block's latents are decoded alongside and only the new block's " - "frames are kept." + "Step within the streaming rollout loop that decodes the finished block to pixels. 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 @@ -627,9 +628,9 @@ def inputs(self) -> list[InputParam]: return [ InputParam("latents", required=True, type_hint=torch.Tensor, description="The denoised block latents"), InputParam( - "previous_latents", - type_hint=torch.Tensor, - description="The previous block's latents, decoded as temporal context; `None` for the first block", + "decode_cache", + type_hint=WanDecodeCache, + description="The VAE's causal-conv cache carried over from the previous blocks; `None` for the first block", ), ] @@ -638,9 +639,9 @@ def intermediate_outputs(self) -> list[OutputParam]: return [ OutputParam("frames", type_hint=np.ndarray, description="This block's decoded frames `[T, H, W, 3]`"), OutputParam( - "previous_latents", - type_hint=torch.Tensor, - description="This block's latents, kept as the next block's decode context", + "decode_cache", + type_hint=WanDecodeCache, + description="The VAE's causal-conv cache after this block, carried to the next block", ), ] @@ -649,26 +650,16 @@ def __call__(self, components, state: PipelineState, k: int): block_state = self.get_block_state(state) vae = components.vae - latents = block_state.latents - if block_state.previous_latents is None: - decode_input = latents - new_frames = None # keep everything - else: - decode_input = torch.cat([block_state.previous_latents, latents], dim=2) - new_frames = latents.shape[2] * 4 - - decode_input = decode_input.to(vae.dtype) - latents_mean = torch.tensor(vae.config.latents_mean, device=decode_input.device, dtype=vae.dtype).view( - 1, -1, 1, 1, 1 - ) - latents_std = torch.tensor(vae.config.latents_std, device=decode_input.device, dtype=vae.dtype).view( + 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 ) - video = vae.decode(decode_input * latents_std + latents_mean, return_dict=False)[0] - video = components.video_processor.postprocess_video(video, output_type="np")[0] + latents_std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=vae.dtype).view(1, -1, 1, 1, 1) - block_state.frames = video if new_frames is None else video[-new_frames:] - block_state.previous_latents = latents + 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 @@ -692,9 +683,9 @@ def description(self) -> str: @property def inputs(self) -> list[InputParam]: - # `action` and `previous_latents` are loop-carried — supplied per iteration by the loop logic (or a live + # `action` and `decode_cache` are loop-carried — supplied per iteration by the loop logic (or a live # driver) and by the decode step of the previous iteration — never user-provided. - inputs = [param for param in super().inputs if param.name not in ("action", "previous_latents")] + inputs = [param for param in super().inputs if param.name not in ("action", "decode_cache")] names = {param.name for param in inputs} # inputs consumed by the loop logic itself loop_inputs = [ 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 index f1de60d97cf3..4e960eb71376 100644 --- a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -166,8 +166,8 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - previous_latents (`Tensor`): - This block's latents, kept as the next block's decode context + decode_cache (`WanDecodeCache`): + The VAE's causal-conv cache after this block, carried to the next block videos (`list`): The generated videos """ @@ -259,8 +259,8 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - previous_latents (`Tensor`): - This block's latents, kept as the next block's decode context + decode_cache (`WanDecodeCache`): + The VAE's causal-conv cache after this block, carried to the next block videos (`list`): The generated videos """ From 206121e08010ddc3451c3ad232ec166657a6da7f Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Fri, 21 Aug 2026 03:06:31 +0000 Subject: [PATCH 5/6] Add AutoencoderTinyVideo (TAEHV / taew2_2); ABot-World decodes per block in both presets AutoencoderTinyVideo ports madebyollin's TAEHV family as one class (config per checkpoint; defaults are taew2_2 for the Wan 2.2 latent space) with a TinyVideoDecodeCache for chunk-by-chunk decode, plus scripts/convert_taehv_to_diffusers.py and a docs page. Parity against the reference: decode 1.5e-8, streaming 2.5e-6, encode exact. The ABot-World presets now mirror the reference's vae_type switch: ABotWorldBlocks decodes every block inside the rollout loop with the full Wan VAE (ABotWorldDecodeStep + WanDecodeCache, as scripts/inference.py does), ABotWorldStreamingBlocks with the tiny VAE (ABotWorldTinyDecodeStep, as the interactive space does). The one-shot end decode is removed; it mirrored an entry point the reference never uses and needed all latents decoded at once. Streaming block time on H100: 2.3 s -> 0.9 s. Co-Authored-By: Claude Fable 5 --- docs/source/en/_toctree.yml | 2 + .../en/api/models/autoencoder_tiny_video.md | 37 +++ scripts/convert_taehv_to_diffusers.py | 103 +++++++ src/diffusers/__init__.py | 2 + src/diffusers/models/__init__.py | 2 + src/diffusers/models/autoencoders/__init__.py | 1 + .../autoencoders/autoencoder_tiny_video.py | 285 ++++++++++++++++++ .../modular_pipelines/abot_world/decoders.py | 95 ------ .../modular_pipelines/abot_world/denoise.py | 231 ++++++++------ .../abot_world/modular_blocks_abot_world.py | 39 +-- src/diffusers/utils/dummy_pt_objects.py | 15 + 11 files changed, 613 insertions(+), 199 deletions(-) create mode 100644 docs/source/en/api/models/autoencoder_tiny_video.md create mode 100644 scripts/convert_taehv_to_diffusers.py create mode 100644 src/diffusers/models/autoencoders/autoencoder_tiny_video.py delete mode 100644 src/diffusers/modular_pipelines/abot_world/decoders.py 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_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 5565ba2a7673..a39267d1e618 100644 --- a/src/diffusers/__init__.py +++ b/src/diffusers/__init__.py @@ -259,6 +259,7 @@ "AutoencoderRAE", "AutoencoderSAME", "AutoencoderTiny", + "AutoencoderTinyVideo", "AutoencoderVidTok", "AutoModel", "BriaFiboTransformer2DModel", @@ -1136,6 +1137,7 @@ AutoencoderRAE, AutoencoderSAME, AutoencoderTiny, + AutoencoderTinyVideo, AutoencoderVidTok, AutoModel, BriaFiboTransformer2DModel, diff --git a/src/diffusers/models/__init__.py b/src/diffusers/models/__init__.py index ec7f66df9c0e..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"] @@ -202,6 +203,7 @@ AutoencoderRAE, AutoencoderSAME, AutoencoderTiny, + AutoencoderTinyVideo, AutoencoderVidTok, ConsistencyDecoderVAE, Cosmos3AVAEAudioTokenizer, 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_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/modular_pipelines/abot_world/decoders.py b/src/diffusers/modular_pipelines/abot_world/decoders.py deleted file mode 100644 index e5d05ded2615..000000000000 --- a/src/diffusers/modular_pipelines/abot_world/decoders.py +++ /dev/null @@ -1,95 +0,0 @@ -# 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 Union - -import numpy as np -import PIL -import torch - -from ...configuration_utils import FrozenDict -from ...models import AutoencoderKLWan -from ...utils import logging -from ...video_processor import VideoProcessor -from ..modular_pipeline import ModularPipelineBlocks, PipelineState -from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam - - -logger = logging.get_logger(__name__) # pylint: disable=invalid-name - - -class ABotWorldDecodeStep(ModularPipelineBlocks): - model_name = "abot-world" - - @property - def description(self) -> str: - return "Step that de-normalizes the rollout's latents and VAE-decodes them into the output video." - - @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("output_type", default="np"), - InputParam( - "video_latents", - required=True, - type_hint=torch.Tensor, - description="The rollout's accumulated latents `[B, C, T, h, w]`", - ), - ] - - @property - def intermediate_outputs(self) -> list[OutputParam]: - return [ - OutputParam( - "videos", - type_hint=Union[list[list[PIL.Image.Image]], list[torch.Tensor], list[np.ndarray]], - description="The generated videos", - ), - ] - - @torch.no_grad() - def __call__(self, components, state: PipelineState) -> PipelineState: - block_state = self.get_block_state(state) - vae = components.vae - - if block_state.output_type == "latent": - block_state.videos = block_state.video_latents - else: - latents = block_state.video_latents.to(vae.dtype) - latents_mean = torch.tensor(vae.config.latents_mean, device=latents.device, dtype=latents.dtype).view( - 1, -1, 1, 1, 1 - ) - latents_std = torch.tensor(vae.config.latents_std, device=latents.device, dtype=latents.dtype).view( - 1, -1, 1, 1, 1 - ) - latents = latents * latents_std + latents_mean - video = vae.decode(latents, return_dict=False)[0] - block_state.videos = components.video_processor.postprocess_video( - video, output_type=block_state.output_type - ) - - 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 index 078fa1069d8e..01abd2e1f6dd 100644 --- a/src/diffusers/modular_pipelines/abot_world/denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -19,8 +19,9 @@ from tqdm import tqdm from ...configuration_utils import FrozenDict -from ...models import ABotWorldTransformer3DModel, AutoencoderKLWan +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 @@ -439,6 +440,131 @@ def __call__(self, components, state: PipelineState, k: int): 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): + 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): + 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" @@ -450,14 +576,15 @@ def loop_variables(self) -> list[str]: 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, and writes the finished " - "block back into the cache. Drive it through `loop_step(components, state, k=k)` to own the iteration — " - "write new `actions` into the state between calls for live interaction." + "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. Drive it through `loop_step(components, state, k=k)` to " + "own the iteration — write new `actions` into the state between calls for live interaction." ) @property def inputs(self) -> list[InputParam]: - inputs = super().inputs + # `decode_cache` is loop-carried — produced by the decode step of the previous iteration, never user-provided + inputs = [param for param in super().inputs if param.name != "decode_cache"] names = {param.name for param in inputs} # `actions` is also consumed by the loop logic itself (the rollout length) loop_inputs = [ @@ -472,13 +599,9 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: - # produced by the loop logic itself, which collects each finished block + # produced by the loop logic itself, which collects each block's decoded frames return super().intermediate_outputs + [ - OutputParam( - "video_latents", - type_hint=torch.Tensor, - description="The rollout's accumulated latents `[B, C, num_blocks * F, h, w]`", - ), + OutputParam("videos", type_hint=list[np.ndarray], description="The generated videos"), ] @torch.no_grad() @@ -490,13 +613,13 @@ def __call__(self, components, state: PipelineState): "interactively, call `loop_step` yourself and write `action` into the state between calls." ) - video_latents = [] + frames = [] with tqdm(total=block_state.actions.shape[0], desc="Rollout") as progress_bar: for k in range(block_state.actions.shape[0]): components, state = self.loop_step(components, state, k=k) - video_latents.append(state.get("latents")) + frames.append(state.get("frames")) progress_bar.update() - state.set("video_latents", torch.cat(video_latents, dim=2)) + state.set("videos", [np.concatenate(frames, axis=0)]) return components, state @@ -509,11 +632,11 @@ def stream(self, components, state: PipelineState): "interactively, call `loop_step` yourself and write `action` into the state between calls." ) - video_latents = [] + frames = [] for k in range(block_state.actions.shape[0]): components, state = yield from self.stream_step(components, state, k=k) - video_latents.append(state.get("latents")) - state.set("video_latents", torch.cat(video_latents, dim=2)) + frames.append(state.get("frames")) + state.set("videos", [np.concatenate(frames, axis=0)]) return components, state @@ -524,15 +647,16 @@ class ABotWorldRolloutStep(ABotWorldRolloutWrapper): ABotWorldPrepareNoiseStep, ABotWorldDenoiseStep, ABotWorldCacheUpdateStep, + ABotWorldDecodeStep, ] - block_names = ["set_action", "prepare_noise", "denoise", "cache_update"] + block_names = ["set_action", "prepare_noise", "denoise", "cache_update", "decode"] @property def description(self) -> str: return ( - "Rollout step that generates the world block by block.\n" + "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." + "cache_update -> decode." ) @@ -600,71 +724,6 @@ def __call__(self, components, state: PipelineState, k: int): return components, state -class ABotWorldStreamingDecodeStep(ModularLoopPipelineBlocks): - model_name = "abot-world" - - @property - def description(self) -> str: - return ( - "Step within the streaming rollout loop that decodes the finished block to pixels. 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): - 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 ABotWorldStreamingRolloutWrapper(IterativePipelineBlocks): model_name = "abot-world" @@ -765,7 +824,7 @@ class ABotWorldStreamingRolloutStep(ABotWorldStreamingRolloutWrapper): ABotWorldPrepareNoiseStep, ABotWorldDenoiseStep, ABotWorldCacheUpdateStep, - ABotWorldStreamingDecodeStep, + ABotWorldTinyDecodeStep, ] block_names = ["set_action", "prepare_noise", "denoise", "cache_update", "decode"] 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 index 4e960eb71376..f3b14e7dcd2e 100644 --- a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -16,7 +16,6 @@ from ..modular_pipeline import SequentialPipelineBlocks from ..modular_pipeline_utils import InsertableDict from .before_denoise import ABotWorldPrepareStep -from .decoders import ABotWorldDecodeStep from .denoise import ABotWorldRolloutStep, ABotWorldStreamingRolloutStep from .encoders import ABotWorldImageEncoderStep, ABotWorldRefImagesEncoderStep, ABotWorldTextEncoderStep @@ -39,7 +38,8 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): block conditioned on the per-block actions. Components: - transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) + transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) vae + (`AutoencoderKLWan`) video_processor (`VideoProcessor`) Inputs: actions (`list`, *optional*): @@ -78,8 +78,12 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): 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` - video_latents (`Tensor`): - The rollout's accumulated latents `[B, C, num_blocks * F, h, w]` + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + decode_cache (`WanDecodeCache`): + The VAE's causal-conv cache after this block, carried to the next block + videos (`list`): + The generated videos """ model_name = "abot-world" @@ -100,7 +104,6 @@ def description(self): ("image_encoder", ABotWorldImageEncoderStep()), ("ref_encoder", ABotWorldRefImagesEncoderStep()), ("denoise", ABotWorldCoreDenoiseStep()), - ("decode", ABotWorldDecodeStep()), ] ) @@ -120,8 +123,8 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): the world out block by block, decoding each block to pixels inside the loop. Components: - transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) vae - (`AutoencoderKLWan`) video_processor (`VideoProcessor`) + transformer (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) tiny_vae + (`AutoencoderTinyVideo`) video_processor (`VideoProcessor`) Inputs: actions (`list`, *optional*): @@ -166,8 +169,8 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`WanDecodeCache`): - The VAE's causal-conv cache after this block, carried to the next block + decode_cache (`TinyVideoDecodeCache`): + The tiny VAE's memory after this block, carried to the next block videos (`list`): The generated videos """ @@ -204,8 +207,8 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): Components: text_encoder (`UMT5EncoderModel`) tokenizer (`AutoTokenizer`) vae (`AutoencoderKLWan`) transformer - (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) video_processor - (`VideoProcessor`) + (`ABotWorldTransformer3DModel`) scheduler (`FlowMatchEulerDiscreteScheduler`) tiny_vae + (`AutoencoderTinyVideo`) video_processor (`VideoProcessor`) Inputs: prompt (`str`): @@ -259,8 +262,8 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`WanDecodeCache`): - The VAE's causal-conv cache after this block, carried to the next block + decode_cache (`TinyVideoDecodeCache`): + The tiny VAE's memory after this block, carried to the next block videos (`list`): The generated videos """ @@ -317,8 +320,6 @@ class ABotWorldBlocks(SequentialPipelineBlocks): Latent frames generated per block (the model was trained with 3) generator (`Generator`, *optional*): Torch generator for deterministic generation. - output_type (`None`, *optional*, defaults to np): - TODO: Add description. Outputs: prompt_embeds (`Tensor`): @@ -341,9 +342,11 @@ class ABotWorldBlocks(SequentialPipelineBlocks): 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` - video_latents (`Tensor`): - The rollout's accumulated latents `[B, C, num_blocks * F, h, w]` - videos (`list | list | list`): + frames (`ndarray`): + This block's decoded frames `[T, H, W, 3]` + decode_cache (`WanDecodeCache`): + The VAE's causal-conv cache after this block, carried to the next block + videos (`list`): The generated videos """ diff --git a/src/diffusers/utils/dummy_pt_objects.py b/src/diffusers/utils/dummy_pt_objects.py index c8a176e9b0f3..44aaafe04300 100644 --- a/src/diffusers/utils/dummy_pt_objects.py +++ b/src/diffusers/utils/dummy_pt_objects.py @@ -930,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"] From d2d0aeab86fd6e17c97d04266ecad5f25c168c24 Mon Sep 17 00:00:00 2001 From: yiyixuxu Date: Sat, 22 Aug 2026 01:05:01 +0000 Subject: [PATCH 6/6] ABot-World: action as a loop variable; one rollout wrapper taking a list or a callable `actions` MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - `action` is a loop variable next to `k`: the rollout loop passes each block's vector to its sub-blocks instead of writing it into the state, so a single `ABotWorldSetActionStep` serves both presets (`ABotWorldCurrentActionStep` removed) and a live driver calls `loop_step(components, state, k=k, action=action)`. - `actions` is either a scripted list (one block per vector) or, with `stream()`, a callable `(block_index) -> vector or None` polled once per block; `action_source` and the separate streaming wrapper are gone. The presets now differ only in their decode step. - Adopt `loop_inputs` / `loop_intermediate_outputs` from the base: `actions` and `denoise_timesteps` are loop inputs, `videos` is written through `set_block_state`. - `decode_cache` is seeded as `None` by the prepare step (created by the decode step on the first block) instead of being filtered out of the loop's inputs; the prepare step no longer validates `actions` — the set-action step checks each block's vector. Co-Authored-By: Claude Fable 5 --- .../abot_world/before_denoise.py | 30 +- .../modular_pipelines/abot_world/denoise.py | 274 ++++-------------- .../abot_world/modular_blocks_abot_world.py | 68 ++--- 3 files changed, 97 insertions(+), 275 deletions(-) diff --git a/src/diffusers/modular_pipelines/abot_world/before_denoise.py b/src/diffusers/modular_pipelines/abot_world/before_denoise.py index 3e354a810129..77a0164a837a 100644 --- a/src/diffusers/modular_pipelines/abot_world/before_denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/before_denoise.py @@ -16,6 +16,8 @@ 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 @@ -33,8 +35,8 @@ class ABotWorldPrepareStep(ModularPipelineBlocks): 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, validates the per-block actions, and allocates the " - "transformer's rolling K/V cache with the reference tokens pinned at its head." + "distilled `denoising_timesteps` through it, and allocates the transformer's rolling K/V cache with " + "the reference tokens pinned at its head." ) @property @@ -47,15 +49,6 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "actions", - type_hint=list[list[int]], - description=( - "Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block " - "(W/A/S/D move, I/J/K/L turn the camera); the scripted rollout generates `len(actions)` " - "blocks. Omit when driving the rollout interactively through `loop_step`." - ), - ), InputParam( "denoising_timesteps", type_hint=list[int], @@ -75,13 +68,18 @@ def inputs(self) -> list[InputParam]: @property def intermediate_outputs(self) -> list[OutputParam]: return [ - OutputParam("actions", type_hint=torch.Tensor, description="The actions as a `[num_blocks, 8]` tensor"), 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() @@ -89,13 +87,6 @@ def __call__(self, components, state: PipelineState) -> PipelineState: block_state = self.get_block_state(state) device = components._execution_device - if block_state.actions is not None: - block_state.actions = torch.tensor(block_state.actions, dtype=torch.float32) - if block_state.actions.ndim != 2 or block_state.actions.shape[1] != 8: - raise ValueError( - f"`actions` must be a list of 8-element vectors, got shape {block_state.actions.shape}" - ) - # 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()) @@ -114,6 +105,7 @@ def __call__(self, components, state: PipelineState) -> PipelineState: * (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], diff --git a/src/diffusers/modular_pipelines/abot_world/denoise.py b/src/diffusers/modular_pipelines/abot_world/denoise.py index 01abd2e1f6dd..ceef7e514c44 100644 --- a/src/diffusers/modular_pipelines/abot_world/denoise.py +++ b/src/diffusers/modular_pipelines/abot_world/denoise.py @@ -42,8 +42,9 @@ 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. An interactive driver overwrites `actions` in the " - "state between `loop_step` calls to steer the world live." + "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 @@ -55,12 +56,6 @@ def expected_components(self) -> list[ComponentSpec]: @property def inputs(self) -> list[InputParam]: return [ - InputParam( - "actions", - required=True, - type_hint=torch.Tensor, - description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", - ), 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( @@ -82,11 +77,16 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): block_state = self.get_block_state(state) device = components._execution_device - action = block_state.actions[k].to(device=device, dtype=components.transformer.dtype) + 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) @@ -148,7 +148,7 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): block_state = self.get_block_state(state) device = components._execution_device @@ -315,15 +315,12 @@ def loop_variables(self) -> list[str]: def description(self) -> str: return ( "Pipeline block that denoises one rollout block over the distilled `denoise_timesteps`. It runs inside " - "the rollout loop and reads the current block index `k` from the loop scope." + "the rollout loop and accepts its loop variables (`k`, `action`)." ) @property - def inputs(self) -> list[InputParam]: - inputs = super().inputs - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself, on top of what the sub-blocks declare - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam( "denoise_timesteps", required=True, @@ -331,17 +328,16 @@ def inputs(self) -> list[InputParam]: description="The warped denoising timesteps the loop iterates", ), ] - return [param for param in loop_inputs if param.name not in names] + inputs @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + 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): + 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) @@ -417,7 +413,7 @@ def inputs(self) -> list[InputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): block_state = self.get_block_state(state) device = components._execution_device @@ -486,7 +482,7 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): block_state = self.get_block_state(state) vae = components.vae @@ -551,7 +547,7 @@ def intermediate_outputs(self) -> list[OutputParam]: ] @torch.no_grad() - def __call__(self, components, state: PipelineState, k: int): + def __call__(self, components, state: PipelineState, k: int, action: torch.Tensor): block_state = self.get_block_state(state) tiny_vae = components.tiny_vae @@ -570,73 +566,77 @@ class ABotWorldRolloutWrapper(IterativePipelineBlocks): @property def loop_variables(self) -> list[str]: - return ["k"] + 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. Drive it through `loop_step(components, state, k=k)` to " - "own the iteration — write new `actions` into the state between calls for live interaction." + "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 inputs(self) -> list[InputParam]: - # `decode_cache` is loop-carried — produced by the decode step of the previous iteration, never user-provided - inputs = [param for param in super().inputs if param.name != "decode_cache"] - names = {param.name for param in inputs} - # `actions` is also consumed by the loop logic itself (the rollout length) - loop_inputs = [ + def loop_inputs(self) -> list[InputParam]: + return [ InputParam( "actions", required=True, - type_hint=torch.Tensor, - description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", + 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`." + ), ), ] - return [param for param in loop_inputs if param.name not in names] + inputs @property - def intermediate_outputs(self) -> list[OutputParam]: - # produced by the loop logic itself, which collects each block's decoded frames - return super().intermediate_outputs + [ - OutputParam("videos", type_hint=list[np.ndarray], description="The generated videos"), - ] + 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 block_state.actions is None: + if callable(block_state.actions): raise ValueError( - "A scripted rollout requires `actions` (one vector per block); to drive the rollout " - "interactively, call `loop_step` yourself and write `action` into the state between calls." + "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=block_state.actions.shape[0], desc="Rollout") as progress_bar: - for k in range(block_state.actions.shape[0]): - components, state = self.loop_step(components, state, k=k) + 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() - state.set("videos", [np.concatenate(frames, axis=0)]) + 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 block_state.actions is None: - raise ValueError( - "A scripted rollout requires `actions` (one vector per block); to drive the rollout " - "interactively, call `loop_step` yourself and write `action` into the state between calls." - ) + if callable(block_state.actions): + next_action = block_state.actions + else: + scripted = block_state.actions - frames = [] - for k in range(block_state.actions.shape[0]): - components, state = yield from self.stream_step(components, state, k=k) + 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")) - state.set("videos", [np.concatenate(frames, axis=0)]) + k += 1 + block_state.videos = [np.concatenate(frames, axis=0)] + self.set_block_state(state, block_state) return components, state @@ -660,167 +660,9 @@ def description(self) -> str: ) -class ABotWorldCurrentActionStep(ModularLoopPipelineBlocks): - model_name = "abot-world" - - @property - def description(self) -> str: - return ( - "Step within the streaming rollout loop that broadcasts the *current* `[W, A, S, D, I, J, K, L]` action " - "vector (the `action` state value) into the conditioning planes. The scripted `__call__`/`stream` set " - "`action` from the `actions` list each iteration; a live driver writes it into the state between " - "`loop_step` calls." - ) - - @property - def expected_components(self) -> list[ComponentSpec]: - return [ - ComponentSpec("transformer", ABotWorldTransformer3DModel), - ] - - @property - def inputs(self) -> list[InputParam]: - return [ - InputParam( - "action", - required=True, - type_hint=torch.Tensor, - description="The current block's `[W, A, S, D, I, J, K, L]` 0/1 action vector", - ), - 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): - block_state = self.get_block_state(state) - device = components._execution_device - - action = torch.as_tensor(block_state.action, dtype=torch.float32) - 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 ABotWorldStreamingRolloutWrapper(IterativePipelineBlocks): - model_name = "abot-world" - - @property - def loop_variables(self) -> list[str]: - return ["k"] - - @property - def description(self) -> str: - return ( - "Pipeline block that rolls the world out block by block and decodes each block to pixels inside the " - "loop. Scripted runs iterate the `actions` list (each iteration writes the block's `action` into the " - "state, exactly as a live driver would); interactive drivers own the loop via " - "`loop_step(components, state, k=k)` and write `action` between calls." - ) - - @property - def inputs(self) -> list[InputParam]: - # `action` and `decode_cache` are loop-carried — supplied per iteration by the loop logic (or a live - # driver) and by the decode step of the previous iteration — never user-provided. - inputs = [param for param in super().inputs if param.name not in ("action", "decode_cache")] - names = {param.name for param in inputs} - # inputs consumed by the loop logic itself - loop_inputs = [ - InputParam( - "actions", - type_hint=torch.Tensor, - description="Per-block action vectors `[num_blocks, 8]`, from the prepare step", - ), - InputParam( - "action_source", - type_hint=Callable, - description=( - "Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` " - "polled once per block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, " - "or `None` to stop. The rollout is unbounded while it returns actions." - ), - ), - ] - return [param for param in loop_inputs if param.name not in names] + inputs - - @property - def intermediate_outputs(self) -> list[OutputParam]: - # produced by the loop logic itself, which collects each block's decoded frames - return super().intermediate_outputs + [ - OutputParam("videos", type_hint=list[np.ndarray], description="The generated videos"), - ] - - def _next_action(self, block_state, k): - if block_state.actions is not None: - return block_state.actions[k] if k < block_state.actions.shape[0] else None - if block_state.action_source is not None: - action = block_state.action_source(k) - return None if action is None else torch.as_tensor(action, dtype=torch.float32) - raise ValueError( - "The streaming rollout needs `actions` (a scripted list) or `action_source` (a callable polled per " - "block); to own the loop yourself instead, call `loop_step` and write `action` into the state " - "between calls." - ) - - @torch.no_grad() - def __call__(self, components, state: PipelineState): - block_state = self.get_block_state(state) - if block_state.actions is None and block_state.action_source is not None: - raise ValueError( - "`action_source` needs `stream()`: with `__call__` no frames flow back between polls, so there " - "is nothing to react to. Use `pipe.stream(...)`, or pass a scripted `actions` list instead." - ) - - frames, k = [], 0 - while (action := self._next_action(block_state, k)) is not None: - state.set("action", action) - components, state = self.loop_step(components, state, k=k) - frames.append(state.get("frames")) - k += 1 - state.set("videos", [np.concatenate(frames, axis=0)]) - - return components, state - - @torch.no_grad() - def stream(self, components, state: PipelineState): - block_state = self.get_block_state(state) - - frames, k = [], 0 - while (action := self._next_action(block_state, k)) is not None: - state.set("action", action) - components, state = yield from self.stream_step(components, state, k=k) - frames.append(state.get("frames")) - k += 1 - state.set("videos", [np.concatenate(frames, axis=0)]) - - return components, state - - -class ABotWorldStreamingRolloutStep(ABotWorldStreamingRolloutWrapper): +class ABotWorldStreamingRolloutStep(ABotWorldRolloutWrapper): block_classes = [ - ABotWorldCurrentActionStep, + ABotWorldSetActionStep, ABotWorldPrepareNoiseStep, ABotWorldDenoiseStep, ABotWorldCacheUpdateStep, 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 index f3b14e7dcd2e..c1d4bb17527c 100644 --- a/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py +++ b/src/diffusers/modular_pipelines/abot_world/modular_blocks_abot_world.py @@ -42,10 +42,6 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): (`AutoencoderKLWan`) video_processor (`VideoProcessor`) Inputs: - actions (`list`, *optional*): - Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout - interactively through `loop_step`. 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): @@ -54,6 +50,10 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): 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*): @@ -66,12 +66,13 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): Per-slot validity mask `[B, K]` for the reference views Outputs: - actions (`Tensor`): - The actions as a `[num_blocks, 8]` tensor 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`): @@ -80,8 +81,6 @@ class ABotWorldCoreDenoiseStep(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`WanDecodeCache`): - The VAE's causal-conv cache after this block, carried to the next block videos (`list`): The generated videos """ @@ -127,10 +126,6 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): (`AutoencoderTinyVideo`) video_processor (`VideoProcessor`) Inputs: - actions (`list`, *optional*): - Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout - interactively through `loop_step`. 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): @@ -139,10 +134,10 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): Width of the generated video in pixels reference_latents (`Tensor`): Normalized VAE latents of the reference views `[B, K, C, 1, h, w]` - action_source (`Callable`, *optional*): - Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` polled once per - block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, or `None` to stop. The - rollout is unbounded while it returns actions. + 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*): @@ -155,12 +150,13 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): Per-slot validity mask `[B, K]` for the reference views Outputs: - actions (`Tensor`): - The actions as a `[num_blocks, 8]` tensor 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`): @@ -169,8 +165,6 @@ class ABotWorldStreamingCoreDenoiseStep(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`TinyVideoDecodeCache`): - The tiny VAE's memory after this block, carried to the next block videos (`list`): The generated videos """ @@ -224,16 +218,12 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): without a reference character. reference_resolution (`int`, *optional*, defaults to 512): Side length the reference views are resized to before encoding - actions (`list`, *optional*): - Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout - interactively through `loop_step`. denoising_timesteps (`list`, *optional*, defaults to [1000, 750, 500, 250]): The distilled student's denoising timesteps, before shift-warping - action_source (`Callable`, *optional*): - Interactive alternative to `actions`: a callable `(block_index) -> action vector or None` polled once per - block — return the current `[W, A, S, D, I, J, K, L]` input to keep rolling, or `None` to stop. The - rollout is unbounded while it returns actions. + 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*): @@ -248,12 +238,13 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): 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 - actions (`Tensor`): - The actions as a `[num_blocks, 8]` tensor 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`): @@ -262,8 +253,6 @@ class ABotWorldStreamingBlocks(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`TinyVideoDecodeCache`): - The tiny VAE's memory after this block, carried to the next block videos (`list`): The generated videos """ @@ -310,12 +299,12 @@ class ABotWorldBlocks(SequentialPipelineBlocks): without a reference character. reference_resolution (`int`, *optional*, defaults to 512): Side length the reference views are resized to before encoding - actions (`list`, *optional*): - Per-block actions, one `[W, A, S, D, I, J, K, L]` 0/1 vector per generated block (W/A/S/D move, I/J/K/L - turn the camera); the scripted rollout generates `len(actions)` blocks. Omit when driving the rollout - interactively through `loop_step`. 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*): @@ -330,12 +319,13 @@ class ABotWorldBlocks(SequentialPipelineBlocks): 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 - actions (`Tensor`): - The actions as a `[num_blocks, 8]` tensor 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`): @@ -344,8 +334,6 @@ class ABotWorldBlocks(SequentialPipelineBlocks): Token offset of this block in the rollout: `k * F * tokens_per_frame` frames (`ndarray`): This block's decoded frames `[T, H, W, 3]` - decode_cache (`WanDecodeCache`): - The VAE's causal-conv cache after this block, carried to the next block videos (`list`): The generated videos """