From c0e64f5eff4ac8f2af5429e87e12cd56d97c8f81 Mon Sep 17 00:00:00 2001 From: youngmagician114514 <97871956+youngmagician114514@users.noreply.github.com> Date: Thu, 23 Jul 2026 08:41:22 +0000 Subject: [PATCH] feat(lingbot): add split async vae streaming --- ...m_lingbot_world_v2_7gpu_split_async_vae.py | 114 ++++ .../pipelines/lingbot_world_fast/async_vae.py | 322 ++++++++++ .../pipelines/lingbot_world_fast/denoising.py | 61 +- .../pipelines/lingbot_world_fast/pipeline.py | 573 ++++++++++++++++-- .../pipelines/lingbot_world_fast/service.py | 132 ++++ .../pipelines/lingbot_world_fast/session.py | 15 + .../lingbot_world_fast/test_pipeline_call.py | 106 +++- .../test_runtime_baseline.py | 15 +- .../test_service_action_loop.py | 47 +- .../lingbot_world_v2/test_service.py | 27 + 10 files changed, 1354 insertions(+), 58 deletions(-) create mode 100644 examples/lingbot/stream_lingbot_world_v2_7gpu_split_async_vae.py create mode 100644 telefuser/pipelines/lingbot_world_fast/async_vae.py diff --git a/examples/lingbot/stream_lingbot_world_v2_7gpu_split_async_vae.py b/examples/lingbot/stream_lingbot_world_v2_7gpu_split_async_vae.py new file mode 100644 index 0000000..c4ee48d --- /dev/null +++ b/examples/lingbot/stream_lingbot_world_v2_7gpu_split_async_vae.py @@ -0,0 +1,114 @@ +"""LingBot-World v2 7-GPU streaming service with split condition/VAE/DiT placement.""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.lingbot_world_fast.service import LingBotWorldFastService +from telefuser.pipelines.lingbot_world_v2.pipeline import LingBotWorldV2Pipeline, LingBotWorldV2PipelineConfig + +RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} + +PPL_CONFIG = dict( + parallelism=5, + total_gpus=7, + condition_device_id=0, + async_vae_device_id=1, + dit_device_ids=[2, 3, 4, 5, 6], + control_mode="cam", + resolution="480p", + target_fps=16, + max_duration_seconds=120.0, + chunk_size=4, + frame_policy="truncate", + sample_shift=10.0, + max_attention_size=None, + attn_impl=AttnImplType.TORCH_SDPA, + enable_fsdp=False, + local_attn_size=18, + sink_size=6, + timestep_indices=(0, 250, 500, 750), + vae_torch_dtype=torch.float32, + torch_dtype=torch.bfloat16, + enable_async_vae=True, + vae_queue_size=1, + enable_condition_prefetch=True, +) + + +def get_pipeline( + model_root: str | None = None, + v2_model_root: str | None = None, +) -> LingBotWorldV2Pipeline: + if model_root is None or v2_model_root is None: + model_zoo_path = Path(os.environ["TF_MODEL_ZOO_PATH"]).expanduser() + default_model_root = str(model_zoo_path / "Wan2.2-I2V-A14B") + default_v2_model_root = str(model_zoo_path / "lingbot" / "lingbot-world-v2-14b-causal-fast" / "transformers") + else: + default_model_root, default_v2_model_root = model_root, v2_model_root + + dtype = PPL_CONFIG["torch_dtype"] + vae_dtype = PPL_CONFIG["vae_torch_dtype"] + pipeline = LingBotWorldV2Pipeline(device="cuda", torch_dtype=dtype) + pipeline.init( + LingBotWorldV2PipelineConfig( + checkpoint_dir=model_root or default_model_root, + fast_checkpoint_path=v2_model_root or default_v2_model_root, + vae_config=ModelRuntimeConfig( + device_type="cuda", + device_id=PPL_CONFIG["condition_device_id"], + torch_dtype=vae_dtype, + ), + async_vae_config=ModelRuntimeConfig( + device_type="cuda", + device_id=PPL_CONFIG["async_vae_device_id"], + torch_dtype=vae_dtype, + ), + text_encoding_config=ModelRuntimeConfig( + device_type="cuda", + device_id=PPL_CONFIG["condition_device_id"], + torch_dtype=dtype, + ), + dit_torch_dtype=dtype, + control_type=PPL_CONFIG["control_mode"], + max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]], + local_attn_size=PPL_CONFIG["local_attn_size"], + sink_size=PPL_CONFIG["sink_size"], + timestep_indices=PPL_CONFIG["timestep_indices"], + attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]), + parallel_config=ParallelConfig( + device_ids=list(PPL_CONFIG["dit_device_ids"]), + sp_ulysses_degree=PPL_CONFIG["parallelism"], + enable_fsdp=PPL_CONFIG["enable_fsdp"], + ), + enable_async_vae=PPL_CONFIG["enable_async_vae"], + vae_queue_size=PPL_CONFIG["vae_queue_size"], + enable_condition_prefetch=PPL_CONFIG["enable_condition_prefetch"], + ) + ) + return pipeline + + +def get_service(gpu_num: int = PPL_CONFIG["total_gpus"]) -> LingBotWorldFastService: + if gpu_num < PPL_CONFIG["total_gpus"]: + raise ValueError( + f"Split async VAE config needs at least {PPL_CONFIG['total_gpus']} visible GPUs, got {gpu_num}" + ) + pipeline = get_pipeline() + return LingBotWorldFastService( + pipeline, + default_fps=PPL_CONFIG["target_fps"], + max_generation_seconds=PPL_CONFIG["max_duration_seconds"], + default_session_config={ + "control_mode": PPL_CONFIG["control_mode"], + "max_duration_seconds": PPL_CONFIG["max_duration_seconds"], + "chunk_size": PPL_CONFIG["chunk_size"], + "frame_policy": PPL_CONFIG["frame_policy"], + "sample_shift": PPL_CONFIG["sample_shift"], + "max_attention_size": PPL_CONFIG["max_attention_size"], + }, + ) diff --git a/telefuser/pipelines/lingbot_world_fast/async_vae.py b/telefuser/pipelines/lingbot_world_fast/async_vae.py new file mode 100644 index 0000000..54cd4f9 --- /dev/null +++ b/telefuser/pipelines/lingbot_world_fast/async_vae.py @@ -0,0 +1,322 @@ +from __future__ import annotations + +import queue +import threading +import time +from dataclasses import dataclass, field +from typing import TYPE_CHECKING + +import torch +from PIL import Image + +from telefuser.utils.logging import logger + +if TYPE_CHECKING: + from .pipeline import LingBotWorldFastPipeline + + +@dataclass +class AsyncVAEChunkHandle: + session_id: str | None + generation_id: int + chunk_id: int + is_last_clip: bool + enqueue_ns: int + denoise_profile: dict[str, object] = field(default_factory=dict) + queue_wait_ms: float = 0.0 + queue_depth_after_enqueue: int = 0 + frames: list[Image.Image] | None = field(default=None, repr=False) + exception: BaseException | None = field(default=None, repr=False) + done: threading.Event = field(default_factory=threading.Event, repr=False) + vae_start_ns: int | None = None + vae_end_ns: int | None = None + output_ready_ns: int | None = None + vae_gpu_ms: float | None = None + overlap_ms: float | None = None + canceled: bool = False + + +@dataclass +class AsyncVAEDecodeTask: + session_id: str | None + generation_id: int + chunk_id: int + latent: torch.Tensor = field(repr=False) + latent_ready_event: torch.cuda.Event | None = field(repr=False) + is_first_clip: bool + is_last_clip: bool + handle: AsyncVAEChunkHandle = field(repr=False) + + +class AsyncVAEManager: + """Single-process, single-VAE async decoder for LingBot streaming.""" + + _SENTINEL = object() + + def __init__(self, pipeline: LingBotWorldFastPipeline, queue_size: int) -> None: + self.pipeline = pipeline + self.queue_size = max(1, int(queue_size)) + self._queue: queue.Queue[AsyncVAEDecodeTask | object] = queue.Queue(maxsize=self.queue_size) + self._lock = threading.Condition() + self._cancelled_generations: set[int] = set() + self._next_chunk_by_generation: dict[int, int] = {} + self._active_generation_id: int | None = None + self._exception: BaseException | None = None + self._closed = False + self._stream: torch.cuda.Stream | None = None + self._worker = threading.Thread( + target=self._worker_loop, + daemon=True, + name="lingbot-async-vae", + ) + self._worker.start() + + def _is_cancelled(self, generation_id: int) -> bool: + with self._lock: + return generation_id in self._cancelled_generations + + def _mark_failed(self, exc: BaseException) -> None: + with self._lock: + if self._exception is None: + self._exception = exc + self._lock.notify_all() + + def raise_if_failed(self) -> None: + with self._lock: + if self._exception is not None: + raise RuntimeError(f"LingBot async VAE worker failed: {self._exception}") from self._exception + if self._closed: + raise RuntimeError("LingBot async VAE worker is closed") + + def enqueue(self, task: AsyncVAEDecodeTask) -> AsyncVAEChunkHandle: + start_ns = time.perf_counter_ns() + while True: + self.raise_if_failed() + try: + self._queue.put(task, timeout=0.1) + break + except queue.Full: + continue + task.handle.queue_wait_ms = (time.perf_counter_ns() - start_ns) / 1_000_000.0 + task.handle.queue_depth_after_enqueue = self._safe_qsize() + logger.info( + "lingbot_async_vae vae_enqueue " + f"session_id={task.session_id} generation_id={task.generation_id} chunk_id={task.chunk_id} " + f"queue_wait_ms={task.handle.queue_wait_ms:.3f} " + f"queue_depth={task.handle.queue_depth_after_enqueue}/{self.queue_size}" + ) + return task.handle + + def cancel_generation(self, generation_id: int, timeout: float = 30.0) -> None: + with self._lock: + self._cancelled_generations.add(generation_id) + self._next_chunk_by_generation.pop(generation_id, None) + self._lock.notify_all() + + kept: list[AsyncVAEDecodeTask | object] = [] + while True: + try: + item = self._queue.get_nowait() + except queue.Empty: + break + if isinstance(item, AsyncVAEDecodeTask) and item.generation_id == generation_id: + self._cancel_handle(item.handle) + else: + kept.append(item) + for item in kept: + try: + self._queue.put_nowait(item) + except queue.Full: + logger.warning("LingBot async VAE queue refilled while cancelling old generation") + break + + deadline = time.perf_counter() + timeout + with self._lock: + while self._active_generation_id == generation_id and self._exception is None: + remaining = deadline - time.perf_counter() + if remaining <= 0: + logger.warning(f"LingBot async VAE cancel timed out for generation_id={generation_id}") + break + self._lock.wait(timeout=remaining) + + def close(self) -> None: + with self._lock: + if self._closed: + return + self._closed = True + self._lock.notify_all() + try: + self._queue.put_nowait(self._SENTINEL) + except queue.Full: + try: + item = self._queue.get_nowait() + if isinstance(item, AsyncVAEDecodeTask): + self._cancel_handle(item.handle) + except queue.Empty: + pass + try: + self._queue.put_nowait(self._SENTINEL) + except queue.Full: + logger.warning("LingBot async VAE worker close could not enqueue sentinel") + self._worker.join(timeout=10.0) + if self._worker.is_alive(): + logger.warning("LingBot async VAE worker did not exit within timeout") + + def _cancel_handle(self, handle: AsyncVAEChunkHandle) -> None: + handle.canceled = True + handle.frames = [] + handle.output_ready_ns = time.perf_counter_ns() + handle.done.set() + + def _safe_qsize(self) -> int: + try: + return self._queue.qsize() + except NotImplementedError: + return -1 + + def _get_stream(self) -> torch.cuda.Stream | None: + device = torch.device(self.pipeline.async_vae_device or self.pipeline.vae_device) + if device.type != "cuda" or not torch.cuda.is_available(): + return None + if self._stream is not None: + return self._stream + with torch.cuda.device(device): + least_priority = 0 + greatest_priority = 0 + try: + least_priority, greatest_priority = torch.cuda.Stream.priority_range() + except Exception: + try: + least_priority, greatest_priority = torch.cuda.priority_range() + except Exception: + pass + self._stream = torch.cuda.Stream(device=device, priority=least_priority) + logger.info( + "lingbot_async_vae stream_created " + f"device={device} least_priority={least_priority} greatest_priority={greatest_priority} " + f"selected_priority={least_priority} selected=least_priority" + ) + return self._stream + + def _decode_task(self, task: AsyncVAEDecodeTask) -> None: + handle = task.handle + if self._is_cancelled(task.generation_id): + self._cancel_handle(handle) + return + with self._lock: + expected_chunk = self._next_chunk_by_generation.get(task.generation_id, task.chunk_id) + if task.chunk_id != expected_chunk: + raise RuntimeError( + f"Async VAE chunk order violation for generation_id={task.generation_id}: " + f"expected {expected_chunk}, got {task.chunk_id}" + ) + self._next_chunk_by_generation[task.generation_id] = task.chunk_id + 1 + + device = torch.device(self.pipeline.async_vae_device or self.pipeline.vae_device) + handle.vae_start_ns = time.perf_counter_ns() + if device.type == "cuda" and torch.cuda.is_available(): + with torch.cuda.device(device): + stream = self._get_stream() + assert stream is not None + start_event = torch.cuda.Event(enable_timing=True) + done_event = torch.cuda.Event(enable_timing=True) + with torch.cuda.stream(stream): + if task.latent_ready_event is not None: + stream.wait_event(task.latent_ready_event) + start_event.record(stream) + frames_tensor = self.pipeline.decode_video_cached_async( + self.pipeline._async_vae_runtime(task.generation_id), + task.latent, + is_first_clip=task.is_first_clip, + is_last_clip=task.is_last_clip, + ) + done_event.record(stream) + done_event.synchronize() + handle.vae_gpu_ms = start_event.elapsed_time(done_event) + else: + frames_tensor = self.pipeline.decode_video_cached_async( + self.pipeline._async_vae_runtime(task.generation_id), + task.latent, + is_first_clip=task.is_first_clip, + is_last_clip=task.is_last_clip, + ) + + if self._is_cancelled(task.generation_id): + self._cancel_handle(handle) + return + + images = self.pipeline.tensor2video(frames_tensor) + handle.vae_end_ns = time.perf_counter_ns() + handle.output_ready_ns = handle.vae_end_ns + handle.frames = images + dit_ms = float(handle.denoise_profile.get("dit_total_ms", 0.0) or 0.0) + vae_ms = (handle.vae_end_ns - handle.vae_start_ns) / 1_000_000.0 if handle.vae_start_ns else 0.0 + if dit_ms > 0 and vae_ms > 0: + overlap_start = max(handle.enqueue_ns, handle.vae_start_ns or handle.enqueue_ns) + overlap_end = min(handle.output_ready_ns, handle.enqueue_ns + int(dit_ms * 1_000_000)) + handle.overlap_ms = max(0.0, (overlap_end - overlap_start) / 1_000_000.0) + handle.done.set() + + memory = self._memory_snapshot(device) + extra_vae_model_copy = bool(getattr(self.pipeline, "has_separate_async_vae", False)) + logger.info( + "lingbot_async_vae output_ready " + f"session_id={task.session_id} generation_id={task.generation_id} chunk_id={task.chunk_id} " + f"frames={len(images)} vae_ms={vae_ms:.3f} vae_gpu_ms={handle.vae_gpu_ms} " + f"vae_device={device} " + f"queue_depth={self._safe_qsize()}/{self.queue_size} overlap_ms={handle.overlap_ms} " + f"vae_peak_allocated={memory['peak_allocated']} vae_peak_reserved={memory['peak_reserved']} " + f"vae_allocated={memory['allocated']} vae_reserved={memory['reserved']} " + f"vae_activation_peak_bytes={memory['vae_activation_peak_bytes']} " + f"extra_vae_model_copy={extra_vae_model_copy}" + ) + + @staticmethod + def _memory_snapshot(device: torch.device) -> dict[str, int]: + if device.type != "cuda" or not torch.cuda.is_available(): + return { + "allocated": 0, + "reserved": 0, + "peak_allocated": 0, + "peak_reserved": 0, + "vae_activation_peak_bytes": 0, + } + allocated = torch.cuda.memory_allocated(device) + reserved = torch.cuda.memory_reserved(device) + peak_allocated = torch.cuda.max_memory_allocated(device) + peak_reserved = torch.cuda.max_memory_reserved(device) + return { + "allocated": int(allocated), + "reserved": int(reserved), + "peak_allocated": int(peak_allocated), + "peak_reserved": int(peak_reserved), + "vae_activation_peak_bytes": int(max(0, peak_allocated - allocated)), + } + + def _worker_loop(self) -> None: + logger.info("LingBot async VAE worker started") + while True: + item = self._queue.get() + if item is self._SENTINEL: + break + if not isinstance(item, AsyncVAEDecodeTask): + continue + with self._lock: + self._active_generation_id = item.generation_id + self._lock.notify_all() + try: + self._decode_task(item) + except BaseException as exc: + item.handle.exception = exc + item.handle.done.set() + self._mark_failed(exc) + logger.exception( + "LingBot async VAE worker failed: " + f"session_id={item.session_id} generation_id={item.generation_id} chunk_id={item.chunk_id}" + ) + finally: + item.latent = torch.empty(0) + with self._lock: + self._active_generation_id = None + self._lock.notify_all() + logger.info("LingBot async VAE worker stopped") diff --git a/telefuser/pipelines/lingbot_world_fast/denoising.py b/telefuser/pipelines/lingbot_world_fast/denoising.py index 3d9d93f..5ccc9d6 100644 --- a/telefuser/pipelines/lingbot_world_fast/denoising.py +++ b/telefuser/pipelines/lingbot_world_fast/denoising.py @@ -1,5 +1,6 @@ from __future__ import annotations +import time from dataclasses import dataclass import torch @@ -217,12 +218,29 @@ def denoise_and_update_cache( control_chunk: torch.Tensor | None, current_start: int, max_attention_size: int, - ) -> torch.Tensor: + chunk_id: int | None = None, + return_profile: bool = False, + ) -> torch.Tensor | tuple[torch.Tensor, dict[str, object]]: """Denoise a chunk and commit its clean KV state inside each worker.""" try: state = self._cache_registry[cache_handle] except KeyError as exc: raise KeyError(f"Unknown cache handle {cache_handle}") from exc + + device = torch.device(self.device) + use_cuda_events = return_profile and device.type == "cuda" and torch.cuda.is_available() + profile: dict[str, object] = {"chunk_id": chunk_id} if return_profile else {} + denoise_start_event = denoise_end_event = None + kv_start_event = kv_end_event = None + if use_cuda_events: + with torch.cuda.device(device): + denoise_start_event = torch.cuda.Event(enable_timing=True) + denoise_end_event = torch.cuda.Event(enable_timing=True) + kv_start_event = torch.cuda.Event(enable_timing=True) + kv_end_event = torch.cuda.Event(enable_timing=True) + denoise_start_event.record() + + denoise_start_ns = time.perf_counter_ns() denoised = self.denoise_chunk( latent_chunk=latent_chunk, condition_chunk=condition_chunk, @@ -236,6 +254,12 @@ def denoise_and_update_cache( max_attention_size=max_attention_size, generator=state.generator, ) + denoise_end_ns = time.perf_counter_ns() + if use_cuda_events and denoise_end_event is not None and kv_start_event is not None: + denoise_end_event.record() + kv_start_event.record() + + kv_start_ns = time.perf_counter_ns() with torch.amp.autocast( self.device.type, dtype=self.torch_dtype, @@ -252,6 +276,41 @@ def denoise_and_update_cache( current_start=current_start, max_attention_size=max_attention_size, ) + kv_end_ns = time.perf_counter_ns() + + if return_profile: + profile.update( + { + "dit_denoise_ms": (denoise_end_ns - denoise_start_ns) / 1_000_000.0, + "kv_update_ms": (kv_end_ns - kv_start_ns) / 1_000_000.0, + "dit_total_ms": (kv_end_ns - denoise_start_ns) / 1_000_000.0, + } + ) + if use_cuda_events and all( + event is not None for event in (denoise_start_event, denoise_end_event, kv_start_event, kv_end_event) + ): + assert denoise_start_event is not None + assert denoise_end_event is not None + assert kv_start_event is not None + assert kv_end_event is not None + kv_end_event.record() + kv_end_event.synchronize() + profile.update( + { + "dit_denoise_gpu_ms": denoise_start_event.elapsed_time(denoise_end_event), + "kv_update_gpu_ms": kv_start_event.elapsed_time(kv_end_event), + "dit_total_gpu_ms": denoise_start_event.elapsed_time(kv_end_event), + } + ) + if chunk_id is not None: + logger.info( + "lingbot_async_vae dit_end " + f"chunk_id={chunk_id} dit_total_ms={profile.get('dit_total_ms'):.3f} " + f"kv_update_ms={profile.get('kv_update_ms'):.3f} " + f"dit_total_gpu_ms={profile.get('dit_total_gpu_ms')} " + f"kv_update_gpu_ms={profile.get('kv_update_gpu_ms')}" + ) + return denoised, profile return denoised def has_cache(self, cache_handle: int) -> bool: diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index 19ad756..1c32697 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -1,6 +1,8 @@ from __future__ import annotations import math +import threading +import time from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path @@ -20,11 +22,13 @@ from telefuser.utils.profiler import ProfilingContext4Debug from telefuser.worker.parallel_worker import ParallelWorker +from .async_vae import AsyncVAEChunkHandle, AsyncVAEDecodeTask, AsyncVAEManager from .control import LingBotWorldFastControlBuilder, LingBotWorldFastControlContext from .denoising import LingBotWorldFastDenoisingStage from .session import ( LingBotWorldFastChunkRequest, LingBotWorldFastChunkResult, + LingBotWorldFastConditionPrefetch, LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, LingBotWorldFastSessionStatus, @@ -37,6 +41,7 @@ class LingBotWorldFastPipelineConfig: checkpoint_dir: str = "" fast_checkpoint_path: str = "lingbot_world_fast" vae_config: ModelRuntimeConfig = field(default_factory=lambda: ModelRuntimeConfig(torch_dtype=torch.float32)) + async_vae_config: ModelRuntimeConfig | None = None text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) dit_torch_dtype: torch.dtype = torch.bfloat16 control_type: str = "cam" @@ -48,6 +53,9 @@ class LingBotWorldFastPipelineConfig: timestep_indices: tuple[int, ...] = (0, 179, 358, 679) parallel_config: ParallelConfig = field(default_factory=ParallelConfig) attention_config: AttentionConfig = field(default_factory=AttentionConfig) + enable_async_vae: bool = False + vae_queue_size: int = 1 + enable_condition_prefetch: bool = False class LingBotWorldFastPipeline(BasePipeline): @@ -59,6 +67,11 @@ def __init__(self, device: str, torch_dtype: torch.dtype = torch.bfloat16) -> No super().__init__(device=device, torch_dtype=torch_dtype) self.height_division_factor = 16 self.width_division_factor = 16 + self._async_vae_manager: AsyncVAEManager | None = None + self._async_vae_runtimes: dict[int, LingBotWorldFastGenerationSession] = {} + self._next_async_vae_generation_id = 1 + self.async_vae: WanVideoVAE | None = None + self.async_vae_device: torch.device | None = None def _get_stages(self) -> list: return [self.denoise_stage] if hasattr(self, "denoise_stage") else [] @@ -83,6 +96,73 @@ def _runtime_device(self, runtime_config: ModelRuntimeConfig) -> torch.device: return torch.device(f"cuda:{runtime_config.device_id}") return torch.device(runtime_config.device_type) + @property + def async_vae_enabled(self) -> bool: + return bool(getattr(getattr(self, "config", None), "enable_async_vae", False)) + + @property + def condition_prefetch_enabled(self) -> bool: + return bool(getattr(getattr(self, "config", None), "enable_condition_prefetch", False)) + + def _ensure_async_vae_manager(self) -> AsyncVAEManager: + if not self.async_vae_enabled: + raise RuntimeError("LingBot async VAE is disabled") + manager = getattr(self, "_async_vae_manager", None) + if manager is None: + queue_size = int(getattr(self.config, "vae_queue_size", 1)) + manager = AsyncVAEManager(self, queue_size=queue_size) + self._async_vae_manager = manager + return manager + + def _assign_async_vae_generation(self, session: LingBotWorldFastGenerationSession) -> int: + generation_id = session.async_vae_generation_id + if generation_id is None: + generation_id = self._next_async_vae_generation_id + self._next_async_vae_generation_id += 1 + session.async_vae_generation_id = generation_id + self._async_vae_runtimes[generation_id] = session + return generation_id + + def _async_vae_runtime(self, generation_id: int) -> LingBotWorldFastGenerationSession: + try: + return self._async_vae_runtimes[generation_id] + except KeyError as exc: + raise RuntimeError(f"Unknown LingBot async VAE generation_id={generation_id}") from exc + + def _cancel_async_vae_generation(self, session: LingBotWorldFastGenerationSession) -> None: + generation_id = session.async_vae_generation_id + if generation_id is None: + return + manager = getattr(self, "_async_vae_manager", None) + if manager is not None: + manager.cancel_generation(generation_id) + self._async_vae_runtimes.pop(generation_id, None) + session.async_vae_generation_id = None + + def _record_latent_ready_event(self, latent: torch.Tensor) -> torch.cuda.Event | None: + if not latent.is_cuda or not torch.cuda.is_available(): + return None + with torch.cuda.device(latent.device): + event = torch.cuda.Event() + event.record(torch.cuda.current_stream(latent.device)) + return event + + @staticmethod + def _build_vae( + vae_state_dict: dict[str, torch.Tensor], + vae_cfg: dict[str, object], + *, + device: torch.device, + torch_dtype: torch.dtype, + ) -> WanVideoVAE: + vae = WanVideoVAE(**vae_cfg) + vae.load_state_dict(vae_state_dict, strict=False) + return vae.to(device=device, dtype=torch_dtype).eval() + + @property + def has_separate_async_vae(self) -> bool: + return getattr(self, "async_vae", None) is not None and self.async_vae is not self.vae + def init(self, config: LingBotWorldFastPipelineConfig) -> None: if config.control_type not in {"cam", "act"}: raise ValueError(f"Unsupported LingBot control_type: {config.control_type!r}") @@ -91,6 +171,10 @@ def init(self, config: LingBotWorldFastPipelineConfig) -> None: self._model_info = [{"name": "lingbot_world_fast", "path": str(checkpoint_root)}] self.text_device = self._runtime_device(config.text_encoding_config) self.vae_device = self._runtime_device(config.vae_config) + async_vae_config = config.async_vae_config + self.async_vae_device = ( + self._runtime_device(async_vae_config) if async_vae_config is not None else self.vae_device + ) self.text_encoder = WanTextEncoder() self.text_encoder.load_state_dict(load_state_dict(str(checkpoint_root / "models_t5_umt5-xxl-enc-bf16.pth"))) @@ -103,9 +187,26 @@ def init(self, config: LingBotWorldFastPipelineConfig) -> None: vae_state_dict, vae_cfg = WanVideoVAE.state_dict_converter().from_official( load_state_dict(str(checkpoint_root / "Wan2.1_VAE.pth")) ) - self.vae = WanVideoVAE(**vae_cfg) - self.vae.load_state_dict(vae_state_dict, strict=False) - self.vae = self.vae.to(device=self.vae_device, dtype=config.vae_config.torch_dtype).eval() + self.vae = self._build_vae( + vae_state_dict, + vae_cfg, + device=self.vae_device, + torch_dtype=config.vae_config.torch_dtype, + ) + if async_vae_config is None: + self.async_vae = self.vae + else: + self.async_vae = self._build_vae( + vae_state_dict, + vae_cfg, + device=self.async_vae_device, + torch_dtype=async_vae_config.torch_dtype, + ) + logger.info( + "LingBot async VAE decoder loaded separately: " + f"encode_device={self.vae_device} decode_device={self.async_vae_device}" + ) + del vae_state_dict fast_path = checkpoint_root / config.fast_checkpoint_path dit_device = "cpu" if config.parallel_config.world_size > 1 else self.device @@ -128,6 +229,9 @@ def init(self, config: LingBotWorldFastPipelineConfig) -> None: denoise_stage = LingBotWorldFastDenoisingStage("lingbot_world_fast_denoise", self.dit, dit_runtime_config) self.denoise_stage = ParallelWorker(denoise_stage) if config.parallel_config.world_size > 1 else denoise_stage self._next_cache_handle = 0 + self._async_vae_manager = None + self._async_vae_runtimes = {} + self._next_async_vae_generation_id = 1 @staticmethod def _build_dit_config(config: LingBotWorldFastPipelineConfig) -> dict[str, object]: @@ -177,6 +281,24 @@ def decode_video_cached( decode_state=session.decoder_state, ) + @torch.inference_mode() + def decode_video_cached_async( + self, + session: LingBotWorldFastGenerationSession, + latents: torch.Tensor, + is_first_clip: bool, + is_last_clip: bool, + ) -> torch.Tensor: + async_vae = self.async_vae or self.vae + async_vae_device = self.async_vae_device or self.vae_device + return async_vae.cached_decode_withflag( + latents, + device=async_vae_device, + is_first_clip=is_first_clip, + is_last_clip=is_last_clip, + decode_state=session.decoder_state, + ) + @staticmethod def _best_output_size(w: int, h: int, expected_area: int, dw: int = 16, dh: int = 16) -> tuple[int, int]: """Match the LingBot source by independently flooring both dimensions.""" @@ -273,9 +395,14 @@ def _prepare_image_tensor(self, image: Image.Image, height: int, width: int) -> tensor = torch.nn.functional.interpolate(tensor.unsqueeze(0), size=(height, width), mode="bicubic").squeeze(0) return tensor.to(self.vae_device, dtype=self.config.vae_config.torch_dtype) - def _encode_condition_chunk(self, session: LingBotWorldFastGenerationSession) -> torch.Tensor: + def _encode_condition_chunk( + self, + session: LingBotWorldFastGenerationSession, + *, + chunk_index: int | None = None, + ) -> torch.Tensor: """Encode only the current image-conditioning chunk with a persistent VAE cache.""" - chunk_index = session.current_chunk_index + chunk_index = session.current_chunk_index if chunk_index is None else int(chunk_index) is_first_chunk = chunk_index == 0 is_last_chunk = chunk_index == session.chunk_count - 1 pixel_frames = 1 + 4 * (session.chunk_size - 1) if is_first_chunk else 4 * session.chunk_size @@ -311,6 +438,133 @@ def _encode_condition_chunk(self, session: LingBotWorldFastGenerationSession) -> session.condition_image = None return torch.cat([mask, latent], dim=0).unsqueeze(0) + def _encode_condition_chunk_profiled( + self, + session: LingBotWorldFastGenerationSession, + *, + chunk_index: int, + progress_callback: Callable[..., None] | None = None, + condition_prefetched: bool = False, + synchronize: bool = False, + ) -> tuple[torch.Tensor, float]: + self._notify_progress( + progress_callback, + "encoding_condition_chunk", + index=chunk_index, + condition_prefetched=condition_prefetched, + ) + condition_start_ns = time.perf_counter_ns() + condition_chunk = self._encode_condition_chunk(session, chunk_index=chunk_index) + if synchronize and self.vae_device.type == "cuda" and torch.cuda.is_available(): + torch.cuda.synchronize(self.vae_device) + condition_end_ns = time.perf_counter_ns() + condition_encode_ms = (condition_end_ns - condition_start_ns) / 1_000_000.0 + self._notify_progress( + progress_callback, + "condition_chunk_encoded", + index=chunk_index, + condition_encode_ms=condition_encode_ms, + condition_prefetched=condition_prefetched, + ) + return condition_chunk, condition_encode_ms + + def _wait_for_condition_prefetch(self, prefetch: LingBotWorldFastConditionPrefetch, timeout: float = 0.1) -> None: + while not prefetch.done.wait(timeout=timeout): + continue + + def _consume_or_encode_condition_chunk( + self, + session: LingBotWorldFastGenerationSession, + *, + chunk_index: int, + progress_callback: Callable[..., None] | None = None, + ) -> tuple[torch.Tensor, dict[str, object]]: + prefetch = session.condition_prefetch + if prefetch is None: + condition_chunk, condition_encode_ms = self._encode_condition_chunk_profiled( + session, + chunk_index=chunk_index, + progress_callback=progress_callback, + ) + return condition_chunk, { + "condition_encode_ms": condition_encode_ms, + "condition_prefetched": False, + "condition_prefetch_wait_ms": 0.0, + } + if prefetch.chunk_index != chunk_index: + raise RuntimeError( + f"LingBot condition prefetch order violation: expected chunk {chunk_index}, " + f"got prefetched chunk {prefetch.chunk_index}" + ) + + wait_start_ns = time.perf_counter_ns() + self._wait_for_condition_prefetch(prefetch) + wait_ms = (time.perf_counter_ns() - wait_start_ns) / 1_000_000.0 + session.condition_prefetch = None + if prefetch.exception is not None: + raise RuntimeError( + f"LingBot condition prefetch failed for chunk {chunk_index}: {prefetch.exception}" + ) from prefetch.exception + if prefetch.condition_chunk is None: + raise RuntimeError(f"LingBot condition prefetch for chunk {chunk_index} produced no tensor") + self._notify_progress( + progress_callback, + "condition_prefetch_consumed", + index=chunk_index, + condition_encode_ms=prefetch.condition_encode_ms, + condition_prefetch_wait_ms=wait_ms, + ) + return prefetch.condition_chunk, { + "condition_encode_ms": prefetch.condition_encode_ms, + "condition_prefetched": True, + "condition_prefetch_wait_ms": wait_ms, + } + + def _start_condition_prefetch( + self, + session: LingBotWorldFastGenerationSession, + *, + chunk_index: int, + progress_callback: Callable[..., None] | None = None, + ) -> None: + if not self.condition_prefetch_enabled or chunk_index >= session.chunk_count: + return + if session.condition_prefetch is not None: + raise RuntimeError( + f"LingBot condition prefetch already in flight for chunk {session.condition_prefetch.chunk_index}" + ) + + prefetch = LingBotWorldFastConditionPrefetch(chunk_index=chunk_index) + session.condition_prefetch = prefetch + + def worker() -> None: + prefetch.start_ns = time.perf_counter_ns() + self._notify_progress(progress_callback, "condition_prefetch_start", index=chunk_index) + try: + condition_chunk, condition_encode_ms = self._encode_condition_chunk_profiled( + session, + chunk_index=chunk_index, + progress_callback=progress_callback, + condition_prefetched=True, + synchronize=True, + ) + prefetch.condition_chunk = condition_chunk + prefetch.condition_encode_ms = condition_encode_ms + except BaseException as exc: + prefetch.exception = exc + finally: + prefetch.end_ns = time.perf_counter_ns() + self._notify_progress( + progress_callback, + "condition_prefetch_done", + index=chunk_index, + condition_encode_ms=prefetch.condition_encode_ms, + failed=prefetch.exception is not None, + ) + prefetch.done.set() + + threading.Thread(target=worker, daemon=True, name=f"lingbot-condition-prefetch-{chunk_index}").start() + def _release_session_cache(self, session: LingBotWorldFastGenerationSession) -> bool: cache_handle = session.cache_handle if cache_handle is None: @@ -329,6 +583,10 @@ def _release_session_cache(self, session: LingBotWorldFastGenerationSession) -> def release_session(self, session: LingBotWorldFastGenerationSession) -> None: """Idempotently release cache and decoder state owned by a session.""" with session.transaction_lock: + self._cancel_async_vae_generation(session) + if session.condition_prefetch is not None: + self._wait_for_condition_prefetch(session.condition_prefetch, timeout=0.1) + session.condition_prefetch = None cache_released = self._release_session_cache(session) session.decoder_state.feat_cache = [] session.decoder_state.feat_idx = [0] @@ -343,7 +601,11 @@ def release_session(self, session: LingBotWorldFastGenerationSession) -> None: session.status = LingBotWorldFastSessionStatus.RELEASED def close(self) -> None: - """Deterministically close the multi-process denoising worker group.""" + """Deterministically close async VAE and multi-process denoising workers.""" + async_vae_manager = getattr(self, "_async_vae_manager", None) + if async_vae_manager is not None: + async_vae_manager.close() + self._async_vae_manager = None denoise_stage = getattr(self, "denoise_stage", None) if isinstance(denoise_stage, ParallelWorker): denoise_stage.close() @@ -632,6 +894,98 @@ def _next_noise_chunk(self, session: LingBotWorldFastGenerationSession) -> torch dtype=torch.float32, ) + def _generate_chunk_latent( + self, + runtime: LingBotWorldFastGenerationSession, + control: torch.Tensor, + progress_callback: Callable[..., None] | None = None, + *, + return_profile: bool = False, + ) -> tuple[int, torch.Tensor, dict[str, object]]: + idx = runtime.current_chunk_index + chunk_start_ns = time.perf_counter_ns() + latent_chunk = self._next_noise_chunk(runtime) + condition_chunk, condition_profile = self._consume_or_encode_condition_chunk( + runtime, + chunk_index=idx, + progress_callback=progress_callback, + ) + self._start_condition_prefetch(runtime, chunk_index=idx + 1, progress_callback=progress_callback) + control_chunk = control + + self._notify_progress(progress_callback, "dit_start", index=idx) + current_start = idx * runtime.chunk_size * runtime.frame_tokens + cached_latent = runtime.world_kv_cached_latents.pop(idx, None) if runtime.world_kv_cached_latents else None + profile: dict[str, object] = { + "chunk_id": idx, + "chunk_start_ns": chunk_start_ns, + "cached_latent": cached_latent is not None, + } + profile.update(condition_profile) + if cached_latent is not None: + self._notify_progress(progress_callback, "decoding_cached_chunk", index=idx) + denoised = cached_latent.to(device=self.device, dtype=self.torch_dtype) + profile["dit_total_ms"] = 0.0 + profile["kv_update_ms"] = 0.0 + else: + with ProfilingContext4Debug("denoise_chunk"): + denoise_kwargs = dict( + cache_handle=runtime.cache_handle, + latent_chunk=latent_chunk, + condition_chunk=condition_chunk, + prompt_emb=runtime.prompt_emb, + control_chunk=control_chunk, + current_start=current_start, + max_attention_size=runtime.max_attention_size, + chunk_id=idx, + return_profile=return_profile, + ) + self._notify_progress(progress_callback, "kv_update_start", index=idx) + dit_start_ns = time.perf_counter_ns() + if isinstance(self.denoise_stage, ParallelWorker): + denoise_result = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs, sync=True) + else: + denoise_result = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs) + dit_end_ns = time.perf_counter_ns() + + if return_profile and isinstance(denoise_result, tuple): + denoised, worker_profile = denoise_result + profile.update(worker_profile) + else: + denoised = denoise_result + profile.setdefault("dit_total_ms", (dit_end_ns - dit_start_ns) / 1_000_000.0) + profile.setdefault("kv_update_ms", profile.get("kv_update_gpu_ms", 0.0)) + self._notify_progress(progress_callback, "dit_end", index=idx, profile=profile) + self._notify_progress(progress_callback, "kv_update_end", index=idx, profile=profile) + + if runtime.world_kv_binding is not None: + try: + runtime.world_kv_binding.on_chunk_finalized(runtime, idx, denoised) + except Exception as exc: + logger.warning(f"world_kv on_chunk_finalized failed at chunk {idx}: {exc}") + + profile["chunk_period_ms"] = (time.perf_counter_ns() - chunk_start_ns) / 1_000_000.0 + return idx, denoised, profile + + def _decode_latent_to_images( + self, + runtime: LingBotWorldFastGenerationSession, + idx: int, + denoised: torch.Tensor, + progress_callback: Callable[..., None] | None = None, + ) -> list[Image.Image]: + self._notify_progress(progress_callback, "decoding_chunk", index=idx, device=str(self.vae_device)) + with ProfilingContext4Debug("vae_decode"): + frames = self.decode_video_cached( + runtime, + denoised, + is_first_clip=(idx == 0), + is_last_clip=(idx == runtime.chunk_count - 1), + ) + images = self.tensor2video(frames) + self._notify_progress(progress_callback, "chunk_decoded", index=idx, frames=len(images)) + return images + @torch.inference_mode() def generate_next_chunk( self, @@ -643,53 +997,166 @@ def generate_next_chunk( return [] with ProfilingContext4Debug("generate_next_chunk"): - idx = runtime.current_chunk_index - latent_chunk = self._next_noise_chunk(runtime) - self._notify_progress(progress_callback, "encoding_condition_chunk", index=idx) - condition_chunk = self._encode_condition_chunk(runtime) - self._notify_progress(progress_callback, "condition_chunk_encoded", index=idx) - control_chunk = control - - self._notify_progress(progress_callback, "denoising_chunk", index=idx) - current_start = idx * runtime.chunk_size * runtime.frame_tokens - cached_latent = runtime.world_kv_cached_latents.pop(idx, None) if runtime.world_kv_cached_latents else None - if cached_latent is not None: - # On a world_kv fast-forward hit, KV is already seeded and the latent comes - # from the cached skeleton. Decode it directly and skip clean-KV rewrite. - self._notify_progress(progress_callback, "decoding_cached_chunk", index=idx) - denoised = cached_latent.to(device=self.device, dtype=self.torch_dtype) - else: - with ProfilingContext4Debug("denoise_chunk"): - denoise_kwargs = dict( - cache_handle=runtime.cache_handle, - latent_chunk=latent_chunk, - condition_chunk=condition_chunk, - prompt_emb=runtime.prompt_emb, - control_chunk=control_chunk, - current_start=current_start, - max_attention_size=runtime.max_attention_size, - ) - if isinstance(self.denoise_stage, ParallelWorker): - denoised = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs, sync=True) - else: - denoised = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs) - - if runtime.world_kv_binding is not None: - try: - runtime.world_kv_binding.on_chunk_finalized(runtime, idx, denoised) - except Exception as exc: - logger.warning(f"world_kv on_chunk_finalized failed at chunk {idx}: {exc}") - - self._notify_progress(progress_callback, "decoding_chunk", index=idx, device=str(self.vae_device)) - with ProfilingContext4Debug("vae_decode"): - frames = self.decode_video_cached( - runtime, - denoised, - is_first_clip=(idx == 0), - is_last_clip=(idx == runtime.chunk_count - 1), - ) - images = self.tensor2video(frames) - self._notify_progress(progress_callback, "chunk_decoded", index=idx, frames=len(images)) + idx, denoised, _profile = self._generate_chunk_latent( + runtime, + control=control, + progress_callback=progress_callback, + ) + images = self._decode_latent_to_images(runtime, idx, denoised, progress_callback) runtime.current_chunk_index += 1 runtime.emitted_frames += len(images) return images + + @torch.inference_mode() + def submit_async_vae_chunk( + self, + session: LingBotWorldFastGenerationSession, + request: LingBotWorldFastChunkRequest, + progress_callback: Callable[..., None] | None = None, + ) -> AsyncVAEChunkHandle: + """Generate one latent chunk, enqueue VAE decode, and return its completion handle.""" + manager = self._ensure_async_vae_manager() + manager.raise_if_failed() + if not session.transaction_lock.acquire(blocking=False): + raise RuntimeError("LingBot session already has a chunk in progress") + try: + self._validate_chunk_request(session, request) + resolved_control: torch.Tensor | None = None + if session.status == LingBotWorldFastSessionStatus.NEW: + try: + + def materialize_first_control() -> None: + nonlocal resolved_control + resolved_control = self._resolve_control(request.control) + + initialized = self._create_initialized_session( + session.config, + progress_callback, + before_cache=materialize_first_control, + ) + session.__dict__.update( + (key, value) for key, value in initialized.__dict__.items() if key != "transaction_lock" + ) + except Exception as exc: + session.status = LingBotWorldFastSessionStatus.POISONED + session.poisoned_reason = f"{type(exc).__name__}: {exc}" + self.release_session(session) + raise + if resolved_control is None: + resolved_control = self._resolve_control(request.control) + self._validate_control(session, resolved_control) + + session.status = LingBotWorldFastSessionStatus.RUNNING + try: + idx, denoised, profile = self._generate_chunk_latent( + session, + control=resolved_control, + progress_callback=progress_callback, + return_profile=True, + ) + generation_id = self._assign_async_vae_generation(session) + enqueue_ns = time.perf_counter_ns() + handle = AsyncVAEChunkHandle( + session_id=request.session_id, + generation_id=generation_id, + chunk_id=idx, + is_last_clip=(idx == session.chunk_count - 1), + enqueue_ns=enqueue_ns, + denoise_profile=profile, + ) + task = AsyncVAEDecodeTask( + session_id=request.session_id, + generation_id=generation_id, + chunk_id=idx, + latent=denoised, + latent_ready_event=self._record_latent_ready_event(denoised), + is_first_clip=(idx == 0), + is_last_clip=(idx == session.chunk_count - 1), + handle=handle, + ) + self._notify_progress(progress_callback, "vae_enqueue", index=idx, generation_id=generation_id) + manager.enqueue(task) + self._notify_progress( + progress_callback, + "vae_enqueued", + index=idx, + generation_id=generation_id, + vae_queue_wait_ms=handle.queue_wait_ms, + vae_queue_depth=handle.queue_depth_after_enqueue, + ) + except Exception as exc: + session.status = LingBotWorldFastSessionStatus.POISONED + session.poisoned_reason = f"{type(exc).__name__}: {exc}" + self.release_session(session) + raise + + session.current_chunk_index += 1 + session.status = LingBotWorldFastSessionStatus.COMMITTED + logger.info( + f"Submitted LingBot async VAE chunk {idx + 1}/{session.chunk_count}: " + f"queue_wait_ms={handle.queue_wait_ms:.3f}" + ) + return handle + finally: + session.transaction_lock.release() + + def wait_async_vae_chunk( + self, + session: LingBotWorldFastGenerationSession, + handle: AsyncVAEChunkHandle, + progress_callback: Callable[..., None] | None = None, + ) -> list[Image.Image]: + """Wait for one async VAE chunk and return frames if it still belongs to the session.""" + manager = getattr(self, "_async_vae_manager", None) + while not handle.done.wait(timeout=0.1): + if manager is not None: + manager.raise_if_failed() + if handle.exception is not None: + session.status = LingBotWorldFastSessionStatus.POISONED + session.poisoned_reason = f"{type(handle.exception).__name__}: {handle.exception}" + self.release_session(session) + raise RuntimeError(session.poisoned_reason) from handle.exception + if handle.canceled or handle.generation_id != session.async_vae_generation_id: + logger.info( + "Dropped stale LingBot async VAE output: " + f"generation_id={handle.generation_id} chunk_id={handle.chunk_id}" + ) + return [] + + frames = handle.frames or [] + session.emitted_frames += len(frames) + output_ready_ns = handle.output_ready_ns or time.perf_counter_ns() + chunk_start_ns = int(handle.denoise_profile.get("chunk_start_ns", handle.enqueue_ns)) + chunk_period_ms = (output_ready_ns - chunk_start_ns) / 1_000_000.0 + vae_ms = None + if handle.vae_start_ns is not None and handle.vae_end_ns is not None: + vae_ms = (handle.vae_end_ns - handle.vae_start_ns) / 1_000_000.0 + overlap_ratio = None + if handle.overlap_ms is not None: + dit_ms = float(handle.denoise_profile.get("dit_total_ms", 0.0) or 0.0) + overlap_vae_ms = vae_ms + if overlap_vae_ms is None: + overlap_vae_ms = ( + (handle.vae_end_ns or output_ready_ns) - (handle.vae_start_ns or handle.enqueue_ns) + ) / 1_000_000.0 + denom = min(dit_ms, overlap_vae_ms) + overlap_ratio = handle.overlap_ms / denom if denom > 0 else None + self._notify_progress( + progress_callback, + "output_ready", + index=handle.chunk_id, + frames=len(frames), + vae_queue_wait_ms=handle.queue_wait_ms, + vae_ms=vae_ms, + vae_gpu_ms=handle.vae_gpu_ms, + chunk_period_ms=chunk_period_ms, + overlap_ms=handle.overlap_ms, + overlap_ratio=overlap_ratio, + ) + logger.info( + "LingBot async VAE chunk ready: " + f"chunk_id={handle.chunk_id} frames={len(frames)} chunk_period_ms={chunk_period_ms:.3f} " + f"queue_wait_ms={handle.queue_wait_ms:.3f} overlap_ms={handle.overlap_ms} " + f"overlap_ratio={overlap_ratio}" + ) + return frames diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index 991abf0..9a4a09e 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -573,6 +573,10 @@ def _run_worker_loop( emit_status: Callable[..., None], ) -> None: try: + if bool(getattr(self.pipeline, "async_vae_enabled", False)): + self._run_worker_loop_async_vae(session_id, state, emit_status) + return + self._emit_preview_frame(state) control_context = state.control_context or self.pipeline.control_context(state.config) control_builder = LingBotWorldFastControlBuilder(control_context) @@ -676,6 +680,134 @@ def _run_worker_loop( self._put_output(state, {"type": "done"}) self._release_generation_session(state) + def _run_worker_loop_async_vae( + self, + session_id: str, + state: LingBotWorldFastSessionState, + emit_status: Callable[..., None], + ) -> None: + self._emit_preview_frame(state) + control_context = state.control_context or self.pipeline.control_context(state.config) + control_builder = LingBotWorldFastControlBuilder(control_context) + state.generation_session = LingBotWorldFastGenerationSession(config=state.config) + runtime = state.generation_session + emit_status( + "runtime_ready", + width=control_context.width, + height=control_context.height, + latent_frames=control_context.latent_frames, + total_chunks=control_context.latent_frames // control_context.chunk_size, + enable_async_vae=True, + enable_condition_prefetch=bool(getattr(self.pipeline, "condition_prefetch_enabled", False)), + ) + + pending_handle = None + pending_controls = None + stop_requested = False + + def runtime_finished() -> bool: + return runtime.chunk_size > 0 and runtime.current_chunk_index >= runtime.chunk_count + + def flush_pending() -> None: + nonlocal pending_handle, pending_controls + if pending_handle is None: + return + frames = self.pipeline.wait_async_vae_chunk(runtime, pending_handle, progress_callback=emit_status) + if frames: + if state.config.show_control_hud: + frames = self._overlay_control_hud(frames, pending_controls) + payload = { + "type": "chunk", + "index": pending_handle.chunk_id, + "fps": state.config.fps, + "timestamp": time.time(), + "frames_b64": self._encode_frames_to_b64(frames), + } + self._put_output(state, payload) + emit_status("chunk_sent", index=pending_handle.chunk_id, frames=len(frames)) + pending_handle = None + pending_controls = None + + while state.active: + if pending_handle is not None: + with state.control_lock: + controls_held = bool(state.pressed_controls) + if not controls_held and state.pending_inputs.empty(): + flush_pending() + if runtime_finished(): + break + + if runtime_finished(): + break + + with state.control_lock: + controls_held = bool(state.pressed_controls) + if controls_held: + try: + incoming = state.pending_inputs.get_nowait() + except queue.Empty: + incoming = {"type": "direction_control"} + else: + incoming = state.pending_inputs.get() + if incoming.get("type") == "stop": + stop_requested = True + break + + explicit_control = incoming if self._is_explicit_control_chunk(incoming) else None + applied_controls = None + while True: + try: + next_item = state.pending_inputs.get_nowait() + except queue.Empty: + break + if next_item.get("type") == "stop": + incoming = next_item + break + if self._is_explicit_control_chunk(next_item): + explicit_control = next_item + + if incoming and incoming.get("type") == "stop": + stop_requested = True + break + if explicit_control is not None: + control = control_builder.defer(explicit_control) + else: + directional_chunk = self._build_directional_control_chunk(state, control_context) + if directional_chunk is None: + continue + applied_controls = directional_chunk["controls"] + emit_status( + "applying_direction_control", + index=runtime.current_chunk_index, + controls=applied_controls, + move_step=state.config.control_move_step, + yaw_step_degrees=state.config.control_yaw_step_degrees, + lateral_step=state.config.control_lateral_step, + pitch_step_degrees=state.config.control_pitch_step_degrees, + ) + control = control_builder.defer(directional_chunk) + + submit_index = runtime.current_chunk_index + emit_status("generating_chunk", index=submit_index) + current_handle = self.pipeline.submit_async_vae_chunk( + runtime, + LingBotWorldFastChunkRequest( + chunk_index=submit_index, + session_id=session_id, + control=control, + ), + progress_callback=emit_status, + ) + if pending_handle is not None: + flush_pending() + pending_handle = current_handle + pending_controls = applied_controls + if current_handle.is_last_clip: + break + + if not stop_requested and pending_handle is not None: + flush_pending() + def push_chunk(self, session_id: str, chunk: dict) -> None: state = self._sessions.get(session_id) if state is None: diff --git a/telefuser/pipelines/lingbot_world_fast/session.py b/telefuser/pipelines/lingbot_world_fast/session.py index f9db47b..dee5f05 100644 --- a/telefuser/pipelines/lingbot_world_fast/session.py +++ b/telefuser/pipelines/lingbot_world_fast/session.py @@ -99,6 +99,19 @@ class LingBotWorldFastSessionStatus(str, Enum): RELEASED = "released" +@dataclass +class LingBotWorldFastConditionPrefetch: + """One in-flight condition chunk encoded ahead of its DiT use.""" + + chunk_index: int + done: threading.Event = field(default_factory=threading.Event, repr=False) + condition_chunk: torch.Tensor | None = field(default=None, repr=False) + condition_encode_ms: float = 0.0 + start_ns: int | None = None + end_ns: int | None = None + exception: BaseException | None = field(default=None, repr=False) + + @dataclass class LingBotWorldFastGenerationSession: """Externally owned state for one chunked LingBot generation.""" @@ -126,6 +139,8 @@ class LingBotWorldFastGenerationSession: # CacheSeek world_kv binding and decode-only latents for fast-forward hits. world_kv_binding: object | None = None world_kv_cached_latents: dict[int, torch.Tensor] = field(default_factory=dict) + async_vae_generation_id: int | None = None + condition_prefetch: LingBotWorldFastConditionPrefetch | None = field(default=None, repr=False) @property def chunk_count(self) -> int: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py index 960d03d..86c7e54 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py @@ -1,4 +1,5 @@ import threading +import time from types import SimpleNamespace from unittest.mock import MagicMock, patch @@ -46,7 +47,15 @@ def _control() -> torch.Tensor: def _pipeline() -> LingBotWorldFastPipeline: pipeline = LingBotWorldFastPipeline(device="cpu") - pipeline.config = SimpleNamespace(control_type="cam") + pipeline.config = SimpleNamespace( + control_type="cam", + enable_async_vae=False, + vae_queue_size=1, + async_vae_config=None, + enable_condition_prefetch=False, + ) + pipeline.vae_device = torch.device("cpu") + pipeline.async_vae_device = torch.device("cpu") return pipeline @@ -290,3 +299,98 @@ def test_pipeline_close_delegates_to_parallel_worker() -> None: pipeline.close() worker.close.assert_called_once_with() + + +def test_async_vae_submit_wait_returns_ordered_frames_without_sync_path_changes() -> None: + pipeline = _pipeline() + pipeline.config.enable_async_vae = True + runtime = _session(chunk_count=2) + expected = [Image.new("RGB", (8, 8), "blue")] + latent = torch.zeros(1, 16, 1, 1, 1) + pipeline._generate_chunk_latent = MagicMock( + return_value=(0, latent, {"dit_total_ms": 1.0, "chunk_start_ns": time.perf_counter_ns()}) + ) + pipeline.decode_video_cached_async = MagicMock(return_value=torch.zeros(3, 1, 8, 8)) + pipeline.tensor2video = MagicMock(return_value=expected) + + try: + handle = pipeline.submit_async_vae_chunk( + runtime, + LingBotWorldFastChunkRequest(chunk_index=0, session_id="session-a", control=_control()), + ) + frames = pipeline.wait_async_vae_chunk(runtime, handle) + finally: + pipeline.close() + + assert frames == expected + assert handle.chunk_id == 0 + assert runtime.current_chunk_index == 1 + assert runtime.emitted_frames == 1 + assert runtime.async_vae_generation_id is not None + pipeline._generate_chunk_latent.assert_called_once() + pipeline.decode_video_cached_async.assert_called_once() + pipeline.tensor2video.assert_called_once() + + +def test_async_decode_uses_separate_vae_decoder_when_configured() -> None: + pipeline = _pipeline() + session = _session() + latent = torch.zeros(1, 16, 1, 1, 1) + expected = torch.ones(3, 1, 8, 8) + pipeline.vae = MagicMock() + pipeline.async_vae = MagicMock() + pipeline.async_vae_device = torch.device("cpu") + pipeline.async_vae.cached_decode_withflag.return_value = expected + + actual = pipeline.decode_video_cached_async(session, latent, is_first_clip=True, is_last_clip=False) + + assert actual is expected + pipeline.async_vae.cached_decode_withflag.assert_called_once_with( + latent, + device=torch.device("cpu"), + is_first_clip=True, + is_last_clip=False, + decode_state=session.decoder_state, + ) + pipeline.vae.cached_decode_withflag.assert_not_called() + + +def test_condition_prefetch_consumes_next_chunk_without_reencoding() -> None: + pipeline = _pipeline() + pipeline.config.enable_condition_prefetch = True + runtime = _session(chunk_count=3) + calls: list[int] = [] + + def encode_condition(_session, *, chunk_index=None): + calls.append(int(chunk_index)) + return torch.full((1,), int(chunk_index), dtype=torch.float32) + + pipeline._encode_condition_chunk = MagicMock(side_effect=encode_condition) + + pipeline._start_condition_prefetch(runtime, chunk_index=1) + condition_chunk, profile = pipeline._consume_or_encode_condition_chunk(runtime, chunk_index=1) + + torch.testing.assert_close(condition_chunk, torch.tensor([1.0])) + assert calls == [1] + assert profile["condition_prefetched"] is True + assert profile["condition_prefetch_wait_ms"] >= 0.0 + assert runtime.condition_prefetch is None + + +def test_release_session_cancels_async_vae_generation_before_clearing_decoder_state() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + pipeline._async_vae_manager = MagicMock() + session = _session() + session.async_vae_generation_id = 44 + pipeline._async_vae_runtimes = {44: session} + session.decoder_state.feat_cache = [object()] + session.decoder_state.feat_idx = [3] + + pipeline.release_session(session) + + pipeline._async_vae_manager.cancel_generation.assert_called_once_with(44) + assert session.async_vae_generation_id is None + assert 44 not in pipeline._async_vae_runtimes + assert session.decoder_state.feat_cache == [] + assert session.decoder_state.feat_idx == [0] diff --git a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py index 0b71a41..e5064e6 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py @@ -58,8 +58,19 @@ def _create_runtime(frame_num: int, seed: int = 42): def test_v1_and_v2_defaults_match_the_shared_source_contract() -> None: image = Image.new("RGB", (16, 16)) - assert LingBotWorldFastPipelineConfig().vae_config.torch_dtype == torch.float32 - assert LingBotWorldV2PipelineConfig().vae_config.torch_dtype == torch.float32 + fast_config = LingBotWorldFastPipelineConfig() + v2_config = LingBotWorldV2PipelineConfig() + + assert fast_config.vae_config.torch_dtype == torch.float32 + assert v2_config.vae_config.torch_dtype == torch.float32 + assert fast_config.enable_async_vae is False + assert v2_config.enable_async_vae is False + assert fast_config.async_vae_config is None + assert v2_config.async_vae_config is None + assert fast_config.enable_condition_prefetch is False + assert v2_config.enable_condition_prefetch is False + assert fast_config.vae_queue_size == 1 + assert v2_config.vae_queue_size == 1 assert LingBotWorldFastSessionConfig(prompt="v1", image=image).frame_policy == "truncate" diff --git a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py index 20fef35..ae5a00a 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -2,7 +2,7 @@ import threading from pathlib import Path from types import SimpleNamespace -from unittest.mock import MagicMock, patch +from unittest.mock import ANY, MagicMock, patch import numpy as np import pytest @@ -75,6 +75,51 @@ def generate_chunk(runtime, request, progress_callback=None): pipeline.release_session.assert_called_once() +def test_async_worker_initializes_session_before_finish_check() -> None: + pipeline = MagicMock() + pipeline.async_vae_enabled = True + control_context = SimpleNamespace( + control_type="cam", + chunk_size=4, + width=8, + height=8, + latent_frames=4, + ) + pipeline.control_context.return_value = control_context + service = LingBotWorldFastService(pipeline) + state = LingBotWorldFastSessionState( + config=LingBotWorldFastSessionConfig( + prompt="test", + image=Image.new("RGB", (8, 8)), + frame_num=13, + chunk_size=4, + show_control_hud=False, + ), + control_context=control_context, + ) + state.pending_inputs.put({"type": "control", "control_tensor": torch.ones(1)}) + handle = SimpleNamespace(chunk_id=0, is_last_clip=True) + + def submit_chunk(runtime, request, progress_callback=None): + assert runtime.chunk_size == 0 + runtime.chunk_size = 4 + runtime.latent_f = 4 + runtime.current_chunk_index = 1 + return handle + + pipeline.submit_async_vae_chunk.side_effect = submit_chunk + pipeline.wait_async_vae_chunk.return_value = [Image.new("RGB", (8, 8))] + builder = MagicMock() + builder.defer.return_value = MagicMock(return_value=torch.ones(1)) + + with patch("telefuser.pipelines.lingbot_world_fast.service.LingBotWorldFastControlBuilder", return_value=builder): + service._run_worker_loop("session-a", state, MagicMock()) + + pipeline.submit_async_vae_chunk.assert_called_once() + pipeline.wait_async_vae_chunk.assert_called_once_with(ANY, handle, progress_callback=ANY) + pipeline.release_session.assert_called_once() + + def test_direction_action_updates_state_and_wakes_worker() -> None: service = LingBotWorldFastService(MagicMock()) state = _state() diff --git a/tests/unit/pipelines/lingbot_world_v2/test_service.py b/tests/unit/pipelines/lingbot_world_v2/test_service.py index 067b9b4..64d843f 100644 --- a/tests/unit/pipelines/lingbot_world_v2/test_service.py +++ b/tests/unit/pipelines/lingbot_world_v2/test_service.py @@ -6,6 +6,7 @@ from examples.lingbot import lingbot_world_v2_image_to_video_h100 as offline_example from examples.lingbot import stream_lingbot_world_v2 as stream_example +from examples.lingbot import stream_lingbot_world_v2_7gpu_split_async_vae as split_async_stream_example from telefuser.core.config import AttnImplType from telefuser.pipelines.lingbot_world_fast.session import LingBotWorldFastSessionConfig @@ -32,6 +33,32 @@ def test_v2_stream_get_pipeline_maps_ppl_config_to_internal_workers() -> None: assert config.parallel_config.device_ids == [0, 1, 2, 3] +def test_v2_split_async_7gpu_stream_config_separates_condition_vae_and_dit() -> None: + pipeline = MagicMock() + + with patch.object(split_async_stream_example, "LingBotWorldV2Pipeline", return_value=pipeline) as pipeline_cls: + result = split_async_stream_example.get_pipeline( + model_root="/models/Wan2.2-I2V-A14B", + v2_model_root="/models/lingbot-world-v2-14b-causal-fast/transformers", + ) + + assert result is pipeline + pipeline_cls.assert_called_once_with(device="cuda", torch_dtype=torch.bfloat16) + config = pipeline.init.call_args.args[0] + assert config.vae_config.device_id == 0 + assert config.async_vae_config.device_id == 1 + assert config.text_encoding_config.device_id == 0 + assert config.parallel_config.device_ids == [2, 3, 4, 5, 6] + assert config.parallel_config.sp_ulysses_degree == 5 + assert config.enable_async_vae is True + assert config.enable_condition_prefetch is True + assert config.vae_queue_size == 1 + assert config.local_attn_size == 18 + assert config.sink_size == 6 + assert config.timestep_indices == (0, 250, 500, 750) + assert config.attention_config.attn_impl == AttnImplType.TORCH_SDPA + + def test_v2_stream_service_constructs_v2_session_from_ppl_config() -> None: pipeline = MagicMock()