From f8c85e4704c14577d01a8e37080ad2493e5e3d83 Mon Sep 17 00:00:00 2001 From: youngmagician114514 <97871956+youngmagician114514@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:19:01 +0000 Subject: [PATCH 1/2] feat(abot): add single-gpu interactive world pipeline --- README.md | 7 +- docs/en/abot_world.md | 101 +++ docs/en/index.md | 1 + examples/abot_world/README.md | 40 + examples/abot_world/_loader.py | 101 +++ .../abot_world/abot_world_interactive_web.py | 852 ++++++++++++++++++ telefuser/models/abot_world_dit.py | 514 +++++++++++ telefuser/models/wan22_video_vae.py | 2 +- telefuser/pipelines/abot_world/__init__.py | 5 + telefuser/pipelines/abot_world/denoising.py | 190 ++++ telefuser/pipelines/abot_world/interactive.py | 163 ++++ telefuser/pipelines/abot_world/pipeline.py | 143 +++ tests/integration/test_abot_world_smoke.py | 56 ++ tests/unit/models/test_wan22_video_vae.py | 50 + tests/unit/pipelines/abot_world/__init__.py | 0 .../pipelines/abot_world/test_interactive.py | 34 + .../abot_world/test_interactive_web.py | 77 ++ tests/unit/pipelines/abot_world/test_model.py | 113 +++ .../pipelines/abot_world/test_pipeline.py | 31 + 19 files changed, 2477 insertions(+), 3 deletions(-) create mode 100644 docs/en/abot_world.md create mode 100644 examples/abot_world/README.md create mode 100644 examples/abot_world/_loader.py create mode 100644 examples/abot_world/abot_world_interactive_web.py create mode 100644 telefuser/models/abot_world_dit.py create mode 100644 telefuser/pipelines/abot_world/__init__.py create mode 100644 telefuser/pipelines/abot_world/denoising.py create mode 100644 telefuser/pipelines/abot_world/interactive.py create mode 100644 telefuser/pipelines/abot_world/pipeline.py create mode 100644 tests/integration/test_abot_world_smoke.py create mode 100644 tests/unit/models/test_wan22_video_vae.py create mode 100644 tests/unit/pipelines/abot_world/__init__.py create mode 100644 tests/unit/pipelines/abot_world/test_interactive.py create mode 100644 tests/unit/pipelines/abot_world/test_interactive_web.py create mode 100644 tests/unit/pipelines/abot_world/test_model.py create mode 100644 tests/unit/pipelines/abot_world/test_pipeline.py diff --git a/README.md b/README.md index b81ab45b..81cdbd66 100644 --- a/README.md +++ b/README.md @@ -24,6 +24,7 @@ TeleFuser is a high-performance runtime for world model inference and multimodal chunk-boundary time slicing, reconnect-friendly browser transport, and server-push/bidirectional contracts. - ✨ **2026-07-22**: Added [**LingBot-Video**](examples/lingbot_video/README.md) support for Dense and MoE T2I/T2V/TI2V generation, native four-GPU CFG/SP execution, and in-memory MoE refinement. - ✨ **2026-07-15**: Added [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) support for offline generation, interactive WebRTC streaming, and multi-GPU inference. +- ✨ **2026-08-06**: Added [**ABot-World 0.5B-LF**](docs/en/abot_world.md) single-GPU browser interaction with persistent causal KV state and bounded RoPE positions. - ✨ **2026-07-06**: Added external **CacheSeek** latent cache integration for service-mode cross-request reuse. Cache hits can skip the first N denoising steps; the Wan2.2 cache-enabled service example snapshots `[5, 10, 15, 20, 25]` by default. See [docs/en/latent_cache.md](docs/en/latent_cache.md). @@ -220,6 +221,7 @@ telefuser/ | Pipeline | Task | Notes | |----------|------|-------| | `LingBot-World v2` | Bidirectional world-model streaming | LiveKit control loop via [examples/lingbot/lingbot_world_v2_image_to_video_h100.py](examples/lingbot/lingbot_world_v2_image_to_video_h100.py) | +| `ABot-World 0.5B-LF` | Single-GPU interactive world model | Direct browser controller via [examples/abot_world/README.md](examples/abot_world/README.md); no LiveKit required | | `LiveAct` | S2V | Speech-driven talking head generation via [examples/liveact/liveact_s2v_h100.py](examples/liveact/liveact_s2v_h100.py) | | `FlashVSR` | VSR | Streaming video super-resolution via [examples/flashvsr/README.md](examples/flashvsr/README.md) | @@ -257,13 +259,14 @@ See [examples/README.md](examples/README.md) for the example runner and baseline - [docs/en/torch_compile_compatibility.md](docs/en/torch_compile_compatibility.md): compile-related constraints - [docs/en/adding_new_model.md](docs/en/adding_new_model.md): integrating new models - [docs/en/adding_new_example.md](docs/en/adding_new_example.md): authoring examples and pipeline contracts +- [docs/en/abot_world.md](docs/en/abot_world.md): ABot-World single-GPU interactive pipeline, controls, and tests ## Known Limitations - `AdaTaylorCache` is only calibrated for selected model families. - `torch.compile` support is still experimental in parts of the stack. - Some optimized paths require specific GPU architectures and CUDA versions. -- World-model examples such as `LingBot-World v2` require external checkpoints and environment setup. +- World-model examples such as `LingBot-World v2` and `ABot-World` require external checkpoints and environment setup. - Multi-machine deployment exists in the architecture but may require project-specific integration and validation. ## Development @@ -274,7 +277,7 @@ pre-commit install pytest tests/ ``` -See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution workflow and [AGENTS.md](AGENTS.md) for project-specific agent guidance. +See [CONTRIBUTING.md](CONTRIBUTING.md) for contribution workflow and [CLAUDE.md](CLAUDE.md) for repository guidance. ## License diff --git a/docs/en/abot_world.md b/docs/en/abot_world.md new file mode 100644 index 00000000..d7dc326a --- /dev/null +++ b/docs/en/abot_world.md @@ -0,0 +1,101 @@ +# ABot-World 0.5B-LF + +TeleFuser provides a single-GPU, direct-browser integration for the public +ABot-World 0.5B-LF long-forcing checkpoint. The supported entry point is the +native HTTP controller: + +```bash +python examples/abot_world/abot_world_interactive_web.py \ + --model-root /path/to/ABot-World-0-5B-LF \ + --host 127.0.0.1 \ + --port 7860 +``` + +This integration does not use `telefuser stream-serve`, LiveKit, WebRTC, or a +TURN server. The browser speaks to the local HTTP controller. For an SSH +session, forward the port from the remote host: + +```bash +ssh -N -p -L 7860:127.0.0.1:7860 @ +``` + +## Checkpoint And Image + +The loader expects the unmodified checkpoint layout: + +```text +ABot-World-0-5B-LF/ + diffusion_pytorch_model.safetensors + Wan2.2_VAE.pth + models_t5_umt5-xxl-enc-bf16.pth +``` + +The browser's default sample image comes from the official ABot-World web +client asset at `../ABot-World/web_client/datasets/images/84b90ad568b693d2.png`. +Pass a different server-side image path from the page when needed. + +## Pipeline Structure + +`ABotWorldPipeline` is a TeleFuser `BasePipeline`. It uses the existing Wan +VAE and text stages plus the model-specific `ABotWorldDenoisingStage`. +`ABotWorldDiT` uses the public TeleFuser attention operations and the official +four-step x0-prediction causal sampler. + +`ABotWorldInteractivePipeline` retains the prompt embedding, initial image +latent, self/cross KV caches, scheduler, RNG, and VAE temporal cache between +control blocks. The initial integration supports one GPU and one retained +causal session. + +## Controls And Idle Behavior + +WASD and arrow keys control movement. IJKL controls camera rotation. Connect +creates the image-conditioned session and displays the input preview; it does +not advance the DiT with an empty action state. A non-empty control snapshot +starts the next three-latent causal block. Releasing all keys stops new model +execution without discarding frames already queued for playback. + +The browser consumes decoded frames in order at 12 FPS. The bounded FIFO +applies producer backpressure when playback is behind, so normal playback does +not drop generated blocks. + +## KV And RoPE + +The default causal window is 18 latent frames: six sink frames plus a +twelve-frame rolling tail. KV cache entries remain unrotated; RoPE is applied +when keys are read using bounded logical positions. Sink positions are `0..5` +and the rolling tail occupies the remaining local window, so the global session +frame number does not grow the RoPE index past the precomputed table. + +This fixed logical position policy is an intentional difference from the +original non-sink ABot baseline and must be evaluated as part of any future +long-horizon quality claim. + +## Tests + +CPU contract tests cover model conversion, sink KV rolling, RoPE boundaries, +session cleanup, idle behavior, FIFO backpressure, and action layout: + +```bash +python -m pytest tests/unit/models/test_wan22_video_vae.py \ + tests/unit/pipelines/abot_world -q +``` + +The opt-in GPU smoke uses the release checkpoint, the public `480x832` shape, +a fixed seed, and 30 control blocks: + +```bash +CUDA_VISIBLE_DEVICES=0 \ +ABOT_WORLD_MODEL_ROOT=/path/to/ABot-World-0-5B-LF \ +ABOT_WORLD_TEST_IMAGE=/path/to/initial.png \ +python -m pytest -m "gpu and slow" \ + tests/integration/test_abot_world_smoke.py -v -s +``` + +The smoke is a generation and cache contract test. It does not establish +visual quality, prompt fidelity, or parity over an unbounded session. + +## Scope + +The first integration is intentionally single-GPU and direct HTTP. LiveKit or +`stream-serve` support requires a separate transport integration and is not +part of this example. diff --git a/docs/en/index.md b/docs/en/index.md index acee67aa..eb444de4 100644 --- a/docs/en/index.md +++ b/docs/en/index.md @@ -58,6 +58,7 @@ Reusable stages, model configs, schedulers, and pipeline orchestration. |-------|-------|-------------| | LingBot-World v2 | Bidirectional streaming | Camera-controlled interactive world model via LiveKit | | LingBot-World-Fast | Bidirectional streaming | Legacy/causal-fast model via LiveKit reliable data messages | +| ABot-World 0.5B-LF | Single-GPU interactive generation | Direct browser control with persistent causal KV state; see the [ABot guide](abot_world/) | ### Video Generation diff --git a/examples/abot_world/README.md b/examples/abot_world/README.md new file mode 100644 index 00000000..d5edf53e --- /dev/null +++ b/examples/abot_world/README.md @@ -0,0 +1,40 @@ +# ABot-World 0.5B-LF + +This example exposes one local single-GPU entry point: + +```bash +python examples/abot_world/abot_world_interactive_web.py \ + --model-root /path/to/ABot-World-0-5B-LF \ + --host 127.0.0.1 \ + --port 7860 +``` + +The browser controls WASD/arrow movement and IJKL camera rotation. Connecting +creates the image-conditioned causal session but does not advance the DiT +until a non-empty control state is received. Generated blocks remain ordered +in a bounded FIFO and the producer waits when the browser is behind. +The six sink latents and rolling tail use fixed logical RoPE positions, so the +global session frame number does not index beyond the trained local window. + +## Test tiers + +CPU contract tests cover action-channel layout, checkpoint conversion, sink +KV rolling, RoPE boundary validation, session cleanup, and the direct runtime +idle/FIFO behavior: + +```bash +pytest tests/unit/pipelines/abot_world +``` + +The 30-block GPU smoke is opt-in because it loads the release checkpoint: + +```bash +ABOT_WORLD_MODEL_ROOT=/path/to/ABot-World-0-5B-LF \ +ABOT_WORLD_TEST_IMAGE=/path/to/initial.png \ +pytest -m "gpu and slow" tests/integration/test_abot_world_smoke.py -v +``` + +The smoke uses the public 480x832 shape, a fixed seed, and a fixed control +state. It checks that every block decodes frames and that the session's +emitted-frame counter matches the observed count. It is a generation contract +test, not a visual-quality or long-horizon parity claim. diff --git a/examples/abot_world/_loader.py b/examples/abot_world/_loader.py new file mode 100644 index 00000000..3cc3e2cc --- /dev/null +++ b/examples/abot_world/_loader.py @@ -0,0 +1,101 @@ +"""Internal ABot-World checkpoint loader for the interactive example.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch + +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + ModelRuntimeConfig, + OffloadConfig, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.models.abot_world_dit import ABotWorldDiT +from telefuser.models.wan22_video_vae import Wan22VideoVAE +from telefuser.models.wan_video_text_encoder import WanTextEncoder +from telefuser.ops.attention.backends import FLASH_ATTN_3_AVAILABLE, FLASH_ATTN_4_AVAILABLE +from telefuser.pipelines.abot_world import ABotWorldPipeline, ABotWorldPipelineConfig + +_PROJECT_ROOT = Path(__file__).resolve().parents[2] +_DEFAULT_MODEL_ROOT = ( + Path(os.environ.get("TF_MODEL_ZOO_PATH", _PROJECT_ROOT.parent / "model_zoo")) / "ABot-World-0-5B-LF" +) +DEFAULT_PROMPT = "A smooth first-person exploration through a vivid natural landscape." + + +def _attention_backend() -> AttnImplType: + if FLASH_ATTN_4_AVAILABLE: + return AttnImplType.FLASH_ATTN_4 + if FLASH_ATTN_3_AVAILABLE: + return AttnImplType.FLASH_ATTN_3 + return AttnImplType.TORCH_SDPA + + +def get_pipeline( + model_root: str | Path = _DEFAULT_MODEL_ROOT, + *, + height: int = 480, + width: int = 832, + latent_frames: int = 31, + pipeline_class: type[ABotWorldPipeline] = ABotWorldPipeline, +) -> ABotWorldPipeline: + """Load the downloaded ABot checkpoint with VAE/T5 model CPU offload.""" + root = Path(model_root).expanduser() + required = ("diffusion_pytorch_model.safetensors", "Wan2.2_VAE.pth", "models_t5_umt5-xxl-enc-bf16.pth") + missing = [name for name in required if not (root / name).is_file()] + if missing: + raise FileNotFoundError(f"ABot model root {root} is missing: {', '.join(missing)}") + + model_manager = ModuleManager(device="cpu") + model_manager.load_model( + str(root / "Wan2.2_VAE.pth"), + name="wan_video_vae", + model_class=Wan22VideoVAE, + torch_dtype=torch.float32, + low_cpu_mem_usage=True, + ) + model_manager.load_model( + str(root / "models_t5_umt5-xxl-enc-bf16.pth"), + name="wan_video_text_encoder", + model_class=WanTextEncoder, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + model_manager.load_model( + str(root / "diffusion_pytorch_model.safetensors"), + name="abot_world_dit", + model_class=ABotWorldDiT, + torch_dtype=torch.bfloat16, + low_cpu_mem_usage=True, + ) + + cpu_offload = OffloadConfig(offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD) + pipeline = pipeline_class(device="cuda", torch_dtype=torch.bfloat16) + pipeline.init( + model_manager, + ABotWorldPipelineConfig( + vae_config=ModelRuntimeConfig( + device_type="cuda", device_id=0, torch_dtype=torch.float32, offload_config=cpu_offload + ), + text_encoding_config=ModelRuntimeConfig( + device_type="cuda", device_id=0, torch_dtype=torch.bfloat16, offload_config=cpu_offload + ), + dit_config=ModelRuntimeConfig( + device_type="cuda", + device_id=0, + torch_dtype=torch.bfloat16, + attention_config=AttentionConfig.dense_attention(_attention_backend()), + ), + height=height, + width=width, + latent_frames=latent_frames, + local_attn_size=18, + sink_size=6, + ), + ) + return pipeline diff --git a/examples/abot_world/abot_world_interactive_web.py b/examples/abot_world/abot_world_interactive_web.py new file mode 100644 index 00000000..a977f666 --- /dev/null +++ b/examples/abot_world/abot_world_interactive_web.py @@ -0,0 +1,852 @@ +"""Direct-loaded, LingBot-style browser controller for ABot-World. + +The page intentionally uses a small native HTTP server instead of a component +framework: control buttons must reflect keyboard state immediately, while the +GPU worker keeps one causal ABot session alive. Press Connect once, then hold +WASD/arrow keys for movement or IJKL for camera rotation; a new causal block +is scheduled as soon as the previous one completes. +""" + +from __future__ import annotations + +import argparse +import json +import queue +import threading +import time +from http import HTTPStatus +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer +from io import BytesIO +from pathlib import Path +from typing import Any +from urllib.parse import urlparse + +from PIL import Image + +try: + from ._loader import DEFAULT_PROMPT, get_pipeline +except ImportError: # Supports direct execution from examples/abot_world. + from _loader import DEFAULT_PROMPT, get_pipeline + +from telefuser.pipelines.abot_world.interactive import ABotWorldInteractivePipeline, ABotWorldInteractiveSession +from telefuser.utils.video import save_video + +_PROJECT_ROOT = Path(__file__).resolve().parents[2] +_OFFICIAL_SAMPLE = _PROJECT_ROOT.parent / "ABot-World" / "web_client" / "datasets" / "images" / "84b90ad568b693d2.png" +_OUTPUT_DIR = _PROJECT_ROOT / "work_dirs" / "abot_world_interactive" +_ACTION_ORDER = ("W", "A", "S", "D", "I", "J", "K", "L") +_DEFAULT_OUTPUT_QUEUE_SIZE = 4 +_PULL_TIMEOUT_SECONDS = 15.0 + + +class InteractiveRuntime: + """Own one causal ABot session plus a LingBot-style output queue.""" + + def __init__( + self, + pipeline: ABotWorldInteractivePipeline, + fps: int, + control_latent_frames: int, + output_queue_size: int = _DEFAULT_OUTPUT_QUEUE_SIZE, + ) -> None: + if control_latent_frames not in {1, 3}: + raise ValueError("control_latent_frames must be 1 or 3") + if output_queue_size <= 0: + raise ValueError("output_queue_size must be positive") + self.pipeline = pipeline + self.fps = fps + self.control_latent_frames = control_latent_frames + self.output_queue_size = output_queue_size + self.lock = threading.RLock() + self.session: ABotWorldInteractiveSession | None = None + self.frames: list[Image.Image] = [] + self.chunk_index = 0 + self.version = 0 + self._encoded_blocks: dict[int, list[bytes]] = {} + self._current_frame_urls: list[str] = [] + self._output_queue: queue.Queue[dict[str, Any]] = queue.Queue(maxsize=output_queue_size) + self._control_event = threading.Event() + self._stop_event = threading.Event() + self._worker: threading.Thread | None = None + self._active_controls: dict[str, bool] = {} + self._control_revision = 0 + self._worker_error: str | None = None + self._queue_high_watermark = 0 + self._dropped_video_blocks = 0 + self._dropped_video_frames = 0 + self._producer_backpressure_events = 0 + self._producer_backpressure_seconds = 0.0 + self.latest_frame_path = _OUTPUT_DIR / "current_frame.jpg" + self.video_path = _OUTPUT_DIR / "current_session.mp4" + + @staticmethod + def _actions(raw_controls: object) -> dict[str, bool]: + if raw_controls is None: + controls: list[object] = [] + elif isinstance(raw_controls, list): + controls = raw_controls + else: + raise ValueError("controls must be an array") + pressed = {str(control).upper() for control in controls} + unknown = pressed.difference(_ACTION_ORDER) + if unknown: + raise ValueError("Unsupported controls: " + ", ".join(sorted(unknown))) + return {control: True for control in pressed} + + def _queue_metrics(self) -> dict[str, int | float]: + return { + "queued_chunks": self._output_queue.qsize(), + "queue_capacity": self.output_queue_size, + "queue_high_watermark": self._queue_high_watermark, + "dropped_video_blocks": self._dropped_video_blocks, + "dropped_video_frames": self._dropped_video_frames, + "producer_backpressure_events": self._producer_backpressure_events, + "producer_backpressure_seconds": round(self._producer_backpressure_seconds, 3), + } + + def _clear_output_queue(self) -> None: + while True: + try: + self._output_queue.get_nowait() + except queue.Empty: + return + + def _cache_block_frames(self, block: list[Image.Image]) -> None: + if not block: + raise RuntimeError("No decoded ABot frames are available") + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + self.version += 1 + encoded: list[bytes] = [] + for frame in block: + buffer = BytesIO() + frame.save(buffer, format="JPEG", quality=92) + encoded.append(buffer.getvalue()) + self._encoded_blocks[self.version] = encoded + for stale_version in sorted(self._encoded_blocks)[:-12]: + del self._encoded_blocks[stale_version] + block[-1].save(self.latest_frame_path, format="JPEG", quality=92) + self._current_frame_urls = [f"/api/block-frame/{self.version}/{index}" for index in range(len(encoded))] + + def _result( + self, + *, + new_frames: int, + controls: dict[str, bool], + status: str, + event_type: str, + control_revision: int, + ) -> dict[str, Any]: + return { + "type": event_type, + "chunk": self.chunk_index, + "new_frames": new_frames, + "total_frames": len(self.frames), + "controls": sorted(controls), + "control_revision": control_revision, + "frame_url": f"/api/frame?v={self.version}", + "frame_urls": list(self._current_frame_urls), + "status": status, + "queue": self._queue_metrics(), + } + + def _enqueue_video_output(self, payload: dict[str, Any]) -> bool: + """Block the producer when the browser has not drained its FIFO. + + This is deliberate backpressure: normal streaming never discards a + generated ABot block merely because the browser is temporarily ahead + of its 12 FPS playback clock. + """ + blocked_started_at: float | None = None + while not self._stop_event.is_set(): + try: + self._output_queue.put(payload, timeout=0.1) + break + except queue.Full: + if blocked_started_at is None: + blocked_started_at = time.monotonic() + with self.lock: + self._producer_backpressure_events += 1 + continue + else: + return False + + if blocked_started_at is not None: + with self.lock: + self._producer_backpressure_seconds += time.monotonic() - blocked_started_at + with self.lock: + self._queue_high_watermark = max(self._queue_high_watermark, self._output_queue.qsize()) + return True + + def _worker_loop(self) -> None: + while not self._stop_event.is_set(): + self._control_event.wait(timeout=0.25) + if self._stop_event.is_set(): + return + with self.lock: + session = self.session + controls = dict(self._active_controls) + revision = self._control_revision + if session is None: + return + if not controls: + continue + try: + block = self.pipeline.generate_next_block(session, controls, self.control_latent_frames) + except Exception as error: + with self.lock: + self._worker_error = str(error) + return + with self.lock: + if session is not self.session: + return + self.frames.extend(block) + self.chunk_index += 1 + self._cache_block_frames(block) + payload = self._result( + new_frames=len(block), + controls=controls, + status="Causal block generated by the background producer.", + event_type="chunk", + control_revision=revision, + ) + should_publish = not self._stop_event.is_set() + if should_publish and not self._enqueue_video_output(payload): + return + + def _stop_worker(self) -> None: + with self.lock: + self._stop_event.set() + self._control_event.set() + worker = self._worker + if worker is not None and worker.is_alive() and worker is not threading.current_thread(): + worker.join(timeout=30.0) + if worker is not None and worker.is_alive(): + raise RuntimeError("Timed out while stopping the ABot background worker") + with self.lock: + if self._worker is worker: + self._worker = None + + def _start_worker_locked(self) -> None: + self._stop_event.clear() + worker = threading.Thread(target=self._worker_loop, daemon=True, name="abot-world-producer") + self._worker = worker + worker.start() + + def start(self, image_path: str, prompt: str, seed: int, raw_controls: object) -> dict[str, Any]: + path = Path(image_path).expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(f"Image does not exist on the server: {path}") + if not prompt.strip(): + raise ValueError("Prompt must not be empty") + controls = self._actions(raw_controls) + self._stop_worker() + with self.lock: + if self.session is not None: + self.pipeline.close_interactive_session(self.session) + self._clear_output_queue() + self._encoded_blocks.clear() + self._current_frame_urls = [] + self.frames = [] + self.chunk_index = 0 + self.version = 0 + self._queue_high_watermark = 0 + self._dropped_video_blocks = 0 + self._dropped_video_frames = 0 + self._producer_backpressure_events = 0 + self._producer_backpressure_seconds = 0.0 + self._worker_error = None + self._control_revision = 0 + self._active_controls = controls + with Image.open(path) as source: + image = source.convert("RGB") + self.session = self.pipeline.create_interactive_session(image, prompt.strip(), seed=seed) + # Session creation prepares causal state only. An explicit + # non-empty control snapshot must drive the first DiT block. + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + image.save(self.latest_frame_path, format="JPEG", quality=92) + result = self._result( + new_frames=0, + controls=controls, + status="Causal session ready; waiting for control input.", + event_type="preview", + control_revision=self._control_revision, + ) + self._start_worker_locked() + if controls: + self._control_event.set() + return result + + def set_controls(self, raw_controls: object, raw_revision: object = None) -> dict[str, Any]: + controls = self._actions(raw_controls) + revision = None if raw_revision is None else int(raw_revision) + with self.lock: + if self.session is None: + raise RuntimeError("No active session; press Connect first") + if revision is not None and revision < self._control_revision: + return { + "type": "status", + "stage": "stale_control_ignored", + "controls": sorted(self._active_controls), + "control_revision": self._control_revision, + "queue": self._queue_metrics(), + } + self._active_controls = controls + self._control_revision = self._control_revision + 1 if revision is None else revision + result = { + "type": "status", + "stage": "control_state", + "controls": sorted(controls), + "control_revision": self._control_revision, + "queue": self._queue_metrics(), + } + if controls: + self._control_event.set() + else: + self._control_event.clear() + return result + + def pull_output(self, timeout_seconds: float = _PULL_TIMEOUT_SECONDS) -> dict[str, Any]: + timeout = min(max(float(timeout_seconds), 0.0), _PULL_TIMEOUT_SECONDS) + try: + payload = dict(self._output_queue.get(timeout=timeout)) + except queue.Empty: + with self.lock: + stage = "waiting_for_chunk" if self._active_controls else "waiting_for_input" + return { + "type": "status", + "stage": stage, + "controls": sorted(self._active_controls), + "control_revision": self._control_revision, + "queue": self._queue_metrics(), + "worker_error": self._worker_error, + } + with self.lock: + payload["queue"] = self._queue_metrics() + return payload + + def next(self, raw_controls: object) -> dict[str, Any]: + self.set_controls(raw_controls) + return self.pull_output() + + def stop(self) -> dict[str, Any]: + self._stop_worker() + with self.lock: + if self.session is not None: + self.pipeline.close_interactive_session(self.session) + self.session = None + self._active_controls = {} + self._clear_output_queue() + if self.frames: + _OUTPUT_DIR.mkdir(parents=True, exist_ok=True) + save_video(self.frames, str(self.video_path), fps=self.fps, quality=8) + return { + "status": "Stopped. The output queue was drained; GPU model weights remain loaded.", + "video_url": "/api/video" if self.video_path.is_file() else None, + "queue": self._queue_metrics(), + } + + def frame_bytes(self) -> bytes | None: + with self.lock: + return self.latest_frame_path.read_bytes() if self.latest_frame_path.is_file() else None + + def block_frame_bytes(self, version: int, index: int) -> bytes | None: + with self.lock: + block = self._encoded_blocks.get(version) + if block is None or not 0 <= index < len(block): + return None + return block[index] + + def video_bytes(self) -> bytes | None: + with self.lock: + return self.video_path.read_bytes() if self.video_path.is_file() else None + + +_HTML = r""" + + + + + ABot-World Single-GPU Controller + + + +
+

