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 <noreply@anthropic.com>

Author: yiyixuxu

scripts/convert_abot_world_to_diffusers.py

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

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,

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,

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