ABot-World Single-GPU Interactive Demo

+
+
+

Server Output

GPU weights are loaded. Ready.
+ Current ABot world frame +
Connect starts a persistent causal session. The browser prebuffers one to two blocks, then consumes frames in order at 12 FPS. A full server FIFO pauses the producer instead of dropping frames.
+
+ +
+
+ + + +""" + + +def _render_html() -> bytes: + return ( + _HTML.replace("__DEFAULT_IMAGE_PATH__", json.dumps(str(_OFFICIAL_SAMPLE))) + .replace("__DEFAULT_PROMPT__", json.dumps(DEFAULT_PROMPT)) + .encode("utf-8") + ) + + +def _make_handler(runtime: InteractiveRuntime) -> type[BaseHTTPRequestHandler]: + class Handler(BaseHTTPRequestHandler): + def _send(self, status: HTTPStatus, payload: bytes, content_type: str) -> None: + self.send_response(status) + self.send_header("Content-Type", content_type) + self.send_header("Content-Length", str(len(payload))) + self.send_header("Cache-Control", "no-store") + self.end_headers() + self.wfile.write(payload) + + def _send_json(self, status: HTTPStatus, payload: dict[str, Any]) -> None: + self._send(status, json.dumps(payload).encode("utf-8"), "application/json; charset=utf-8") + + def _read_json(self) -> dict[str, Any]: + length = int(self.headers.get("Content-Length") or "0") + if length > 1_000_000: + raise ValueError("Request body is too large") + value = json.loads(self.rfile.read(length) or b"{}") + if not isinstance(value, dict): + raise ValueError("JSON request body must be an object") + return value + + def do_GET(self) -> None: # noqa: N802 + path = urlparse(self.path).path + if path == "/": + self._send(HTTPStatus.OK, _render_html(), "text/html; charset=utf-8") + elif path == "/sample-image": + self._send(HTTPStatus.OK, _OFFICIAL_SAMPLE.read_bytes(), "image/jpeg") + elif path.startswith("/api/block-frame/"): + parts = path.split("/") + if len(parts) != 5: + self.send_error(HTTPStatus.NOT_FOUND) + return + try: + version, index = int(parts[3]), int(parts[4]) + except ValueError: + self.send_error(HTTPStatus.NOT_FOUND) + return + payload = runtime.block_frame_bytes(version, index) + if payload is None: + self.send_error(HTTPStatus.NOT_FOUND, "Frame is no longer available") + else: + self._send(HTTPStatus.OK, payload, "image/jpeg") + elif path == "/api/frame": + payload = runtime.frame_bytes() + if payload is None: + self.send_error(HTTPStatus.NO_CONTENT) + else: + self._send(HTTPStatus.OK, payload, "image/jpeg") + elif path == "/api/video": + payload = runtime.video_bytes() + if payload is None: + self.send_error(HTTPStatus.NOT_FOUND, "No stopped session video is available") + else: + self._send(HTTPStatus.OK, payload, "video/mp4") + else: + self.send_error(HTTPStatus.NOT_FOUND) + + def do_POST(self) -> None: # noqa: N802 + path = urlparse(self.path).path + try: + payload = self._read_json() + if path == "/api/start": + result = runtime.start( + str(payload.get("image_path") or ""), + str(payload.get("prompt") or ""), + int(payload.get("seed", 42)), + payload.get("controls"), + ) + elif path == "/api/control": + result = runtime.set_controls(payload.get("controls"), payload.get("revision")) + elif path == "/api/pull": + result = runtime.pull_output(float(payload.get("timeout_ms", 15_000)) / 1000.0) + elif path == "/api/next": + result = runtime.next(payload.get("controls")) + elif path == "/api/stop": + result = runtime.stop() + else: + self.send_error(HTTPStatus.NOT_FOUND) + return + except (FileNotFoundError, RuntimeError, ValueError, OSError) as error: + self._send_json(HTTPStatus.BAD_REQUEST, {"detail": str(error)}) + return + except Exception as error: # pragma: no cover - protects the local HTTP boundary. + self._send_json(HTTPStatus.INTERNAL_SERVER_ERROR, {"detail": str(error)}) + return + self._send_json(HTTPStatus.OK, result) + + def log_message(self, _format: str, *_args: object) -> None: + return + + return Handler + + +def main() -> None: + parser = argparse.ArgumentParser(description="Direct-loaded ABot-World single-GPU interactive controller") + parser.add_argument("--model-root", default=None) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--latent-frames", type=int, default=31) + parser.add_argument("--fps", type=int, default=12) + parser.add_argument( + "--control-latent-frames", + type=int, + choices=(1, 3), + default=3, + help="Causal latents per control update: 3 matches the official ABot streaming checkpoint; 1 is experimental.", + ) + parser.add_argument("--output-queue-size", type=int, default=_DEFAULT_OUTPUT_QUEUE_SIZE) + parser.add_argument("--host", default="127.0.0.1") + parser.add_argument("--port", type=int, default=7860) + args = parser.parse_args() + if args.height % 32 or args.width % 32: + parser.error("height and width must be divisible by 32") + if args.latent_frames < 1 or (args.latent_frames - 1) % 3: + parser.error("latent-frames must be positive and equal to 1 mod 3") + model_root = args.model_root or str(_PROJECT_ROOT.parent / "model_zoo" / "ABot-World-0-5B-LF") + print("Loading VAE, T5, and ABot DiT onto GPU before accepting browser requests.", flush=True) + pipeline = get_pipeline( + model_root, + height=args.height, + width=args.width, + latent_frames=args.latent_frames, + pipeline_class=ABotWorldInteractivePipeline, + ) + pipeline.preload_models() + runtime = InteractiveRuntime( + pipeline, + args.fps, + args.control_latent_frames, + output_queue_size=args.output_queue_size, + ) + server = ThreadingHTTPServer((args.host, args.port), _make_handler(runtime)) + print(f"ABot controller ready at http://{args.host}:{args.port}", flush=True) + try: + server.serve_forever() + except KeyboardInterrupt: + print("Stopping ABot controller.", flush=True) + finally: + server.server_close() + runtime.stop() + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/telefuser/models/abot_world_dit.py b/telefuser/models/abot_world_dit.py new file mode 100644 index 00000000..98ff3920 --- /dev/null +++ b/telefuser/models/abot_world_dit.py @@ -0,0 +1,514 @@ +"""ABot-World causal Wan 2.2 DiT. + +This is the inference-only, single-card implementation for the public +``ABot-World-0-5B-LF`` checkpoint. Its parameter names intentionally match +the official checkpoint. The causal cache stores unrotated keys and applies +RoPE at read time. A sink configuration uses fixed logical positions for the +sink prefix and rolling tail, so long sessions do not grow RoPE indices. +""" + +from __future__ import annotations + +import math +from typing import Any + +import torch +import torch.nn as nn +from einops import rearrange + +from telefuser.core.base_model import BaseModel +from telefuser.core.config import AttentionConfig +from telefuser.ops.attention import attention as attention_fn +from telefuser.ops.normalization import LayerNorm, RMSNorm + +from .wan_video_dit import precompute_freqs_cis_3d, sinusoidal_embedding_1d + + +def _rope_apply( + x: torch.Tensor, + grid_size: tuple[int, int, int], + freqs: torch.Tensor, + frame_indices: torch.Tensor, +) -> torch.Tensor: + """Apply Wan's 3D complex RoPE to ``[B, S, heads, head_dim]`` tensors.""" + frames, height, width = grid_size + sequence_length = frames * height * width + if x.shape[1] != sequence_length: + raise ValueError(f"RoPE expected {sequence_length} tokens, got {x.shape[1]}") + half_dim = x.shape[-1] // 2 + time_dim = half_dim - 2 * (half_dim // 3) + height_dim = half_dim // 3 + width_dim = half_dim // 3 + freq_t, freq_h, freq_w = freqs.split([time_dim, height_dim, width_dim], dim=1) + indices = frame_indices.to(device=x.device, dtype=torch.long) + if indices.numel() != frames: + raise ValueError(f"RoPE frame index count {indices.numel()} does not match {frames}") + if indices.min().item() < 0 or indices.max().item() >= freq_t.shape[0]: + raise ValueError(f"RoPE frame indices must be in [0, {freq_t.shape[0]}), got {indices.tolist()}") + expanded = torch.cat( + [ + freq_t[indices].view(frames, 1, 1, -1).expand(frames, height, width, -1), + freq_h[:height].view(1, height, 1, -1).expand(frames, height, width, -1), + freq_w[:width].view(1, 1, width, -1).expand(frames, height, width, -1), + ], + dim=-1, + ).reshape(sequence_length, 1, -1) + rotated = torch.view_as_real( + torch.view_as_complex(x.float().reshape(x.shape[0], sequence_length, x.shape[2], -1, 2)) * expanded + ).flatten(3) + return rotated.to(dtype=x.dtype) + + +class _ResidualBlock(nn.Module): + def __init__(self, dim: int) -> None: + 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: torch.Tensor) -> torch.Tensor: + return x + self.conv2(self.relu(self.conv1(x))) + + +class SimpleAdapter(nn.Module): + """Official ABot action adapter from pixel-space key presses to DiT tokens.""" + + def __init__( + self, + in_dim: int, + out_dim: int, + kernel_size: tuple[int, int], + stride: tuple[int, int], + downscale_factor: int = 16, + ) -> None: + super().__init__() + self.pixel_unshuffle = nn.PixelUnshuffle(downscale_factor=downscale_factor) + self.conv = nn.Conv2d( + in_dim * downscale_factor * downscale_factor, + out_dim, + kernel_size=kernel_size, + stride=stride, + padding=0, + ) + self.residual_blocks = nn.Sequential(_ResidualBlock(out_dim)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + batch, channels, frames, height, width = x.shape + pixels = x.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width) + features = self.residual_blocks(self.conv(self.pixel_unshuffle(pixels))) + return features.reshape(batch, frames, *features.shape[1:]).permute(0, 2, 1, 3, 4) + + +class CausalWanSelfAttention(nn.Module): + """Single-GPU causal self-attention with ABot's rolling cache semantics.""" + + def __init__( + self, + dim: int, + num_heads: int, + local_attn_size: int = -1, + sink_size: int = 0, + eps: float = 1e-6, + use_relative_rope: bool = True, + ) -> None: + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.use_relative_rope = use_relative_rope + self.attention_config = AttentionConfig() + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + + @staticmethod + def _cursor(cache: dict[str, Any], name: str) -> int: + value = cache[name] + return int(value.item()) if isinstance(value, torch.Tensor) else int(value) + + @staticmethod + def _set_cursor(cache: dict[str, Any], name: str, value: int) -> None: + cursor = cache[name] + if isinstance(cursor, torch.Tensor): + cursor.fill_(value) + else: + cache[name] = value + + def _update_cache( + self, + cache: dict[str, Any], + key: torch.Tensor, + value: torch.Tensor, + current_start: int, + frame_tokens: int, + ) -> tuple[int, int]: + """Store a block and return its local start/end positions.""" + token_count = key.shape[1] + current_end = current_start + token_count + global_end = self._cursor(cache, "global_end_index") + local_end = self._cursor(cache, "local_end_index") + capacity = cache["k"].shape[1] + sink_tokens = self.sink_size * frame_tokens + if self.local_attn_size != -1 and current_end > global_end and token_count + local_end > capacity: + evicted = token_count + local_end - capacity + rolled = local_end - evicted - sink_tokens + if rolled < 0: + raise RuntimeError("ABot KV cache is smaller than its causal attention window") + cache["k"][:, sink_tokens : sink_tokens + rolled].copy_( + cache["k"][:, sink_tokens + evicted : sink_tokens + evicted + rolled].clone() + ) + cache["v"][:, sink_tokens : sink_tokens + rolled].copy_( + cache["v"][:, sink_tokens + evicted : sink_tokens + evicted + rolled].clone() + ) + local_end = local_end + current_end - global_end - evicted + else: + local_end = local_end + current_end - global_end + local_start = local_end - token_count + if local_start < 0 or local_end > capacity: + raise RuntimeError(f"ABot KV cache write [{local_start}:{local_end}] exceeds capacity {capacity}") + cache["k"][:, local_start:local_end].copy_(key.detach()) + cache["v"][:, local_start:local_end].copy_(value) + self._set_cursor(cache, "global_end_index", current_end) + self._set_cursor(cache, "local_end_index", local_end) + return local_start, local_end + + def forward( + self, + x: torch.Tensor, + grid_size: tuple[int, int, int], + freqs: torch.Tensor, + kv_cache: dict[str, Any], + current_start: int, + ) -> torch.Tensor: + batch, tokens, _ = x.shape + frames, height, width = grid_size + frame_tokens = height * width + if tokens != frames * frame_tokens: + raise ValueError("ABot causal attention received inconsistent grid size") + query = rearrange(self.norm_q(self.q(x)), "b s (h d) -> b s h d", h=self.num_heads) + key = rearrange(self.norm_k(self.k(x)), "b s (h d) -> b s h d", h=self.num_heads) + value = rearrange(self.v(x), "b s (h d) -> b s h d", h=self.num_heads) + local_start, local_end = self._update_cache(kv_cache, key, value, current_start, frame_tokens) + # A non-zero sink occupies the prefix of the fixed KV allocation. The + # remainder has already been compacted by _update_cache into a rolling + # tail, so both pieces must stay visible to attention. + if self.local_attn_size == -1 or self.sink_size: + visible_start = 0 + else: + visible_start = max(0, local_end - self.local_attn_size * frame_tokens) + visible_tokens = local_end - visible_start + if visible_tokens % frame_tokens: + raise RuntimeError("ABot causal KV cache lost frame alignment") + visible_frames = visible_tokens // frame_tokens + cached_key = kv_cache["k"][:, visible_start:local_end] + cached_value = kv_cache["v"][:, visible_start:local_end] + if self.use_relative_rope: + if self.sink_size: + # The cache layout is always [sink, rolling tail]. Keep all + # positions inside this trained local window; raw K is rotated + # afresh after each eviction. + cache_indices = torch.arange(visible_frames, device=x.device) + cached_key = _rope_apply(cached_key, (visible_frames, height, width), freqs, cache_indices) + query_start = local_start // frame_tokens + query = _rope_apply( + query, grid_size, freqs, torch.arange(query_start, query_start + frames, device=x.device) + ) + else: + cached_key = _rope_apply( + cached_key, + (visible_frames, height, width), + freqs, + torch.arange(visible_frames, device=x.device), + ) + query_start = visible_frames - frames + if query_start < 0: + raise RuntimeError("ABot query block is larger than its visible causal window") + query = _rope_apply( + query, + grid_size, + freqs, + torch.arange(query_start, visible_frames, device=x.device), + ) + else: + start_frame = current_start // frame_tokens + query = _rope_apply( + query, grid_size, freqs, torch.arange(start_frame, start_frame + frames, device=x.device) + ) + output = attention_fn( + query, + cached_key, + cached_value, + attention_config=self.attention_config, + input_layout="BSND", + output_layout="BSND", + ) + return self.o(rearrange(output, "b s h d -> b s (h d)")) + + +class CausalWanCrossAttention(nn.Module): + def __init__(self, dim: int, num_heads: int, eps: float = 1e-6) -> None: + super().__init__() + self.num_heads = num_heads + self.q = nn.Linear(dim, dim) + self.k = nn.Linear(dim, dim) + self.v = nn.Linear(dim, dim) + self.o = nn.Linear(dim, dim) + self.norm_q = RMSNorm(dim, eps=eps) + self.norm_k = RMSNorm(dim, eps=eps) + self.attention_config = AttentionConfig() + + def forward(self, x: torch.Tensor, context: torch.Tensor, cache: dict[str, Any]) -> torch.Tensor: + query = rearrange(self.norm_q(self.q(x)), "b s (h d) -> b s h d", h=self.num_heads) + if not bool(cache["is_init"]): + key = rearrange(self.norm_k(self.k(context)), "b s (h d) -> b s h d", h=self.num_heads) + value = rearrange(self.v(context), "b s (h d) -> b s h d", h=self.num_heads) + if cache["k"].shape[1] < key.shape[1]: + raise ValueError("ABot cross-attention cache is shorter than the text sequence") + cache["k"][:, : key.shape[1]].copy_(key) + cache["v"][:, : value.shape[1]].copy_(value) + cache["sequence_length"] = key.shape[1] + cache["is_init"] = True + length = int(cache["sequence_length"]) + output = attention_fn( + query, + cache["k"][:, :length], + cache["v"][:, :length], + attention_config=self.attention_config, + input_layout="BSND", + output_layout="BSND", + ) + return self.o(rearrange(output, "b s h d -> b s (h d)")) + + +class CausalWanAttentionBlock(nn.Module): + def __init__( + self, + dim: int, + ffn_dim: int, + num_heads: int, + local_attn_size: int, + sink_size: int, + eps: float, + use_relative_rope: bool, + ) -> None: + super().__init__() + self.norm1 = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.self_attn = CausalWanSelfAttention( + dim, num_heads, local_attn_size, sink_size, eps, use_relative_rope=use_relative_rope + ) + self.norm3 = LayerNorm(dim, eps=eps) + self.cross_attn = CausalWanCrossAttention(dim, num_heads, eps) + self.norm2 = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.ffn = nn.Sequential(nn.Linear(dim, ffn_dim), nn.GELU(approximate="tanh"), nn.Linear(ffn_dim, dim)) + self.modulation = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) + + @staticmethod + def _expand(value: torch.Tensor, frames: int, frame_tokens: int) -> torch.Tensor: + return value.unflatten(1, (frames, frame_tokens)).flatten(1, 2) + + def forward( + self, + x: torch.Tensor, + time_modulation: torch.Tensor, + context: torch.Tensor, + grid_size: tuple[int, int, int], + freqs: torch.Tensor, + kv_cache: dict[str, Any], + crossattn_cache: dict[str, Any], + current_start: int, + ) -> torch.Tensor: + frames, height, width = grid_size + frame_tokens = height * width + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = ( + self.modulation.to(device=x.device, dtype=x.dtype).unsqueeze(0) + time_modulation + ).chunk(6, dim=2) + normed = self.norm1(x).unflatten(1, (frames, frame_tokens)) + attention_input = (normed * (1 + scale_msa) + shift_msa).flatten(1, 2) + attention = self.self_attn(attention_input, grid_size, freqs, kv_cache, current_start) + x = x + (attention.unflatten(1, (frames, frame_tokens)) * gate_msa).flatten(1, 2) + x = x + self.cross_attn(self.norm3(x), context, crossattn_cache) + normed = self.norm2(x).unflatten(1, (frames, frame_tokens)) + ffn_input = (normed * (1 + scale_mlp) + shift_mlp).flatten(1, 2) + ffn_output = self.ffn(ffn_input) + return x + (ffn_output.unflatten(1, (frames, frame_tokens)) * gate_mlp).flatten(1, 2) + + +class CausalHead(nn.Module): + def __init__(self, dim: int, out_dim: int, patch_size: tuple[int, int, int], eps: float) -> None: + super().__init__() + self.patch_size = patch_size + self.norm = LayerNorm(dim, eps=eps, elementwise_affine=False) + self.head = nn.Linear(dim, out_dim * math.prod(patch_size)) + self.modulation = nn.Parameter(torch.randn(1, 2, dim) / dim**0.5) + + def forward(self, x: torch.Tensor, time_embedding: torch.Tensor) -> torch.Tensor: + frames = time_embedding.shape[1] + if x.shape[1] % frames: + raise ValueError("ABot head requires an integral number of tokens per frame") + tokens_per_frame = x.shape[1] // frames + shift, scale = (self.modulation.to(device=x.device, dtype=x.dtype).unsqueeze(1) + time_embedding).chunk( + 2, dim=2 + ) + value = self.norm(x).unflatten(1, (frames, tokens_per_frame)) * (1 + scale) + shift + return self.head(value) + + +class ABotWorldDiT(BaseModel): + """Checkpoint-compatible ABot-World 0.5B long-forcing backbone.""" + + def __init__( + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + text_len: int = 512, + in_dim: int = 48, + dim: int = 3072, + ffn_dim: int = 14336, + freq_dim: int = 256, + text_dim: int = 4096, + out_dim: int = 48, + num_heads: int = 24, + num_layers: int = 30, + local_attn_size: int = 18, + sink_size: int = 6, + eps: float = 1e-6, + use_relative_rope: bool = True, + downscale_factor_control_adapter: int = 16, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.text_len = text_len + self.in_dim = in_dim + self.dim = dim + self.num_heads = num_heads + self.num_layers = num_layers + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.freq_dim = freq_dim + self.out_dim = out_dim + self.use_relative_rope = use_relative_rope + self.layer_name_list = ["blocks"] + self.patch_embedding = nn.Conv3d(in_dim, dim, kernel_size=patch_size, stride=patch_size) + self.text_embedding = nn.Sequential(nn.Linear(text_dim, dim), nn.GELU(approximate="tanh"), nn.Linear(dim, dim)) + self.time_embedding = nn.Sequential(nn.Linear(freq_dim, dim), nn.SiLU(), nn.Linear(dim, dim)) + self.time_projection = nn.Sequential(nn.SiLU(), nn.Linear(dim, dim * 6)) + self.blocks = nn.ModuleList( + [ + CausalWanAttentionBlock(dim, ffn_dim, num_heads, local_attn_size, sink_size, eps, use_relative_rope) + for _ in range(num_layers) + ] + ) + self.head = CausalHead(dim, out_dim, patch_size, eps) + self.act_control_adapter = SimpleAdapter( + 32, + dim, + kernel_size=patch_size[1:], + stride=patch_size[1:], + downscale_factor=downscale_factor_control_adapter, + ) + self.freqs = torch.cat(precompute_freqs_cis_3d(dim // num_heads), dim=1) + self._freqs_by_device: dict[tuple[str, int | None], torch.Tensor] = {} + + def set_causal_attention_window(self, local_attn_size: int, sink_size: int = 0) -> None: + if local_attn_size < 1: + raise ValueError("ABot single-card inference requires a positive local_attn_size") + if not 0 <= sink_size < local_attn_size: + raise ValueError("ABot sink_size must be non-negative and smaller than local_attn_size") + self.local_attn_size = local_attn_size + self.sink_size = sink_size + for block in self.blocks: + block.self_attn.local_attn_size = local_attn_size + block.self_attn.sink_size = sink_size + + def set_attention_config(self, attention_config: AttentionConfig) -> None: + for block in self.blocks: + block.self_attn.attention_config = attention_config + block.cross_attn.attention_config = attention_config + + def _frequencies(self, device: torch.device) -> torch.Tensor: + key = (device.type, device.index) + if key not in self._freqs_by_device: + self._freqs_by_device[key] = self.freqs.to(device=device) + return self._freqs_by_device[key] + + def _unpatchify(self, x: torch.Tensor, grid_size: tuple[int, int, int]) -> torch.Tensor: + frames, height, width = grid_size + patch_t, patch_h, patch_w = self.patch_size + if x.ndim == 4: + # CausalHead retains [frame, tokens-per-frame] for frame-wise modulation. + x = x.flatten(1, 2) + return rearrange( + x, + "b (f h w) (pt ph pw c) -> b c (f pt) (h ph) (w pw)", + f=frames, + h=height, + w=width, + pt=patch_t, + ph=patch_h, + pw=patch_w, + c=self.out_dim, + ) + + def forward( + self, + x: torch.Tensor, + timestep: torch.Tensor, + context: torch.Tensor, + act_context: torch.Tensor, + kv_cache: list[dict[str, Any]], + crossattn_cache: list[dict[str, Any]], + current_start: int, + act_context_scale: float = 1.0, + ) -> torch.Tensor: + if x.ndim != 5 or timestep.ndim != 2: + raise ValueError("ABot expects x=[B,C,F,H,W] and timestep=[B,F]") + if x.shape[2] != timestep.shape[1] or x.shape[0] != timestep.shape[0]: + raise ValueError("ABot timestep shape must match the latent batch and frame dimensions") + if len(kv_cache) != self.num_layers or len(crossattn_cache) != self.num_layers: + raise ValueError("ABot cache lists must contain one entry per transformer layer") + embedded = self.patch_embedding(x) + action = self.act_control_adapter(act_context.to(device=x.device, dtype=embedded.dtype)) + if action.shape != embedded.shape: + raise ValueError( + f"ABot action adapter output {tuple(action.shape)} does not match latent tokens {tuple(embedded.shape)}" + ) + embedded = embedded + action * act_context_scale + grid_size = tuple(int(value) for value in embedded.shape[2:]) + tokens = rearrange(embedded, "b c f h w -> b (f h w) c") + time_embedding = self.time_embedding( + sinusoidal_embedding_1d(self.freq_dim, timestep.flatten()).to(tokens.dtype) + ) + time_embedding = time_embedding.unflatten(0, timestep.shape) + time_modulation = self.time_projection(time_embedding).unflatten(2, (6, self.dim)) + context = self.text_embedding(context) + freqs = self._frequencies(tokens.device) + for index, block in enumerate(self.blocks): + tokens = block( + tokens, + time_modulation, + context, + grid_size, + freqs, + kv_cache[index], + crossattn_cache[index], + current_start, + ) + return self._unpatchify(self.head(tokens, time_embedding.unsqueeze(2)), grid_size) + + @staticmethod + def state_dict_converter() -> "ABotWorldDiTStateDictConverter": + return ABotWorldDiTStateDictConverter() + + +class ABotWorldDiTStateDictConverter: + """The published ABot safetensor already uses native parameter names.""" + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + return state_dict, {} + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + return state_dict, {} diff --git a/telefuser/models/wan22_video_vae.py b/telefuser/models/wan22_video_vae.py index 9cfbf445..35e722e2 100644 --- a/telefuser/models/wan22_video_vae.py +++ b/telefuser/models/wan22_video_vae.py @@ -1458,7 +1458,7 @@ def cached_decode_withflag( x[:, :, i : i + 1, :, :], feat_cache=self._feat_cache, feat_idx=self._feat_idx, - first_chunk=True, + first_chunk=is_first_clip and i == 0, ) else: out_ = self.model.decoder( diff --git a/telefuser/pipelines/abot_world/__init__.py b/telefuser/pipelines/abot_world/__init__.py new file mode 100644 index 00000000..f449745c --- /dev/null +++ b/telefuser/pipelines/abot_world/__init__.py @@ -0,0 +1,5 @@ +"""Single-card ABot-World pipeline.""" + +from .pipeline import ABotWorldPipeline, ABotWorldPipelineConfig + +__all__ = ["ABotWorldPipeline", "ABotWorldPipelineConfig"] diff --git a/telefuser/pipelines/abot_world/denoising.py b/telefuser/pipelines/abot_world/denoising.py new file mode 100644 index 00000000..514dfcb3 --- /dev/null +++ b/telefuser/pipelines/abot_world/denoising.py @@ -0,0 +1,190 @@ +"""Causal long-forcing denoising stage for ABot-World.""" + +from __future__ import annotations + +from typing import Any + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.abot_world_dit import ABotWorldDiT +from telefuser.schedulers.flow_match import FlowMatchScheduler + + +class ABotWorldDenoisingStage(BaseStage): + """Run ABot's published four-step, x0-prediction causal sampler on one GPU.""" + + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + dit = module_manager.fetch_module("abot_world_dit") + if dit is None or not isinstance(dit, ABotWorldDiT): + raise ValueError("ABot-World requires a loaded abot_world_dit module") + self.dit = dit + self.model_names = ["dit"] + + def parallel_models(self) -> None: + if self.model_runtime_config.parallel_config.world_size != 1: + raise ValueError("ABot-World initial integration supports exactly one DiT GPU") + self.dit.set_attention_config(self.model_runtime_config.attention_config) + + def _new_cache(self, batch_size: int, height: int, width: int) -> tuple[list[dict[str, Any]], list[dict[str, Any]]]: + frame_tokens = (height // self.dit.patch_size[1]) * (width // self.dit.patch_size[2]) + # The fixed window includes both the retained sink and rolling tail, + # matching LingBot-World v2: 6 + 12 = 18 latent frames by default. + kv_size = self.dit.local_attn_size * frame_tokens + dtype = self.torch_dtype + self_cache: list[dict[str, Any]] = [] + cross_cache: list[dict[str, Any]] = [] + head_dim = self.dit.dim // self.dit.num_heads + for _ in range(self.dit.num_layers): + self_cache.append( + { + "k": torch.zeros( + (batch_size, kv_size, self.dit.num_heads, head_dim), dtype=dtype, device=self.device + ), + "v": torch.zeros( + (batch_size, kv_size, self.dit.num_heads, head_dim), dtype=dtype, device=self.device + ), + "global_end_index": torch.zeros(1, dtype=torch.long, device=self.device), + "local_end_index": torch.zeros(1, dtype=torch.long, device=self.device), + } + ) + cross_cache.append( + { + "k": torch.zeros( + (batch_size, self.dit.text_len, self.dit.num_heads, head_dim), dtype=dtype, device=self.device + ), + "v": torch.zeros( + (batch_size, self.dit.text_len, self.dit.num_heads, head_dim), dtype=dtype, device=self.device + ), + "is_init": False, + "sequence_length": 0, + } + ) + return self_cache, cross_cache + + @staticmethod + def _x0_prediction( + flow_prediction: torch.Tensor, + latent: torch.Tensor, + timestep: torch.Tensor, + scheduler: FlowMatchScheduler, + ) -> torch.Tensor: + if timestep.shape != (latent.shape[0], latent.shape[2]): + raise ValueError("ABot timestep must have one value per latent frame") + flat_timestep = timestep.flatten().float() + sigma_index = ( + (scheduler.timesteps.to(flat_timestep.device).unsqueeze(0) - flat_timestep.unsqueeze(1)).abs().argmin(dim=1) + ) + sigma = scheduler.sigmas.to(device=latent.device, dtype=torch.float64)[sigma_index] + sigma = sigma.view(latent.shape[0], 1, latent.shape[2], 1, 1) + return (latent.double() - sigma * flow_prediction.double()).to(flow_prediction.dtype) + + @staticmethod + def _scheduler() -> FlowMatchScheduler: + scheduler = FlowMatchScheduler(template="Wan") + # ABot's official wrapper creates the full Wan training schedule, then + # uses these exact four training timesteps. + scheduler.set_timesteps(1000, training=True, shift=5.0) + return scheduler + + @staticmethod + def _official_denoising_timesteps(scheduler: FlowMatchScheduler) -> torch.Tensor: + """Return the four warped training times from ABot's published config.""" + # ``warp_denoising_step: true`` indexes the 1,000-step shifted Wan + # schedule with ``1000 - [1000, 750, 500, 250]``. + return scheduler.timesteps[torch.tensor((0, 250, 500, 750), dtype=torch.long)] + + def _denoise_block( + self, + latent: torch.Tensor, + prompt_emb: torch.Tensor, + action_context: torch.Tensor, + first_frame_latent: torch.Tensor | None, + self_cache: list[dict[str, Any]], + cross_cache: list[dict[str, Any]], + current_start: int, + generator: torch.Generator, + scheduler: FlowMatchScheduler, + ) -> torch.Tensor: + current = latent + batch, _, frames, height, width = current.shape + replace_first = current_start == 0 and first_frame_latent is not None + if replace_first: + current = current.clone() + current[:, :, :1].copy_(first_frame_latent) + frame_tokens = (height // self.dit.patch_size[1]) * (width // self.dit.patch_size[2]) + timesteps = self._official_denoising_timesteps(scheduler).to(device=self.device) + for index, current_timestep in enumerate(timesteps): + timestep = torch.full((batch, frames), current_timestep, dtype=timesteps.dtype, device=self.device) + if replace_first: + timestep[:, 0] = 0 + with torch.autocast(self.device.type, dtype=self.torch_dtype, enabled=self.device.type == "cuda"): + flow_prediction = self.dit( + x=current.to(dtype=self.torch_dtype), + timestep=timestep, + context=prompt_emb, + act_context=action_context, + kv_cache=self_cache, + crossattn_cache=cross_cache, + current_start=current_start * frame_tokens, + ) + x0 = self._x0_prediction(flow_prediction, current, timestep, scheduler) + if index < len(timesteps) - 1: + noise = torch.randn(x0.shape, generator=generator, dtype=x0.dtype, device=self.device) + current = scheduler.add_noise(x0, noise, timesteps[index + 1]) + else: + current = x0 + if replace_first: + current[:, :, :1].copy_(first_frame_latent) + context_timestep = torch.zeros_like(timestep) + self.dit( + x=current.to(dtype=self.torch_dtype), + timestep=context_timestep, + context=prompt_emb, + act_context=action_context, + kv_cache=self_cache, + crossattn_cache=cross_cache, + current_start=current_start * frame_tokens, + ) + return current + + @with_model_offload(["dit"]) + @torch.inference_mode() + def process( + self, + noise: torch.Tensor, + prompt_emb: torch.Tensor, + action_context: torch.Tensor, + first_frame_latent: torch.Tensor, + seed: int, + ) -> torch.Tensor: + """Generate a ``1 mod 3`` latent-frame video from a starting image.""" + if noise.ndim != 5 or noise.shape[2] < 1 or (noise.shape[2] - 1) % 3: + raise ValueError("ABot latent frame count must be positive and equal to 1 mod 3") + if action_context.shape[:3] != (noise.shape[0], 32, noise.shape[2]): + raise ValueError("ABot action context must be [batch, 32, latent_frames, height, width]") + if first_frame_latent.shape != noise[:, :, :1].shape: + raise ValueError("ABot starting-image latent must be [batch, 48, 1, latent_height, latent_width]") + self.dit.set_causal_attention_window(self.dit.local_attn_size, self.dit.sink_size) + self_cache, cross_cache = self._new_cache(noise.shape[0], noise.shape[-2], noise.shape[-1]) + scheduler = self._scheduler() + generator = torch.Generator(device=self.device).manual_seed(seed) + output = [] + for start in range(0, noise.shape[2], 3): + frames = 1 if start == 0 else 3 + block = self._denoise_block( + noise[:, :, start : start + frames].to(device=self.device, dtype=self.torch_dtype), + prompt_emb.to(device=self.device, dtype=self.torch_dtype), + action_context[:, :, start : start + frames].to(device=self.device, dtype=self.torch_dtype), + first_frame_latent.to(device=self.device, dtype=self.torch_dtype) if start == 0 else None, + self_cache, + cross_cache, + start, + generator, + scheduler, + ) + output.append(block) + return torch.cat(output, dim=2) diff --git a/telefuser/pipelines/abot_world/interactive.py b/telefuser/pipelines/abot_world/interactive.py new file mode 100644 index 00000000..e9acbca3 --- /dev/null +++ b/telefuser/pipelines/abot_world/interactive.py @@ -0,0 +1,163 @@ +"""Persistent single-session interaction for ABot-World on one GPU. + +The runtime mirrors LingBot's important session invariant: text embeddings, +causal DiT KV caches, scheduler state, RNG state, and VAE temporal decode +cache all remain resident between control blocks. It intentionally supports +one local session; LiveKit admission and multi-session scheduling are a later +transport/service layer rather than a prerequisite for browser testing. +""" + +from __future__ import annotations + +import threading +from dataclasses import dataclass +from typing import Any, Mapping + +import torch +from PIL import Image + +from telefuser.core.config import WeightOffloadType + +from .pipeline import ABotWorldPipeline + + +@dataclass +class ABotWorldInteractiveSession: + """State retained across causally generated ABot action blocks.""" + + prompt_emb: torch.Tensor + first_frame_latent: torch.Tensor + self_cache: list[dict[str, Any]] + cross_cache: list[dict[str, Any]] + scheduler: Any + generator: torch.Generator + next_latent_frame: int = 0 + emitted_frames: int = 0 + closed: bool = False + + +class ABotWorldInteractivePipeline(ABotWorldPipeline): + """ABot pipeline whose model weights and one generation session stay on GPU.""" + + def __init__(self, *args: object, **kwargs: object) -> None: + super().__init__(*args, **kwargs) + self._interactive_lock = threading.RLock() + self._interactive_session: ABotWorldInteractiveSession | None = None + self._models_preloaded = False + + def preload_models(self) -> None: + """Place VAE, T5, and DiT on the configured GPU before accepting controls.""" + with self._interactive_lock: + if self._models_preloaded: + return + for stage in self._get_stages(): + stage.model_runtime_config.offload_config.offload_type = WeightOffloadType.NO_CPU_OFFLOAD + stage.onload_models() + stage.onload_models_flag = True + self._models_preloaded = True + + @torch.inference_mode() + def create_interactive_session( + self, + image: Image.Image, + prompt: str, + *, + seed: int = 42, + ) -> ABotWorldInteractiveSession: + """Encode the start image and allocate session-owned causal caches.""" + if not isinstance(image, Image.Image): + raise TypeError("image must be a PIL Image") + with self._interactive_lock: + if self._interactive_session is not None: + self.close_interactive_session(self._interactive_session) + self.preload_models() + pixels = self.preprocess_image(image.convert("RGB"), self.config.height, self.config.width) + start_latent, _ = self.vae_stage.process("encode_image", pixels, None, 1, concat_mask=False) + first_frame_latent = start_latent.unsqueeze(0).to(device=self.device, dtype=self.torch_dtype) + prompt_emb = self.text_encoding_stage.process([prompt])[0].to(device=self.device, dtype=self.torch_dtype) + self_cache, cross_cache = self.denoise_stage._new_cache( + first_frame_latent.shape[0], first_frame_latent.shape[-2], first_frame_latent.shape[-1] + ) + session = ABotWorldInteractiveSession( + prompt_emb=prompt_emb, + first_frame_latent=first_frame_latent, + self_cache=self_cache, + cross_cache=cross_cache, + scheduler=self.denoise_stage._scheduler(), + generator=torch.Generator(device=self.device).manual_seed(seed), + ) + self._interactive_session = session + return session + + @torch.inference_mode() + def generate_next_block( + self, + session: ABotWorldInteractiveSession, + actions: Mapping[str, bool] | None = None, + control_latent_frames: int = 3, + ) -> list[Image.Image]: + """Generate one causal action-controlled latent group.""" + if control_latent_frames not in {1, 3}: + raise ValueError("control_latent_frames must be 1 or 3") + with self._interactive_lock: + if session is not self._interactive_session or session.closed: + raise RuntimeError("ABot interactive session is no longer active") + frame_count = control_latent_frames + latent_shape = session.first_frame_latent.shape + noise = torch.randn( + (latent_shape[0], latent_shape[1], frame_count, latent_shape[3], latent_shape[4]), + generator=session.generator, + device=self.device, + dtype=torch.float32, + ) + action_context = self.build_action_context( + actions, + latent_frames=frame_count, + height=self.config.height, + width=self.config.width, + device=self.device, + dtype=self.torch_dtype, + ) + latents = self.denoise_stage._denoise_block( + noise.to(dtype=self.torch_dtype), + session.prompt_emb, + action_context, + session.first_frame_latent if session.next_latent_frame == 0 else None, + session.self_cache, + session.cross_cache, + session.next_latent_frame, + session.generator, + session.scheduler, + ) + decoded = self.vae_stage.process( + "decode_video_cached", + latents[0], + session.next_latent_frame == 0, + False, + ) + if decoded.ndim == 5: + decoded = decoded[0] + frames = self.tensor2video(decoded) + session.next_latent_frame += frame_count + session.emitted_frames += len(frames) + return frames + + def close_interactive_session(self, session: ABotWorldInteractiveSession | None = None) -> None: + """Release retained cache references and reset the model-specific VAE stream cache.""" + with self._interactive_lock: + target = self._interactive_session if session is None else session + if target is None or target.closed: + return + target.closed = True + target.self_cache.clear() + target.cross_cache.clear() + vae = self.vae_stage.vae + if hasattr(vae, "_feat_cache"): + vae._feat_cache = [] + vae._feat_idx = [0] + if target is self._interactive_session: + self._interactive_session = None + + def close(self) -> None: + self.close_interactive_session() + super().close() diff --git a/telefuser/pipelines/abot_world/pipeline.py b/telefuser/pipelines/abot_world/pipeline.py new file mode 100644 index 00000000..86f2d79e --- /dev/null +++ b/telefuser/pipelines/abot_world/pipeline.py @@ -0,0 +1,143 @@ +"""Native TeleFuser single-GPU pipeline for ABot-World.""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Mapping + +import torch +from PIL import Image + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.wan_video.text_encoding import TextEncodingStage +from telefuser.pipelines.wan_video.vae import VAEStage + +from .denoising import ABotWorldDenoisingStage + + +@dataclass +class ABotWorldPipelineConfig: + """Runtime settings for the public ABot-World 0.5B long-forcing checkpoint.""" + + vae_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + dit_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + height: int = 480 + width: int = 832 + latent_frames: int = 31 + # Match LingBot-World v2: six fixed sink latents plus a twelve-latent rolling tail. + local_attn_size: int = 18 + sink_size: int = 6 + + +class ABotWorldPipeline(BasePipeline): + """Image-conditioned, action-controlled ABot-World inference on one GPU. + + ``latent_frames`` is the causal Wan latent-frame count (the shipped model + configuration uses 31). It must be ``1 mod 3`` because the first image + context is generated as one block and later blocks contain three latents. + """ + + clear_memory_after_call = False + _ACTION_ORDER = ("W", "A", "S", "D", "I", "J", "K", "L") + + def __init__(self, device: str | torch.device = "cuda", torch_dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__(device=device, torch_dtype=torch_dtype) + # VAE spatial compression (16) plus DiT spatial patching (2). + self.height_division_factor = 32 + self.width_division_factor = 32 + + def _get_stages(self) -> list: + return [self.vae_stage, self.text_encoding_stage, self.denoise_stage] + + def init(self, module_manager: ModuleManager, config: ABotWorldPipelineConfig) -> None: + if config.dit_config.parallel_config.world_size != 1: + raise ValueError("ABot-World initial integration supports exactly one GPU") + if config.local_attn_size < 1: + raise ValueError("local_attn_size must be positive") + if not 0 <= config.sink_size < config.local_attn_size: + raise ValueError("sink_size must be non-negative and smaller than local_attn_size") + if config.latent_frames < 1 or (config.latent_frames - 1) % 3: + raise ValueError("latent_frames must be positive and equal to 1 mod 3") + height, width = self.check_resize_height_width(config.height, config.width) + if (height, width) != (config.height, config.width): + raise ValueError("ABot height and width must already be divisible by 32") + self._model_info = module_manager.get_model_info() + self.config = config + self.vae_stage = VAEStage("abot_world_vae", module_manager, config.vae_config) + self.text_encoding_stage = TextEncodingStage( + "abot_world_text_encoding", module_manager, config.text_encoding_config + ) + self.denoise_stage = ABotWorldDenoisingStage("abot_world_denoise", module_manager, config.dit_config) + self.denoise_stage.parallel_models() + self.denoise_stage.dit.set_causal_attention_window(config.local_attn_size, config.sink_size) + + @classmethod + def build_action_context( + cls, + keys: Mapping[str, bool] | None, + *, + latent_frames: int, + height: int, + width: int, + device: torch.device | str, + dtype: torch.dtype, + ) -> torch.Tensor: + """Create the official 32-channel action map from WASD/IJKL key state.""" + if latent_frames < 1: + raise ValueError("latent_frames must be positive") + key_state = {} if keys is None else keys + unknown = set(key_state).difference(cls._ACTION_ORDER) + if unknown: + raise ValueError(f"Unknown ABot action keys: {sorted(unknown)}") + values = [float(bool(key_state.get(key, False))) for key in cls._ACTION_ORDER] + base = torch.tensor(values, device=device, dtype=dtype).view(1, 8, 1, 1, 1) + return base.expand(1, 8, latent_frames, height, width).repeat_interleave(4, dim=1).contiguous() + + @torch.inference_mode() + def __call__( + self, + image: Image.Image, + prompt: str, + actions: Mapping[str, bool] | None = None, + seed: int = 42, + ) -> list[Image.Image]: + if not isinstance(image, Image.Image): + raise TypeError("image must be a PIL Image") + image = image.convert("RGB") + pixels = self.preprocess_image(image, self.config.height, self.config.width) + start_latent, _ = self.vae_stage.process( + "encode_image", + pixels, + None, + 1, + concat_mask=False, + ) + first_frame_latent = start_latent.unsqueeze(0).to(device=self.device, dtype=self.torch_dtype) + latent_height, latent_width = first_frame_latent.shape[-2:] + noise = self.generate_noise( + (1, first_frame_latent.shape[1], self.config.latent_frames, latent_height, latent_width), + seed=seed, + device=self.device, + dtype=torch.float32, + ) + action_context = self.build_action_context( + actions, + latent_frames=self.config.latent_frames, + height=self.config.height, + width=self.config.width, + device=self.device, + dtype=self.torch_dtype, + ) + prompt_emb = self.text_encoding_stage.process([prompt])[0] + latents = self.denoise_stage.process(noise, prompt_emb, action_context, first_frame_latent, seed) + frames = self.vae_stage.process("decode_video", latents) + return self.tensor2video(frames[0]) + + def close(self) -> None: + """Release stage references explicitly for scripts that run once.""" + for name in ("vae_stage", "text_encoding_stage", "denoise_stage"): + if hasattr(self, name): + getattr(self, name).offload_models() diff --git a/tests/integration/test_abot_world_smoke.py b/tests/integration/test_abot_world_smoke.py new file mode 100644 index 00000000..1204a6e9 --- /dev/null +++ b/tests/integration/test_abot_world_smoke.py @@ -0,0 +1,56 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +from PIL import Image + +from examples.abot_world._loader import get_pipeline +from telefuser.pipelines.abot_world.interactive import ABotWorldInteractivePipeline + + +@pytest.mark.gpu +@pytest.mark.slow +def test_abot_world_thirty_control_blocks_preserve_session_frame_count() -> None: + """Run the requested long-forcing smoke against a local checkpoint. + + Set ``ABOT_WORLD_MODEL_ROOT`` and ``ABOT_WORLD_TEST_IMAGE`` to run this + test. It is skipped on ordinary CPU CI because loading the release + checkpoint is intentionally expensive. + """ + model_root = os.environ.get("ABOT_WORLD_MODEL_ROOT") + image_path = os.environ.get("ABOT_WORLD_TEST_IMAGE") + if not model_root or not image_path: + pytest.skip("set ABOT_WORLD_MODEL_ROOT and ABOT_WORLD_TEST_IMAGE to run the ABot GPU smoke test") + if not Path(model_root).is_dir() or not Path(image_path).is_file(): + pytest.skip("ABot checkpoint or test image is unavailable") + + pipeline = get_pipeline( + model_root, + height=480, + width=832, + latent_frames=31, + pipeline_class=ABotWorldInteractivePipeline, + ) + try: + pipeline.preload_models() + with Image.open(image_path) as source: + session = pipeline.create_interactive_session( + source.convert("RGB"), + "A smooth first-person exploration through a vivid natural landscape.", + seed=42, + ) + total_frames = 0 + try: + for _ in range(30): + frames = pipeline.generate_next_block(session, {"W": True}, control_latent_frames=3) + assert frames + total_frames += len(frames) + assert total_frames > 0 + assert session.emitted_frames == total_frames + assert not session.closed + finally: + pipeline.close_interactive_session(session) + finally: + pipeline.close() diff --git a/tests/unit/models/test_wan22_video_vae.py b/tests/unit/models/test_wan22_video_vae.py new file mode 100644 index 00000000..3e2b1d7b --- /dev/null +++ b/tests/unit/models/test_wan22_video_vae.py @@ -0,0 +1,50 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +import torch.nn as nn + +from telefuser.models import wan22_video_vae + + +class _RecordingDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.first_chunk_flags: list[bool] = [] + + def forward(self, x, feat_cache, feat_idx, first_chunk: bool = False): + self.first_chunk_flags.append(first_chunk) + return x + + +def test_cached_decode_marks_only_the_first_frame_of_the_first_clip(monkeypatch) -> None: + decoder = _RecordingDecoder() + fake_vae = SimpleNamespace( + model=SimpleNamespace(conv2=lambda value: value, decoder=decoder), + z_dim=1, + _feat_cache=[], + _feat_idx=[0], + _get_scale_on_device=lambda _device, _dtype: [torch.zeros(1), torch.ones(1)], + ) + monkeypatch.setattr(wan22_video_vae, "_count_conv3d", lambda _decoder: 1) + monkeypatch.setattr(wan22_video_vae, "unpatchify", lambda video, patch_size: video) + + first = torch.ones(1, 1, 2, 1, 1) + second = torch.ones(1, 1, 1, 1, 1) + wan22_video_vae.Wan22VideoVAE.cached_decode_withflag( + fake_vae, + first, + device=torch.device("cpu"), + is_first_clip=True, + is_last_clip=False, + ) + wan22_video_vae.Wan22VideoVAE.cached_decode_withflag( + fake_vae, + second, + device=torch.device("cpu"), + is_first_clip=False, + is_last_clip=False, + ) + + assert decoder.first_chunk_flags == [True, False, False] diff --git a/tests/unit/pipelines/abot_world/__init__.py b/tests/unit/pipelines/abot_world/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/unit/pipelines/abot_world/test_interactive.py b/tests/unit/pipelines/abot_world/test_interactive.py new file mode 100644 index 00000000..3588d50a --- /dev/null +++ b/tests/unit/pipelines/abot_world/test_interactive.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from telefuser.pipelines.abot_world.interactive import ( + ABotWorldInteractivePipeline, + ABotWorldInteractiveSession, +) + + +def test_close_interactive_session_clears_all_retained_state() -> None: + pipeline = ABotWorldInteractivePipeline(device="cpu") + vae = SimpleNamespace(_feat_cache=[torch.ones(1)], _feat_idx=[3]) + pipeline.vae_stage = SimpleNamespace(vae=vae) + session = ABotWorldInteractiveSession( + prompt_emb=torch.ones(1), + first_frame_latent=torch.ones(1), + self_cache=[{"k": torch.ones(1)}], + cross_cache=[{"k": torch.ones(1)}], + scheduler=object(), + generator=torch.Generator(device="cpu"), + ) + pipeline._interactive_session = session + + pipeline.close_interactive_session(session) + + assert session.closed + assert session.self_cache == [] + assert session.cross_cache == [] + assert vae._feat_cache == [] + assert vae._feat_idx == [0] + assert pipeline._interactive_session is None diff --git a/tests/unit/pipelines/abot_world/test_interactive_web.py b/tests/unit/pipelines/abot_world/test_interactive_web.py new file mode 100644 index 00000000..13efb8cf --- /dev/null +++ b/tests/unit/pipelines/abot_world/test_interactive_web.py @@ -0,0 +1,77 @@ +from __future__ import annotations + +import threading +import time +from pathlib import Path + +from PIL import Image + +from examples.abot_world.abot_world_interactive_web import InteractiveRuntime + + +class _FakePipeline: + def __init__(self) -> None: + self.generate_calls = 0 + self.closed_sessions: list[object] = [] + + def create_interactive_session(self, image: Image.Image, prompt: str, *, seed: int) -> object: + assert image.mode == "RGB" + assert prompt + return object() + + def close_interactive_session(self, session: object) -> None: + self.closed_sessions.append(session) + + def generate_next_block(self, session: object, controls: dict[str, bool], control_latent_frames: int) -> list: + self.generate_calls += 1 + assert controls + assert control_latent_frames == 3 + return [object()] + + +def _write_image(path: Path) -> None: + Image.new("RGB", (8, 8), color=(32, 64, 96)).save(path) + + +def test_connect_with_empty_controls_does_not_advance_dit(tmp_path: Path) -> None: + image_path = tmp_path / "initial.png" + _write_image(image_path) + pipeline = _FakePipeline() + runtime = InteractiveRuntime(pipeline, fps=12, control_latent_frames=3, output_queue_size=2) + + result = runtime.start(str(image_path), "test prompt", seed=42, raw_controls=[]) + try: + assert result["new_frames"] == 0 + assert result["status"] == "Causal session ready; waiting for control input." + assert pipeline.generate_calls == 0 + finally: + runtime.stop() + + +def test_full_fifo_applies_backpressure_without_reordering() -> None: + pipeline = _FakePipeline() + runtime = InteractiveRuntime(pipeline, fps=12, control_latent_frames=3, output_queue_size=1) + first = {"type": "chunk", "index": 1} + second = {"type": "chunk", "index": 2} + runtime._output_queue.put(first) + + producer_done = threading.Event() + + def produce() -> None: + assert runtime._enqueue_video_output(second) + producer_done.set() + + producer = threading.Thread(target=produce) + producer.start() + deadline = time.monotonic() + 1.0 + while runtime._producer_backpressure_events == 0 and time.monotonic() < deadline: + time.sleep(0.01) + assert runtime._producer_backpressure_events > 0 + assert not producer_done.is_set() + + assert runtime._output_queue.get(timeout=1.0) is first + producer.join(timeout=1.0) + assert producer_done.is_set() + assert runtime._output_queue.get(timeout=1.0) is second + assert runtime._dropped_video_blocks == 0 + assert runtime._dropped_video_frames == 0 diff --git a/tests/unit/pipelines/abot_world/test_model.py b/tests/unit/pipelines/abot_world/test_model.py new file mode 100644 index 00000000..5655867f --- /dev/null +++ b/tests/unit/pipelines/abot_world/test_model.py @@ -0,0 +1,113 @@ +from __future__ import annotations + +import pytest +import torch + +from telefuser.models.abot_world_dit import ( + ABotWorldDiT, + CausalWanSelfAttention, + _rope_apply, +) +from telefuser.models.wan_video_dit import precompute_freqs_cis_3d + + +def _tiny_dit() -> ABotWorldDiT: + return ABotWorldDiT( + patch_size=(1, 2, 2), + text_len=4, + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=2, + downscale_factor_control_adapter=2, + ) + + +def test_official_state_dict_converter_is_native() -> None: + state_dict = {"blocks.0.self_attn.q.weight": torch.zeros(4, 4)} + + converted, metadata = ABotWorldDiT.state_dict_converter().from_official(state_dict) + + assert converted is state_dict + assert metadata == {} + + +def test_causal_window_updates_every_transformer_block() -> None: + model = _tiny_dit() + + model.set_causal_attention_window(local_attn_size=18, sink_size=6) + + assert model.local_attn_size == 18 + assert model.sink_size == 6 + assert all(block.self_attn.local_attn_size == 18 for block in model.blocks) + assert all(block.self_attn.sink_size == 6 for block in model.blocks) + + +def test_causal_window_rejects_invalid_sink_configuration() -> None: + model = _tiny_dit() + + with pytest.raises(ValueError, match="sink_size"): + model.set_causal_attention_window(local_attn_size=6, sink_size=6) + + +def test_sink_cache_retains_prefix_and_rolls_tail() -> None: + attention = CausalWanSelfAttention(dim=4, num_heads=1, local_attn_size=3, sink_size=1) + cache = { + "k": torch.zeros(1, 3, 1, 4), + "v": torch.zeros(1, 3, 1, 4), + "global_end_index": torch.zeros(1, dtype=torch.long), + "local_end_index": torch.zeros(1, dtype=torch.long), + } + + for frame in range(4): + value = torch.full((1, 1, 1, 4), float(frame)) + attention._update_cache(cache, value, value, current_start=frame, frame_tokens=1) + + # The first frame is the sink; the rolling tail contains the newest frames. + assert torch.equal(cache["k"][0, :, 0, 0], torch.tensor([0.0, 2.0, 3.0])) + assert int(cache["global_end_index"].item()) == 4 + assert int(cache["local_end_index"].item()) == 3 + + +def test_rope_applies_frame_indices_at_the_supported_boundary() -> None: + freqs = torch.cat(precompute_freqs_cis_3d(8), dim=1) + values = torch.randn(1, 4, 2, 8) + + output = _rope_apply(values, (2, 1, 2), freqs, torch.tensor([0, 1023])) + + assert output.shape == values.shape + assert torch.isfinite(output).all() + + +def test_rope_rejects_positions_outside_the_precomputed_table() -> None: + freqs = torch.cat(precompute_freqs_cis_3d(8), dim=1) + values = torch.randn(1, 2, 1, 8) + + with pytest.raises(ValueError, match="frame indices"): + _rope_apply(values, (2, 1, 1), freqs, torch.tensor([0, 1024])) + + +def test_sink_attention_uses_bounded_positions_for_long_sessions() -> None: + attention = CausalWanSelfAttention(dim=8, num_heads=1, local_attn_size=3, sink_size=1) + freqs = torch.cat(precompute_freqs_cis_3d(8), dim=1) + cache = { + "k": torch.zeros(1, 3, 1, 8), + "v": torch.zeros(1, 3, 1, 8), + "global_end_index": torch.tensor([2048]), + "local_end_index": torch.tensor([3]), + } + + output = attention( + torch.randn(1, 1, 8), + (1, 1, 1), + freqs, + cache, + current_start=2048, + ) + + assert output.shape == (1, 1, 8) + assert torch.isfinite(output).all() diff --git a/tests/unit/pipelines/abot_world/test_pipeline.py b/tests/unit/pipelines/abot_world/test_pipeline.py new file mode 100644 index 00000000..04f37786 --- /dev/null +++ b/tests/unit/pipelines/abot_world/test_pipeline.py @@ -0,0 +1,31 @@ +from __future__ import annotations + +import pytest +import torch + +from telefuser.pipelines.abot_world import ABotWorldPipeline + + +def test_action_context_uses_official_wasd_ijkl_channel_layout() -> None: + action = ABotWorldPipeline.build_action_context( + {"W": True, "D": True, "L": True}, + latent_frames=4, + height=32, + width=64, + device="cpu", + dtype=torch.float32, + ) + assert action.shape == (1, 32, 4, 32, 64) + # Every key is expanded into four contiguous channels, in W,A,S,D,I,J,K,L order. + assert torch.all(action[:, 0:4] == 1) + assert torch.all(action[:, 12:16] == 1) + assert torch.all(action[:, 28:32] == 1) + assert torch.all(action[:, 4:12] == 0) + assert torch.all(action[:, 16:28] == 0) + + +def test_action_context_rejects_unknown_keys() -> None: + with pytest.raises(ValueError, match="Unknown ABot action keys"): + ABotWorldPipeline.build_action_context( + {"SPACE": True}, latent_frames=1, height=32, width=32, device="cpu", dtype=torch.float32 + ) From d686623e22811adb779ab5e74b4e4c66f2484f45 Mon Sep 17 00:00:00 2001 From: youngmagician114514 <97871956+youngmagician114514@users.noreply.github.com> Date: Thu, 6 Aug 2026 03:49:07 +0000 Subject: [PATCH 2/2] test(abot): add layered model and checkpoint contracts --- .../integration/test_abot_world_checkpoint.py | 34 +++++++ .../pipelines/abot_world/test_denoising.py | 89 +++++++++++++++++++ tests/unit/pipelines/abot_world/test_model.py | 43 +++++++++ .../pipelines/abot_world/test_pipeline.py | 27 ++++++ 4 files changed, 193 insertions(+) create mode 100644 tests/integration/test_abot_world_checkpoint.py create mode 100644 tests/unit/pipelines/abot_world/test_denoising.py diff --git a/tests/integration/test_abot_world_checkpoint.py b/tests/integration/test_abot_world_checkpoint.py new file mode 100644 index 00000000..86f002b3 --- /dev/null +++ b/tests/integration/test_abot_world_checkpoint.py @@ -0,0 +1,34 @@ +from __future__ import annotations + +import os +from pathlib import Path + +import pytest +import torch + +try: + from safetensors import safe_open +except ImportError: + pytest.skip("safetensors is required for ABot checkpoint contract tests", allow_module_level=True) + +from telefuser.models.abot_world_dit import ABotWorldDiT + + +@pytest.mark.filesystem +def test_official_abot_checkpoint_has_the_exact_native_dit_contract() -> None: + model_root = os.environ.get("ABOT_WORLD_MODEL_ROOT") + if not model_root: + pytest.skip("set ABOT_WORLD_MODEL_ROOT to run the ABot checkpoint contract test") + checkpoint = Path(model_root) / "diffusion_pytorch_model.safetensors" + if not checkpoint.is_file(): + pytest.skip("ABot DiT checkpoint is unavailable") + + with safe_open(checkpoint, framework="pt", device="cpu") as handle: + official = {name: tuple(handle.get_slice(name).get_shape()) for name in handle.keys()} + with torch.device("meta"): + model = ABotWorldDiT() + expected = {name: tuple(value.shape) for name, value in model.state_dict().items()} + + assert set(official) == set(expected) + mismatches = {name: (official[name], expected[name]) for name in expected if official[name] != expected[name]} + assert not mismatches diff --git a/tests/unit/pipelines/abot_world/test_denoising.py b/tests/unit/pipelines/abot_world/test_denoising.py new file mode 100644 index 00000000..2224ae34 --- /dev/null +++ b/tests/unit/pipelines/abot_world/test_denoising.py @@ -0,0 +1,89 @@ +from __future__ import annotations + +from types import MethodType + +import torch + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.models.abot_world_dit import ABotWorldDiT +from telefuser.pipelines.abot_world.denoising import ABotWorldDenoisingStage + + +def _stage_with_recording_dit() -> tuple[ABotWorldDenoisingStage, list[torch.Tensor]]: + dit = ABotWorldDiT( + patch_size=(1, 2, 2), + text_len=4, + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=2, + downscale_factor_control_adapter=2, + ) + manager = ModuleManager(torch_dtype=torch.float32, device="cpu") + manager.add_module(dit, "abot_world_dit") + stage = ABotWorldDenoisingStage( + "abot-world-test", + manager, + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + ) + stage.parallel_models() + observed_timesteps: list[torch.Tensor] = [] + + def zero_flow_prediction(model: ABotWorldDiT, **kwargs: object) -> torch.Tensor: + del model + observed_timesteps.append(kwargs["timestep"].detach().clone()) + return torch.zeros_like(kwargs["x"]) + + dit.forward = MethodType(zero_flow_prediction, dit) + return stage, observed_timesteps + + +def test_official_four_step_schedule_matches_warped_wan_training_indices() -> None: + scheduler = ABotWorldDenoisingStage._scheduler() + + actual = ABotWorldDenoisingStage._official_denoising_timesteps(scheduler) + + torch.testing.assert_close(actual, torch.tensor([1000.0, 937.5, 833.3333, 625.0]), rtol=1e-4, atol=1e-4) + + +def test_x0_prediction_uses_the_scheduler_sigma_for_each_frame() -> None: + scheduler = ABotWorldDenoisingStage._scheduler() + timestep = ABotWorldDenoisingStage._official_denoising_timesteps(scheduler)[1].reshape(1, 1) + latent = torch.full((1, 1, 1, 1, 1), 4.0) + flow_prediction = torch.full_like(latent, 2.0) + + actual = ABotWorldDenoisingStage._x0_prediction(flow_prediction, latent, timestep, scheduler) + + torch.testing.assert_close(actual, torch.full_like(latent, 2.125)) + + +def test_denoising_block_runs_four_model_updates_then_issues_context_cache_update() -> None: + stage, observed_timesteps = _stage_with_recording_dit() + self_cache, cross_cache = stage._new_cache(batch_size=1, height=8, width=8) + scheduler = stage._scheduler() + generator = torch.Generator(device="cpu").manual_seed(42) + noise = torch.randn(1, 4, 3, 8, 8, generator=generator) + + output = stage._denoise_block( + latent=noise, + prompt_emb=torch.randn(1, 4, 16), + action_context=torch.randn(1, 32, 3, 16, 16), + first_frame_latent=None, + self_cache=self_cache, + cross_cache=cross_cache, + current_start=3, + generator=generator, + scheduler=scheduler, + ) + + expected = ABotWorldDenoisingStage._official_denoising_timesteps(scheduler) + assert output.shape == noise.shape + assert len(observed_timesteps) == 5 + for observed, timestep in zip(observed_timesteps[:4], expected, strict=True): + torch.testing.assert_close(observed, torch.full((1, 3), timestep)) + assert torch.equal(observed_timesteps[-1], torch.zeros(1, 3)) diff --git a/tests/unit/pipelines/abot_world/test_model.py b/tests/unit/pipelines/abot_world/test_model.py index 5655867f..c22f59ea 100644 --- a/tests/unit/pipelines/abot_world/test_model.py +++ b/tests/unit/pipelines/abot_world/test_model.py @@ -111,3 +111,46 @@ def test_sink_attention_uses_bounded_positions_for_long_sessions() -> None: assert output.shape == (1, 1, 8) assert torch.isfinite(output).all() + + +def test_small_dit_forward_preserves_latent_and_cache_contract() -> None: + """Exercise the complete DiT path with the same cache shape as the stage.""" + model = _tiny_dit().eval() + model.set_causal_attention_window(local_attn_size=3, sink_size=1) + batch, frames, latent_height, latent_width = 1, 1, 8, 8 + frame_tokens = (latent_height // 2) * (latent_width // 2) + self_cache = [] + cross_cache = [] + for _ in range(model.num_layers): + self_cache.append( + { + "k": torch.zeros(batch, 3 * frame_tokens, model.num_heads, model.dim // model.num_heads), + "v": torch.zeros(batch, 3 * frame_tokens, model.num_heads, model.dim // model.num_heads), + "global_end_index": torch.zeros(1, dtype=torch.long), + "local_end_index": torch.zeros(1, dtype=torch.long), + } + ) + cross_cache.append( + { + "k": torch.zeros(batch, model.text_len, model.num_heads, model.dim // model.num_heads), + "v": torch.zeros(batch, model.text_len, model.num_heads, model.dim // model.num_heads), + "is_init": False, + "sequence_length": 0, + } + ) + + output = model( + x=torch.randn(batch, model.in_dim, frames, latent_height, latent_width), + timestep=torch.tensor([[0.5]]), + context=torch.randn(batch, model.text_len, 16), + # The action map is pixel-space, while x is VAE-latent space. + act_context=torch.randn(batch, 32, frames, latent_height * 2, latent_width * 2), + kv_cache=self_cache, + crossattn_cache=cross_cache, + current_start=0, + ) + + assert output.shape == (batch, model.out_dim, frames, latent_height, latent_width) + assert torch.isfinite(output).all() + assert all(int(cache["global_end_index"].item()) == frame_tokens for cache in self_cache) + assert all(bool(cache["is_init"]) for cache in cross_cache) diff --git a/tests/unit/pipelines/abot_world/test_pipeline.py b/tests/unit/pipelines/abot_world/test_pipeline.py index 04f37786..144ae8e6 100644 --- a/tests/unit/pipelines/abot_world/test_pipeline.py +++ b/tests/unit/pipelines/abot_world/test_pipeline.py @@ -3,7 +3,10 @@ import pytest import torch +from telefuser.core.config import ModelRuntimeConfig, ParallelConfig +from telefuser.core.module_manager import ModuleManager from telefuser.pipelines.abot_world import ABotWorldPipeline +from telefuser.pipelines.abot_world.pipeline import ABotWorldPipelineConfig def test_action_context_uses_official_wasd_ijkl_channel_layout() -> None: @@ -29,3 +32,27 @@ def test_action_context_rejects_unknown_keys() -> None: ABotWorldPipeline.build_action_context( {"SPACE": True}, latent_frames=1, height=32, width=32, device="cpu", dtype=torch.float32 ) + + +@pytest.mark.parametrize( + ("config", "message"), + [ + ( + ABotWorldPipelineConfig( + dit_config=ModelRuntimeConfig( + device_type="cpu", + parallel_config=ParallelConfig(device_ids=[0, 1], dp_degree=2), + ) + ), + "exactly one GPU", + ), + (ABotWorldPipelineConfig(latent_frames=2), "1 mod 3"), + (ABotWorldPipelineConfig(local_attn_size=6, sink_size=6), "smaller than local_attn_size"), + (ABotWorldPipelineConfig(height=481), "divisible by 32"), + ], +) +def test_pipeline_rejects_release_incompatible_configuration(config: ABotWorldPipelineConfig, message: str) -> None: + pipeline = ABotWorldPipeline(device="cpu") + + with pytest.raises(ValueError, match=message): + pipeline.init(ModuleManager(device="cpu"), config)