# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM-Omni project """Request-scoped LingBot-World v2 causal DMD pipeline.""" from __future__ import annotations import math import os from collections.abc import Callable, Iterable, Iterator, Sequence from contextlib import contextmanager from dataclasses import dataclass, replace from typing import TYPE_CHECKING, Any, ClassVar, cast import numpy as np import PIL.Image import torch import torch.nn.functional as F from diffusers.utils.torch_utils import randn_tensor from torch import nn from transformers import AutoTokenizer, UMT5EncoderModel from vllm.distributed import get_tensor_model_parallel_world_size from vllm.model_executor.models.utils import AutoWeightsLoader from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig from vllm_omni.diffusion.distributed.autoencoders.autoencoder_kl_wan import DistributedAutoencoderKLWan from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader from vllm_omni.diffusion.model_loader.hub_prefetch import from_pretrained_with_prefetch, prefetch_subfolders from vllm_omni.diffusion.models.interface import SupportImageInput, SupportsComponentDiscovery, SupportsStepExecution from vllm_omni.diffusion.models.lingbot_world.actions import ( LINGBOT_CAMERA_ACTION_SCHEMA, LINGBOT_CAMERA_TRAJECTORY_SCHEMA, LingBotCameraActionFrames, LingBotCameraActionScript, as_camera_action_frames, as_camera_action_script, integrate_lingbot_camera_actions, parse_lingbot_camera_action_frames, parse_lingbot_camera_action_script, ) from vllm_omni.diffusion.models.lingbot_world.camera import ( CameraTrajectory, build_plucker_embedding, interpolate_camera_trajectory, load_camera_trajectory, resolve_trusted_action_directory, ) from vllm_omni.diffusion.models.lingbot_world.dmd_block import ARBlockContext, LingBotDMDBlockRunner from vllm_omni.diffusion.models.lingbot_world.transformer import ( CausalLingBotWorldTransformer3DModel, LingBotAttentionCache, LingBotTransformerCache, ) from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.utils import _load_json from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import load_transformer_config, retrieve_latents from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.diffusion.worker.request_batch import DiffusionRequestBatch from vllm_omni.experimental.ar_diffusion.capability import ( ARDiffusionCrossAttentionKVSpec, ARDiffusionKVBranchSpec, ARDiffusionKVCacheSpec, ) from vllm_omni.experimental.ar_diffusion.tick_protocol import ( ARDiffusionChunkMetadata, ARDiffusionTickRequest, ) if TYPE_CHECKING: from diffusers.video_processor import VideoProcessor from tqdm.std import tqdm as TqdmProgressBar from vllm_omni.diffusion.worker.input_batch import InputBatch from vllm_omni.diffusion.worker.utils import StepRequestState from vllm_omni.experimental.ar_diffusion.kv_cache.state import ( ARDiffusionKVState, ) LINGBOT_DMD_TIMESTEPS = (1000, 750, 500, 250) _CAMERA_SPATIAL_FOLD = 8 _MAX_PIXEL_AREA = 480 * 832 _MAX_SOURCE_IMAGE_PIXELS = 4096 * 4096 _MAX_RAW_FRAMES = 117 _MAX_SEQUENCE_LENGTH = 512 _ACTION_ROOT_ENV = "VLLM_OMNI_LINGBOT_ACTION_ROOT" _PREPROCESSED_CAMERA_KEY = "_lingbot_camera_trajectory" _PREPROCESSED_CAMERA_ACTIONS_KEY = "_lingbot_camera_actions" _PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY = "_lingbot_camera_action_script" _SOURCE_IMAGE_ERROR = ( "Unable to load multi_modal_data.image; expected a decodable image within 4096 * 4096 source pixels." ) @dataclass(frozen=True) class _LingBotRequestInputs: """Validated model-specific request boundary. Pre-processing materializes exactly one of the three camera fields, chosen from the shape of the incoming control. They are not alternative spellings of one input: ``camera_action_script`` is a *container* whose element is a ``camera_actions``, and the two are live at once during step execution. - ``camera_trajectory`` is explicit per-frame poses and intrinsics. Offline replay loads it from ``extra_args["action_path"]``; the tick entry point receives the same content as a ``lingbot.camera_trajectory.v1`` control. Integrating WASD actions also produces one, written back onto this field. - ``camera_actions`` is WASD key state for *one* three-latent-frame block: the unit every path reduces to before the camera is embedded. The tick entry point fills it directly, because one of its requests is one block; step execution derives it per block by slicing ``camera_action_script``. ``_prepare_camera`` also reads it as the signal that the trajectory it was handed is already latent-aligned. - ``camera_action_script`` holds one such action list *per generated chunk* for the whole rollout, which only step execution needs, because only it runs every block inside a single ``generate()``. The tick entry point has no use for it and ``forward()`` rejects a request that carries it. """ prompt: str image: PIL.Image.Image | torch.Tensor camera_trajectory: CameraTrajectory | None camera_actions: LingBotCameraActionFrames | None camera_action_script: LingBotCameraActionScript | None height: int width: int num_frames: int num_latent_frames: int output_type: str max_sequence_length: int flow_shift: float generator: torch.Generator @dataclass class _LingBotARSessionState: """Small model-owned state; attention tensors remain runner-owned.""" next_chunk_index: int = 0 prompt: str | None = None generator_state: torch.Tensor | None = None image_condition: torch.Tensor | None = None camera_tail: CameraTrajectory | None = None camera_pitch: float = 0.0 def _positive_finite_flow_shift(value: Any) -> float: if isinstance(value, bool): raise ValueError("flow_shift must be a positive finite number.") try: flow_shift = float(value) except (TypeError, ValueError, OverflowError) as exc: raise ValueError("flow_shift must be a positive finite number.") from exc if not math.isfinite(flow_shift) or flow_shift <= 0: raise ValueError("flow_shift must be a positive finite number.") return flow_shift def _build_shifted_flow_schedule( *, flow_shift: float, ) -> tuple[tuple[float, float], ...]: """Map checkpoint DMD labels to request-local warped timestep/sigma pairs.""" # The checkpoint labels index the unshifted [1, ..., 1 / N] training # lattice. Apply the request's flow shift before using the value both as # the transformer timestep and the sampling sigma. base_sigmas = torch.tensor(LINGBOT_DMD_TIMESTEPS, dtype=torch.float64) / 1000 shifted_numerators = flow_shift * base_sigmas shifted_sigmas = shifted_numerators / ((1.0 - base_sigmas) + shifted_numerators) warped_timesteps = shifted_sigmas * 1000 return tuple( (float(timestep), float(sigma)) for timestep, sigma in zip(warped_timesteps.tolist(), shifted_sigmas.tolist(), strict=True) ) def _validate_scheduler_config(config: dict[str, Any]) -> None: contract = { "_class_name": "UniPCMultistepScheduler", "num_train_timesteps": 1000, "prediction_type": "flow_prediction", "predict_x0": True, "use_flow_sigmas": True, "use_dynamic_shifting": False, "use_beta_sigmas": False, "use_exponential_sigmas": False, "use_karras_sigmas": False, "final_sigmas_type": "zero", "timestep_spacing": "linspace", } for name, expected in contract.items(): if config.get(name) != expected: raise ValueError(f"LingBot scheduler config {name} must be {expected!r}, got {config.get(name)!r}.") def _validate_parallel_config(od_config: OmniDiffusionConfig) -> None: if getattr(od_config, "quantization_config", None) is not None: raise NotImplementedError("LingBot World v1 does not support quantization.") parallel_config = getattr(od_config, "parallel_config", None) if parallel_config is None: return unsupported_sizes = { "pipeline_parallel_size": "pipeline parallelism", "sequence_parallel_size": "sequence parallelism", "cfg_parallel_size": "CFG parallelism", "vae_patch_parallel_size": "VAE parallelism", } for field, feature in unsupported_sizes.items(): size = getattr(parallel_config, field, 1) or 1 if size > 1: raise NotImplementedError(f"LingBot World v1 does not support {feature} ({field}={size}).") if getattr(parallel_config, "use_hsdp", False): raise NotImplementedError("LingBot World v1 does not support HSDP.") if getattr(parallel_config, "enable_expert_parallel", False): raise NotImplementedError("LingBot World v1 does not support expert parallelism.") def _validate_source_image_size(image: PIL.Image.Image) -> None: width, height = image.size if ( isinstance(width, bool) or not isinstance(width, int) or width <= 0 or isinstance(height, bool) or not isinstance(height, int) or height <= 0 ): raise ValueError("source image width and height must be positive integers.") if width * height > _MAX_SOURCE_IMAGE_PIXELS: raise ValueError("source image pixel count must not exceed 4096 * 4096.") def _decode_source_image(image: PIL.Image.Image) -> PIL.Image.Image: _validate_source_image_size(image) try: return image.convert("RGB") except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError): raise ValueError(_SOURCE_IMAGE_ERROR) from None def _load_source_image(path: str | os.PathLike[str]) -> PIL.Image.Image: try: source_image = PIL.Image.open(path) except (OSError, SyntaxError, ValueError, PIL.Image.DecompressionBombError): raise ValueError(_SOURCE_IMAGE_ERROR) from None try: return _decode_source_image(source_image) finally: source_image.close() def get_lingbot_world_pre_process_func( od_config: OmniDiffusionConfig, ) -> Callable[[OmniDiffusionRequest], OmniDiffusionRequest]: """Materialize request-local files once before dispatching GPU workers.""" model_config = getattr(od_config, "model_config", None) or {} configured_action_root = model_config.get("lingbot_action_root") or os.environ.get(_ACTION_ROOT_ENV) def pre_process_func(request: OmniDiffusionRequest) -> OmniDiffusionRequest: prompt = request.prompt if not isinstance(prompt, dict): raise ValueError("LingBot World requires a prompt mapping containing multi_modal_data.image.") multi_modal_data = prompt.get("multi_modal_data") or {} additional_information = prompt.get("additional_information") or {} if not isinstance(multi_modal_data, dict): raise ValueError("prompt.multi_modal_data must be a mapping containing image.") if not isinstance(additional_information, dict): raise ValueError("prompt.additional_information must be a mapping.") if "action_path" in additional_information: raise ValueError( "prompt.additional_information.action_path is not supported; use " "sampling_params.extra_args.action_path." ) image = multi_modal_data.get("image") if isinstance(image, list): raise ValueError("LingBot World requires one image and does not accept an image list.") if image is None: raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.") if isinstance(image, (str, os.PathLike)): image = _load_source_image(image) elif isinstance(image, PIL.Image.Image): image = _decode_source_image(image) elif not isinstance(image, torch.Tensor): raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.") extra_args = getattr(request.sampling_params, "extra_args", None) or {} if not isinstance(extra_args, dict): raise ValueError("sampling_params.extra_args must be a mapping.") tick = ARDiffusionTickRequest.from_extra_args(extra_args) if tick is not None: camera_controls = [control for control in tick.controls if control.track == "camera"] if len(camera_controls) != 1: raise ValueError("LingBot realtime ticks require exactly one camera control.") camera_control = camera_controls[0] if camera_control.schema == LINGBOT_CAMERA_TRAJECTORY_SCHEMA: camera_data = camera_control.data try: poses = torch.as_tensor(camera_data["poses"], dtype=torch.float32) intrinsics = torch.as_tensor( camera_data["intrinsics"], dtype=torch.float32, ) except (KeyError, TypeError, ValueError): raise ValueError( "lingbot.camera_trajectory.v1 requires numeric poses and intrinsics fields." ) from None if poses.ndim != 3 or poses.shape[1:] != (4, 4): raise ValueError("camera control poses must have shape [frames, 4, 4].") if intrinsics.ndim != 2 or intrinsics.shape[1:] != (4,): raise ValueError("camera control intrinsics must have shape [frames, 4].") if poses.shape[0] != intrinsics.shape[0] or poses.shape[0] == 0: raise ValueError( "camera control poses and intrinsics must contain the same positive number of frames." ) if not torch.isfinite(poses).all() or not torch.isfinite(intrinsics).all(): raise ValueError("camera control values must all be finite.") trajectory = CameraTrajectory( poses=poses, intrinsics=intrinsics, ) camera_actions = None elif camera_control.schema == LINGBOT_CAMERA_ACTION_SCHEMA: trajectory = None camera_actions = parse_lingbot_camera_action_frames( camera_control.data, expected_frames=3, ) else: raise ValueError( "LingBot realtime camera controls require schema " "'lingbot.camera_trajectory.v1' or 'lingbot.camera_actions.v1'." ) camera_action_script = None elif extra_args.get("camera_action_script") is not None: if extra_args.get("action_path"): raise ValueError("camera_action_script cannot be combined with action_path.") camera_action_script = parse_lingbot_camera_action_script( extra_args.get("camera_action_script"), frames_per_chunk=3, ) trajectory = None camera_actions = None else: action_path = extra_args.get("action_path") if not isinstance(action_path, (str, os.PathLike)) or not str(action_path): raise ValueError("action_path is required in sampling_params.extra_args.action_path.") if not configured_action_root: raise ValueError( "sampling_params.extra_args.action_path requires a trusted action root configured by " f"model_config.lingbot_action_root or {_ACTION_ROOT_ENV}." ) action_directory = resolve_trusted_action_directory( action_path, configured_action_root, ) try: trajectory = load_camera_trajectory(action_directory) except OSError: raise ValueError( "Unable to load camera trajectory from action_path; expected poses.npy and intrinsics.npy." ) from None camera_actions = None camera_action_script = None updated_prompt = dict(prompt) updated_multi_modal_data = dict(multi_modal_data) updated_multi_modal_data["image"] = image updated_prompt["multi_modal_data"] = updated_multi_modal_data request.prompt = updated_prompt request.sampling_params.extra_args = { **extra_args, _PREPROCESSED_CAMERA_KEY: trajectory, _PREPROCESSED_CAMERA_ACTIONS_KEY: camera_actions, _PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY: camera_action_script, } return request return pre_process_func def _fold_camera_embedding( camera_embedding: torch.Tensor, ) -> torch.Tensor: """Pixel-unshuffle ``[frames, 6, H, W]`` onto the Wan latent grid.""" folded = F.pixel_unshuffle(camera_embedding, _CAMERA_SPATIAL_FOLD) return folded.permute(1, 0, 2, 3).unsqueeze(0).contiguous() def get_lingbot_world_post_process_func(od_config: OmniDiffusionConfig) -> Callable[..., Any]: del od_config from diffusers.video_processor import VideoProcessor video_processor = VideoProcessor(vae_scale_factor=8) def post_process_func( video: torch.Tensor, output_type: str = "np", sampling_params: Any | None = None, ) -> Any: if isinstance(video, dict) and isinstance(video.get("payload"), dict): return video if sampling_params is not None: output_type = getattr(sampling_params, "output_type", None) or output_type if output_type == "latent": return video return { "payload": {"video": video_processor.postprocess_video(video, output_type=output_type)}, "metadata": {}, } return post_process_func class LingBotWorldCausalDMDPipeline( nn.Module, SupportImageInput, SupportsComponentDiscovery, SupportsStepExecution, ProgressBarMixin, DiffusionPipelineProfilerMixin, ): """LingBot-World v2 I2V generation with a request-local causal cache.""" supports_step_execution: ClassVar[bool] = True _dit_modules: ClassVar[list[str]] = ["transformer"] _encoder_modules: ClassVar[list[str]] = ["text_encoder"] _vae_modules: ClassVar[list[str]] = ["vae"] # Generic warmup cannot synthesize the required camera action directory. dummy_run_num_frames: ClassVar[int] = 0 _AR_BRANCH = "main" _AR_TEXT_CACHE = "text" def __init__( self, *, od_config: OmniDiffusionConfig, prefix: str = "", ) -> None: super().__init__() del prefix _validate_parallel_config(od_config) self.od_config = od_config self.device = get_local_device() dtype = getattr(od_config, "dtype", torch.bfloat16) model = od_config.model local_files_only = os.path.exists(model) managed_component_placement = bool( getattr(od_config, "enable_cpu_offload", False) or getattr(od_config, "enable_layerwise_offload", False) ) # Standard components use from_pretrained; the custom transformer uses the loader. self.weights_sources = [ DiffusersPipelineLoader.ComponentSource( model_or_path=model, subfolder="transformer", revision=None, prefix="transformer.", fall_back_to_pt=True, ) ] subfolders = ["tokenizer", "text_encoder", "vae"] prefetch_subfolders(model, subfolders, local_files_only=local_files_only) self.tokenizer = from_pretrained_with_prefetch( AutoTokenizer.from_pretrained, model, subfolder="tokenizer", prefetch_list=subfolders, local_files_only=local_files_only, ) self.text_encoder = from_pretrained_with_prefetch( UMT5EncoderModel.from_pretrained, model, subfolder="text_encoder", prefetch_list=subfolders, local_files_only=local_files_only, torch_dtype=dtype, ) if not managed_component_placement: self.text_encoder = self.text_encoder.to(self.device) self.vae = from_pretrained_with_prefetch( DistributedAutoencoderKLWan.from_pretrained, model, subfolder="vae", prefetch_list=subfolders, local_files_only=local_files_only, torch_dtype=dtype, ) if not managed_component_placement: self.vae = self.vae.to(self.device) transformer_config = load_transformer_config(model, "transformer", local_files_only) self.transformer = CausalLingBotWorldTransformer3DModel.from_config( transformer_config, quant_config=getattr(od_config, "quantization_config", None), prefix="transformer", ) scheduler_config = _load_json(model, "scheduler/scheduler_config.json", local_files_only) _validate_scheduler_config(scheduler_config) configured_shift = getattr(od_config, "flow_shift", None) checkpoint_shift = scheduler_config.get("flow_shift", scheduler_config.get("shift", 5.0)) scheduler_shift = _positive_finite_flow_shift( checkpoint_shift if configured_shift is None else configured_shift ) scheduler_keys = { "num_train_timesteps", "solver_order", "prediction_type", "use_dynamic_shifting", "thresholding", "dynamic_thresholding_ratio", "sample_max_value", "predict_x0", "solver_type", "lower_order_final", "disable_corrector", "timestep_spacing", "steps_offset", "final_sigmas_type", } scheduler_kwargs = {name: value for name, value in scheduler_config.items() if name in scheduler_keys} scheduler_kwargs["shift"] = scheduler_shift self.scheduler = FlowUniPCMultistepScheduler(**scheduler_kwargs) self.vae_scale_factor_temporal = int(getattr(self.vae.config, "scale_factor_temporal", 4)) self.vae_scale_factor_spatial = int(getattr(self.vae.config, "scale_factor_spatial", 8)) model_config = getattr(od_config, "model_config", None) or {} self._ar_height = int(model_config.get("ar_diffusion_height", 480)) self._ar_width = int(model_config.get("ar_diffusion_width", 832)) self._ar_diffusion_kv_state: ARDiffusionKVState | None = None self._ar_sessions: dict[str, _LingBotARSessionState] = {} self.setup_diffusion_pipeline_profiler( profiler_targets=[ "vae.encode", "vae.decode", "_generate_block", "text_encoder.forward", "tokenizer.forward", ], enable_diffusion_pipeline_profiler=od_config.enable_diffusion_pipeline_profiler, ) def ar_diffusion_kv_cache_spec(self) -> ARDiffusionKVCacheSpec: """Describe the fixed worker-local cache geometry for realtime ticks.""" patch_frames, patch_height, patch_width = self.transformer.config.patch_size if patch_frames != 1: raise ValueError("LingBot AR-Diffusion requires temporal patch size 1.") spatial = self.vae_scale_factor_spatial if self._ar_height % (spatial * patch_height) or self._ar_width % (spatial * patch_width): raise ValueError("ar_diffusion_height/width must align with VAE and DiT patch sizes.") latent_height = self._ar_height // spatial latent_width = self._ar_width // spatial tokens_per_frame = (latent_height // patch_height) * (latent_width // patch_width) tp_size = get_tensor_model_parallel_world_size() num_local_heads = int(self.transformer.config.num_attention_heads) // tp_size total_window_frames = ( int(self.transformer.config.local_attn_size) if int(self.transformer.config.local_attn_size) != -1 else int(self.transformer.config.sliding_window_num_frames) ) sink_frames = int(self.transformer.config.sink_size) recent_window_frames = total_window_frames - sink_frames if recent_window_frames <= 0: raise ValueError("LingBot AR-Diffusion cache needs a positive recent window after reserving sink frames.") horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 condition_channels = self.vae_scale_factor_temporal + int(self.transformer.config.out_channels) condition_bytes_per_session = ( condition_channels * horizon_latent_frames * latent_height * latent_width * torch.empty((), dtype=self.transformer.dtype).element_size() ) return ARDiffusionKVCacheSpec( num_layers=int(self.transformer.config.num_layers), num_kv_heads=num_local_heads, head_size=int(self.transformer.config.attention_head_dim), tokens_per_frame=tokens_per_frame, frames_per_block=int(self.transformer.config.num_frames_per_block), window_frames=recent_window_frames, sink_frames=sink_frames, kv_branches=(ARDiffusionKVBranchSpec(self._AR_BRANCH, 0),), session_capacity=2, cross_attention=( ARDiffusionCrossAttentionKVSpec( self._AR_TEXT_CACHE, _MAX_SEQUENCE_LENGTH, ), ), model_owned_state_bytes_per_session=condition_bytes_per_session, ) @contextmanager def bind_ar_diffusion_state( self, session_id: str, state: ARDiffusionKVState, ) -> Iterator[None]: if self._ar_diffusion_kv_state is not None: raise RuntimeError("LingBot AR-Diffusion state is already bound.") if state.session_id != session_id: raise ValueError(f"LingBot bound session mismatch: {state.session_id!r} != {session_id!r}.") self._ar_diffusion_kv_state = state try: yield finally: self._ar_diffusion_kv_state = None def reset_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) def close_ar_diffusion_session(self, session_id: str) -> None: self._ar_sessions.pop(session_id, None) def _parse_request(self, req: DiffusionRequestBatch) -> _LingBotRequestInputs: if req.num_reqs != 1 or len(req.prompts) != 1: raise ValueError("LingBot World supports a single prompt request, not request batching.") sampling = req.sampling_params if int(sampling.num_outputs_per_prompt or 1) != 1: raise ValueError("LingBot World requires num_outputs_per_prompt=1.") generator = getattr(sampling, "generator", None) if isinstance(generator, list): raise ValueError("LingBot World accepts one torch.Generator, not a generator list.") if not isinstance(generator, torch.Generator): raise ValueError("LingBot World requires the runner-provided torch.Generator.") if getattr(sampling, "latents", None) is not None: raise ValueError("LingBot World does not support caller-provided latents.") prompt_value = req.prompts[0] multi_modal_data: dict[str, Any] if isinstance(prompt_value, str): prompt = prompt_value multi_modal_data = {} elif isinstance(prompt_value, dict): prompt = prompt_value.get("prompt") or "" multi_modal_data = prompt_value.get("multi_modal_data") or {} else: raise ValueError("prompt must be a string or prompt mapping.") if not isinstance(prompt, str) or not prompt.strip(): raise ValueError("prompt must contain non-empty text.") if not isinstance(multi_modal_data, dict): raise ValueError("prompt.multi_modal_data must be a mapping containing image.") image = multi_modal_data.get("image") if isinstance(image, list): raise ValueError("LingBot World requires one image and does not accept an image list.") if image is None: raise ValueError("LingBot World requires exactly one image in multi_modal_data.image.") if isinstance(image, (str, os.PathLike)): raise ValueError("file-path images must be materialized by the LingBot pre-process function.") if not isinstance(image, (PIL.Image.Image, torch.Tensor)): raise ValueError("multi_modal_data.image must be a PIL image, tensor, or file path.") extra_args = getattr(sampling, "extra_args", None) or {} if not isinstance(extra_args, dict): raise ValueError("sampling_params.extra_args must be a mapping.") camera_trajectory = extra_args.get(_PREPROCESSED_CAMERA_KEY) camera_actions = extra_args.get(_PREPROCESSED_CAMERA_ACTIONS_KEY) camera_action_script = extra_args.get(_PREPROCESSED_CAMERA_ACTION_SCRIPT_KEY) if camera_trajectory is not None and not isinstance(camera_trajectory, CameraTrajectory): raise ValueError("LingBot camera trajectory must be materialized by the pre-process function.") if camera_actions is not None: camera_actions = as_camera_action_frames(camera_actions) if camera_action_script is not None: camera_action_script = as_camera_action_script(camera_action_script) camera_sources = ( int(camera_trajectory is not None) + int(camera_actions is not None) + int(camera_action_script is not None) ) if camera_sources != 1: raise ValueError("LingBot pre-processing must materialize exactly one camera input.") request_flow_shift = ( extra_args["flow_shift"] if "flow_shift" in extra_args else getattr(self.scheduler.config, "shift", 5.0) ) flow_shift = _positive_finite_flow_shift(request_flow_shift) height = getattr(sampling, "height", None) width = getattr(sampling, "width", None) if isinstance(image, PIL.Image.Image): image_width, image_height = image.size elif image.ndim == 3 and image.shape[0] == 3: image_height, image_width = image.shape[-2:] elif image.ndim == 4 and image.shape[:2] == (1, 3): image_height, image_width = image.shape[-2:] else: raise ValueError("tensor image must have shape [3, height, width] or [1, 3, height, width].") if image_height <= 0 or image_width <= 0: raise ValueError("source image width and height must be positive integers.") if image_height * image_width > _MAX_SOURCE_IMAGE_PIXELS: raise ValueError("source image pixel count must not exceed 4096 * 4096.") if (height is None) != (width is None): raise ValueError("height and width must either both be provided or both be omitted.") patch_size = tuple(self.transformer.config.patch_size) if len(patch_size) != 3 or patch_size[0] != 1: raise RuntimeError( "transformer.config.patch_size must be a three-dimensional tuple with temporal patch size 1." ) height_divisor = self.vae_scale_factor_spatial * patch_size[1] width_divisor = self.vae_scale_factor_spatial * patch_size[2] if height is None or width is None: aspect_ratio = image_height / image_width latent_height = round( math.sqrt(_MAX_PIXEL_AREA * aspect_ratio) // self.vae_scale_factor_spatial // patch_size[1] * patch_size[1] ) latent_width = round( math.sqrt(_MAX_PIXEL_AREA / aspect_ratio) // self.vae_scale_factor_spatial // patch_size[2] * patch_size[2] ) height = latent_height * self.vae_scale_factor_spatial width = latent_width * self.vae_scale_factor_spatial if isinstance(height, bool) or not isinstance(height, int) or height <= 0: raise ValueError(f"height must be a positive integer, got {height!r}.") if isinstance(width, bool) or not isinstance(width, int) or width <= 0: raise ValueError(f"width must be a positive integer, got {width!r}.") if height % height_divisor: raise ValueError(f"height must be divisible by {height_divisor}, got {height}.") if width % width_divisor: raise ValueError(f"width must be divisible by {width_divisor}, got {width}.") if height * width > _MAX_PIXEL_AREA: raise ValueError("height * width pixel area must not exceed 480 * 832.") num_inference_steps = getattr(sampling, "num_inference_steps", None) if num_inference_steps is None: num_inference_steps = len(LINGBOT_DMD_TIMESTEPS) if num_inference_steps != len(LINGBOT_DMD_TIMESTEPS): raise ValueError( "num_inference_steps must be 4 for LingBot World causal DMD " f"timesteps {list(LINGBOT_DMD_TIMESTEPS)}, got {num_inference_steps}." ) num_frames = getattr(sampling, "num_frames", None) if isinstance(num_frames, bool) or not isinstance(num_frames, int) or num_frames <= 0: raise ValueError(f"num_frames must be a positive integer, got {num_frames!r}.") if num_frames > _MAX_RAW_FRAMES: raise ValueError(f"num_frames must not exceed {_MAX_RAW_FRAMES}.") temporal_factor = self.vae_scale_factor_temporal if (num_frames - 1) % temporal_factor: raise ValueError( "num_frames must satisfy the causal Wan VAE geometry " f"(num_frames - 1) divisible by {temporal_factor}, got {num_frames}." ) num_latent_frames = (num_frames - 1) // temporal_factor + 1 block_frames = int(self.transformer.config.num_frames_per_block) if num_latent_frames % block_frames: raise ValueError( "num_frames must map to a whole number of configured three-frame latent blocks; " f"got num_frames={num_frames}, latent_frames={num_latent_frames}, block_frames={block_frames}." ) max_sequence_length = getattr(sampling, "max_sequence_length", None) if max_sequence_length is None: max_sequence_length = _MAX_SEQUENCE_LENGTH if isinstance(max_sequence_length, bool) or not isinstance(max_sequence_length, int): raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.") if max_sequence_length != _MAX_SEQUENCE_LENGTH: raise ValueError(f"max_sequence_length must be exactly {_MAX_SEQUENCE_LENGTH}.") return _LingBotRequestInputs( prompt=prompt.strip(), image=image, camera_trajectory=camera_trajectory, camera_actions=camera_actions, camera_action_script=camera_action_script, height=height, width=width, num_frames=num_frames, num_latent_frames=num_latent_frames, output_type=getattr(sampling, "output_type", None) or "np", max_sequence_length=max_sequence_length, flow_shift=flow_shift, generator=generator, ) def _prepare_image_tensor(self, image: PIL.Image.Image | torch.Tensor, *, height: int, width: int) -> torch.Tensor: if isinstance(image, PIL.Image.Image): array = np.asarray(image, dtype=np.float32).copy() image_tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) / 255.0 else: image_tensor = image.detach() if image_tensor.ndim == 3: image_tensor = image_tensor.unsqueeze(0) image_tensor = image_tensor.to(dtype=torch.float32) if not torch.isfinite(image_tensor).all(): raise ValueError("tensor image values must all be finite.") minimum = image_tensor.min().item() maximum = image_tensor.max().item() if minimum >= 0.0 and maximum > 255.0: raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].") if minimum < -1.0 or (minimum < 0.0 and maximum > 1.0): raise ValueError("tensor image values must be in [0, 1], [0, 255], or [-1, 1].") if maximum > 1.0: image_tensor = image_tensor / 255.0 if image_tensor.min().item() >= 0.0: image_tensor = image_tensor * 2.0 - 1.0 if image_tensor.shape[-2:] != (height, width): image_tensor = F.interpolate( image_tensor, size=(height, width), mode="bicubic", align_corners=False, ) return image_tensor.to(device=self.device, dtype=torch.float32) def _vae_latent_stats(self, reference: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: shape = (1, -1, 1, 1, 1) latent_mean = torch.as_tensor( self.vae.config.latents_mean, device=reference.device, dtype=reference.dtype, ).view(*shape) latent_std = torch.as_tensor( self.vae.config.latents_std, device=reference.device, dtype=reference.dtype, ).view(*shape) return latent_mean, latent_std def _prepare_condition(self, inputs: _LingBotRequestInputs, *, dtype: torch.dtype) -> torch.Tensor: """Encode the first frame as ``[mask4, image_latent16]``.""" image = self._prepare_image_tensor(inputs.image, height=inputs.height, width=inputs.width) video_condition = image.new_zeros(1, 3, inputs.num_frames, inputs.height, inputs.width) video_condition[:, :, 0] = image latent_condition = retrieve_latents( self.vae.encode(video_condition.to(dtype=self.vae.dtype)), sample_mode="argmax", ) if latent_condition.shape != ( 1, self.transformer.config.out_channels, inputs.num_latent_frames, inputs.height // self.vae_scale_factor_spatial, inputs.width // self.vae_scale_factor_spatial, ): raise RuntimeError( f"vae.encode returned an incompatible image latent shape: got {tuple(latent_condition.shape)}." ) latent_mean, latent_std = self._vae_latent_stats(latent_condition) latent_condition = (latent_condition - latent_mean) / latent_std temporal_mask = latent_condition.new_zeros( 1, self.vae_scale_factor_temporal, inputs.num_latent_frames, latent_condition.shape[-2], latent_condition.shape[-1], ) temporal_mask[:, :, 0] = 1 condition = torch.cat((temporal_mask, latent_condition), dim=1).to(dtype=dtype) if condition.shape[1] != 20: raise RuntimeError( "LingBot image condition must contain 4 temporal-mask then 16 image-latent channels, " f"got {condition.shape[1]}." ) return condition def _prepare_camera( self, inputs: _LingBotRequestInputs, *, dtype: torch.dtype, previous: CameraTrajectory | None = None, latent_aligned: bool = False, ) -> tuple[torch.Tensor, CameraTrajectory]: """Convert raw camera frames to a latent-aligned ray tensor.""" trajectory = inputs.camera_trajectory if trajectory is None: raise RuntimeError("LingBot camera trajectory was not prepared before camera embedding.") available_frames = int(trajectory.poses.shape[0]) if inputs.camera_actions is not None or latent_aligned: if available_frames != inputs.num_latent_frames: raise ValueError( "camera actions must produce exactly one pose per latent frame; " f"got camera_frames={available_frames}, latent_frames={inputs.num_latent_frames}." ) else: if available_frames < inputs.num_frames: raise ValueError( "camera trajectory frames must be at least num_frames; " f"got camera_frames={available_frames}, num_frames={inputs.num_frames}." ) trajectory = CameraTrajectory( poses=trajectory.poses[: inputs.num_frames], intrinsics=trajectory.intrinsics[: inputs.num_frames], ) trajectory = interpolate_camera_trajectory(trajectory, inputs.num_latent_frames) embedding_trajectory = trajectory drop_anchor = False if previous is not None: embedding_trajectory = CameraTrajectory( poses=torch.cat((previous.poses, trajectory.poses), dim=0), intrinsics=torch.cat( (previous.intrinsics, trajectory.intrinsics), dim=0, ), ) drop_anchor = True elif inputs.camera_actions is not None: # The action integrator returns post-action poses. Keep the first # action visible to framewise-delta conditioning by prepending its # known pre-action state (identity for a new realtime session). identity = torch.eye( 4, device=trajectory.poses.device, dtype=trajectory.poses.dtype, ).unsqueeze(0) embedding_trajectory = CameraTrajectory( poses=torch.cat((identity, trajectory.poses), dim=0), intrinsics=torch.cat( (trajectory.intrinsics[:1], trajectory.intrinsics), dim=0, ), ) drop_anchor = True camera_embedding = build_plucker_embedding( embedding_trajectory, height=inputs.height, width=inputs.width, target_height=inputs.height, target_width=inputs.width, device=self.device, dtype=dtype, ) if drop_anchor: camera_embedding = camera_embedding[1:] tail = CameraTrajectory( poses=trajectory.poses[-1:].clone(), intrinsics=trajectory.intrinsics[-1:].clone(), ) return _fold_camera_embedding(camera_embedding), tail def _ar_text_caches( self, prompt_embeds: torch.Tensor, *, invalidate: bool, ) -> list[LingBotAttentionCache]: state = self._ar_diffusion_kv_state if state is None: raise RuntimeError("LingBot AR text cache requested without a bound state.") if invalidate: state.clear_cross_attention() if not state.is_cross_attention_populated( self._AR_BRANCH, self._AR_TEXT_CACHE, ): projected_text = self.transformer.text_embedding(prompt_embeds) def layer_kv() -> Iterator[tuple[torch.Tensor, torch.Tensor]]: for block in self.transformer.blocks: cross_attention = block.cross_attn key = cross_attention.norm_k(cross_attention.k(projected_text)).unflatten( 2, ( cross_attention.num_local_heads, cross_attention.head_dim, ), ) value = cross_attention.v(projected_text).unflatten( 2, ( cross_attention.num_local_heads, cross_attention.head_dim, ), ) yield key, value state.populate_cross_attention( self._AR_BRANCH, self._AR_TEXT_CACHE, layer_kv(), ) return [ LingBotAttentionCache( key=layer["k"], value=layer["v"], end=prompt_embeds.shape[1], absolute_end=prompt_embeds.shape[1], last_start=0, ) for layer in state.get_cross_attention_kv( self._AR_BRANCH, self._AR_TEXT_CACHE, ) ] # ── DMD block helpers ───────────────────────────────────────────────── # # The DMD math for one AR block lives in ``LingBotDMDBlockRunner`` so that # request mode (``_generate_block`` in a loop) and stepwise execution # (``denoise_step`` / ``step_scheduler`` / ``post_decode``) share one copy. # The pipeline's only jobs here are building the runner and translating # the bound AR session into an ``ARBlockContext``. @property def _dmd_blocks(self) -> LingBotDMDBlockRunner: runner = getattr(self, "_dmd_block_runner", None) if runner is None: runner = LingBotDMDBlockRunner( self.transformer, device=self.device, enforce_eager=bool(self.od_config.enforce_eager), ) self._dmd_block_runner = runner return runner def _ar_block_context( self, cross_attention: list[LingBotAttentionCache] | None, ) -> ARBlockContext | None: """Wrap the bound session for the block runner; ``None`` keeps request-local KV.""" if cross_attention is None: return None state = self._ar_diffusion_kv_state if state is None: raise RuntimeError("LingBot AR cache requested without a bound state.") return ARBlockContext(state=state, cross_attention=cross_attention, branch=self._AR_BRANCH) def _generate_block( self, *, condition: torch.Tensor, camera: torch.Tensor, prompt_embeds: torch.Tensor, cache: LingBotTransformerCache | None, ar_cross_attention: list[LingBotAttentionCache] | None, start_frame: int, schedule: tuple[tuple[float, float], ...], generator: torch.Generator, progress_bar: TqdmProgressBar[Any], ) -> torch.Tensor: return self._dmd_blocks.generate_block( condition=condition, camera=camera, prompt_embeds=prompt_embeds, cache=cache, ar=self._ar_block_context(ar_cross_attention), start_frame=start_frame, schedule=schedule, generator=generator, progress_bar=progress_bar, ) def encode_prompt( self, prompt: str, *, max_sequence_length: int, dtype: torch.dtype, ) -> torch.Tensor: text_inputs = self.tokenizer( [" ".join(prompt.strip().split())], padding="max_length", max_length=max_sequence_length, truncation=True, add_special_tokens=True, return_attention_mask=True, return_tensors="pt", ) input_ids = text_inputs.input_ids.to(self.device) attention_mask = text_inputs.attention_mask.to(self.device) prompt_embeds = self.text_encoder(input_ids, attention_mask).last_hidden_state prompt_embeds = prompt_embeds.to(device=self.device, dtype=dtype) return prompt_embeds * attention_mask.unsqueeze(-1).to(dtype=prompt_embeds.dtype) def forward(self, req: DiffusionRequestBatch) -> DiffusionOutput: inputs = self._parse_request(req) if inputs.camera_action_script is not None: # Only step execution walks a whole rollout inside one request, so # a per-chunk script here would be silently dropped rather than # steering the camera. raise ValueError( "camera_action_script is read only by LingBot step execution; " "request mode takes one three-frame camera_actions control per " "block, or a camera_trajectory for the request." ) tick = ARDiffusionTickRequest.from_extra_args(req.sampling_params.extra_args) if tick is not None and self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot typed ticks require ARDiffusionEngine session binding.") if tick is None and self._ar_diffusion_kv_state is not None: raise ValueError("LingBot ARDiffusionEngine requests must carry ar_diffusion_tick.") session_state: _LingBotARSessionState | None = None if tick is not None: if tick.prompt is not None and " ".join(tick.prompt.split()) != " ".join(inputs.prompt.split()): raise ValueError("ar_diffusion_tick.prompt must match the standard request prompt.") block_frames = int(self.transformer.config.num_frames_per_block) if inputs.num_latent_frames != block_frames: raise ValueError("A LingBot AR-Diffusion request must generate exactly one three-latent-frame block.") if (inputs.height, inputs.width) != ( self._ar_height, self._ar_width, ): raise ValueError( "LingBot AR-Diffusion request resolution must match the " "fixed cache geometry " f"{self._ar_height}x{self._ar_width}." ) horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 max_realtime_ticks = horizon_latent_frames // block_frames if tick.chunk_index >= max_realtime_ticks: raise ValueError( "LingBot realtime generation currently supports at most " f"{max_realtime_ticks} ticks per generation epoch " f"(chunk_index 0 through {max_realtime_ticks - 1}) because " f"the image-condition horizon is {_MAX_RAW_FRAMES} pixel " "frames; reset or create a session to start a new world." ) session_state = self._ar_sessions.setdefault( tick.session_id, _LingBotARSessionState(), ) if tick.chunk_index != session_state.next_chunk_index: raise ValueError( "LingBot chunk_index must be contiguous within a session: " f"expected {session_state.next_chunk_index}, got " f"{tick.chunk_index}." ) schedule = _build_shifted_flow_schedule( flow_shift=inputs.flow_shift, ) dtype = self.transformer.dtype # Phase 1: turn all three user inputs into DiT-ready conditions. prompt_embeds = self.encode_prompt( inputs.prompt, max_sequence_length=inputs.max_sequence_length, dtype=dtype, ) if tick is None: condition = self._prepare_condition(inputs, dtype=dtype) else: assert session_state is not None if session_state.image_condition is None: horizon_latent_frames = (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 session_state.image_condition = self._prepare_condition( replace( inputs, num_frames=_MAX_RAW_FRAMES, num_latent_frames=horizon_latent_frames, ), dtype=dtype, ) block_frames = int(self.transformer.config.num_frames_per_block) condition_start = tick.chunk_index * block_frames condition_stop = condition_start + block_frames if condition_stop > session_state.image_condition.shape[2]: raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") condition = session_state.image_condition[ :, :, condition_start:condition_stop, ] camera_pitch: float | None = None if inputs.camera_actions is not None: assert session_state is not None action_trajectory, camera_pitch = integrate_lingbot_camera_actions( inputs.camera_actions, width=inputs.width, height=inputs.height, initial_pose=(session_state.camera_tail.poses[-1] if session_state.camera_tail is not None else None), initial_pitch=session_state.camera_pitch, ) inputs = replace(inputs, camera_trajectory=action_trajectory) camera, camera_tail = self._prepare_camera( inputs, dtype=dtype, previous=session_state.camera_tail if session_state is not None else None, ) if camera.shape[2:] != condition.shape[2:]: raise RuntimeError( "folded camera and image condition must share latent frame/height/width geometry; " f"got camera={tuple(camera.shape)}, condition={tuple(condition.shape)}." ) block_frames = int(self.transformer.config.num_frames_per_block) cache: LingBotTransformerCache | None = None ar_cross_attention: list[LingBotAttentionCache] | None = None if tick is None: # Direct/offline oracle keeps request-local contiguous K/V. cache = self.transformer.allocate_cache( batch_size=1, latent_height=condition.shape[-2], latent_width=condition.shape[-1], device=self.device, dtype=dtype, ) else: assert session_state is not None prompt_changed = session_state.prompt is not None and session_state.prompt != inputs.prompt if session_state.generator_state is not None: inputs.generator.set_state(session_state.generator_state) ar_cross_attention = self._ar_text_caches( prompt_embeds, invalidate=prompt_changed, ) generated_blocks: list[torch.Tensor] = [] total_steps = (inputs.num_latent_frames // block_frames) * len(LINGBOT_DMD_TIMESTEPS) first_start_frame = tick.chunk_index * block_frames if tick is not None else 0 with self.progress_bar(total=total_steps) as progress_bar: # Phase 3: generate left-to-right in latent-frame blocks. The # monotonically increasing ``start_frame`` is also the temporal # position and cache offset used by the Transformer. for local_start_frame in range( 0, inputs.num_latent_frames, block_frames, ): start_frame = first_start_frame + local_start_frame stop_frame = local_start_frame + block_frames generated_blocks.append( self._generate_block( condition=condition[:, :, local_start_frame:stop_frame], camera=camera[:, :, local_start_frame:stop_frame], prompt_embeds=prompt_embeds, cache=cache, ar_cross_attention=ar_cross_attention, start_frame=start_frame, schedule=schedule, generator=inputs.generator, progress_bar=progress_bar, ) ) generated_latents = torch.cat(generated_blocks, dim=2) # The direct/offline cache is request-local and no longer needed once # all latent blocks have been generated. Release it before VAE decode # so the two large allocations do not overlap. cache = None if tick is not None: assert session_state is not None session_state.prompt = inputs.prompt session_state.generator_state = inputs.generator.get_state() session_state.camera_tail = camera_tail if camera_pitch is not None: session_state.camera_pitch = camera_pitch session_state.next_chunk_index += 1 # Phase 4: either expose model-space latents or invert the checkpoint's # latent normalization and decode to pixel-space video. if tick is not None: if inputs.output_type != "latent": raise ValueError( "LingBot realtime ticks currently require output_type='latent'; " "stateful streaming VAE decode is a separate integration step." ) output = { "payload": {"latents": generated_latents}, "metadata": {"ar_diffusion": ARDiffusionChunkMetadata.from_tick(tick).to_dict()}, } elif inputs.output_type == "latent": output = generated_latents else: latent_mean, latent_std = self._vae_latent_stats(generated_latents) vae_latents = (generated_latents * latent_std + latent_mean).to(dtype=self.vae.dtype) output = self.vae.decode(vae_latents, return_dict=False)[0] if output.shape[2] != inputs.num_frames: raise RuntimeError( "vae.decode returned an incompatible temporal geometry: " f"expected {inputs.num_frames} frames, got {output.shape[2]}." ) return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, ) def _decode_chunk_to_pixels( self, latents: torch.Tensor, *, output_type: str ) -> torch.Tensor | np.ndarray | list[list[PIL.Image.Image]]: """Decode one AR block so streaming consumers receive pixels, not latents. Blocks are decoded independently, which keeps the shared VAE stateless and lets a rollout be served without a per-session temporal cache. That is the same shape Helios streams today; cross-block ``feat_cache`` continuity is a separate decision and is not made here. The registered post-process hook returns a ``{"payload", "metadata"}`` envelope untouched, so the pixel conversion it would do for a bare tensor has to happen here instead. """ latent_mean, latent_std = self._vae_latent_stats(latents) vae_latents = (latents * latent_std + latent_mean).to(dtype=self.vae.dtype) video = self.vae.decode(vae_latents, return_dict=False)[0] return self._video_processor().postprocess_video(video, output_type=output_type) def _video_processor(self) -> VideoProcessor: processor = getattr(self, "_cached_video_processor", None) if processor is None: from diffusers.video_processor import VideoProcessor processor = VideoProcessor(vae_scale_factor=self.vae_scale_factor_spatial) self._cached_video_processor = processor return processor def _horizon_latent_frames(self) -> int: return (_MAX_RAW_FRAMES - 1) // self.vae_scale_factor_temporal + 1 def _max_realtime_chunks(self) -> int: return self._horizon_latent_frames() // int(self.transformer.config.num_frames_per_block) def _require_bound_ar_state(self) -> None: if self._ar_diffusion_kv_state is None: raise RuntimeError("LingBot step execution requires AR-Diffusion session binding.") def prepare_encode(self, state: StepRequestState, **kwargs: Any) -> StepRequestState: del kwargs self._require_bound_ar_state() req = DiffusionRequestBatch( requests=[ OmniDiffusionRequest( prompt=state.prompt, sampling_params=state.sampling, request_id=state.request_id, ) ] ) inputs = self._parse_request(req) if ARDiffusionTickRequest.from_extra_args(state.sampling.extra_args) is not None: raise ValueError("LingBot step execution does not accept AR-Diffusion ticks.") if (inputs.height, inputs.width) != (self._ar_height, self._ar_width): raise ValueError( "LingBot AR-Diffusion request resolution must match the " f"fixed cache geometry {self._ar_height}x{self._ar_width}." ) block_frames = int(self.transformer.config.num_frames_per_block) total_chunks = inputs.num_latent_frames // block_frames max_chunks = self._max_realtime_chunks() if total_chunks > max_chunks: raise ValueError( "LingBot step execution currently supports at most " f"{max_chunks} chunks because the image-condition horizon is " f"{_MAX_RAW_FRAMES} pixel frames." ) if inputs.camera_action_script is not None and len(inputs.camera_action_script) != total_chunks: raise ValueError( "camera_action_script must contain one action list per generated chunk; " f"got {len(inputs.camera_action_script)} chunks for {total_chunks} requested chunks." ) dtype = self.transformer.dtype prompt_embeds = self.encode_prompt( inputs.prompt, max_sequence_length=inputs.max_sequence_length, dtype=dtype, ) image_condition = self._prepare_condition( replace( inputs, num_frames=_MAX_RAW_FRAMES, num_latent_frames=self._horizon_latent_frames(), ), dtype=dtype, ) camera_trajectory_cache = None camera_embedding_cache = None if inputs.camera_trajectory is not None: available_frames = int(inputs.camera_trajectory.poses.shape[0]) if available_frames < inputs.num_frames: raise ValueError( "camera trajectory frames must be at least num_frames; " f"got camera_frames={available_frames}, num_frames={inputs.num_frames}." ) truncated = CameraTrajectory( poses=inputs.camera_trajectory.poses[: inputs.num_frames], intrinsics=inputs.camera_trajectory.intrinsics[: inputs.num_frames], ) camera_trajectory_cache = interpolate_camera_trajectory(truncated, inputs.num_latent_frames) # Embed the whole trajectory once, exactly as request mode does, and # slice it per block in ``_prepare_next_chunk``. Embedding each chunk # on its own would normalize that block's framewise translations by # its own largest step, so a slow block would be conditioned as # full-speed motion and drift from offline replay of the same path. camera_embedding_cache, _ = self._prepare_camera(inputs, dtype=dtype) state.prompt_embeds = prompt_embeds state.chunk_index = 0 state.step_index = 0 state.step_in_chunk = 0 state.total_chunks = total_chunks state.chunk_num_steps = len(LINGBOT_DMD_TIMESTEPS) state.extra = { "inputs": inputs, "image_condition": image_condition, "schedule": _build_shifted_flow_schedule(flow_shift=inputs.flow_shift), "dtype": dtype, "block_frames": block_frames, "camera_pitch": 0.0, "camera_tail": None, "camera_action_script": inputs.camera_action_script, "camera_trajectory_cache": camera_trajectory_cache, "camera_embedding_cache": camera_embedding_cache, } self._ar_text_caches(prompt_embeds, invalidate=False) self._prepare_next_chunk(state) return state def _prepare_next_chunk(self, state: StepRequestState) -> None: extra = state.extra inputs: _LingBotRequestInputs = extra["inputs"] block_frames = int(extra["block_frames"]) start_frame = state.chunk_index * block_frames stop_frame = start_frame + block_frames image_condition = extra["image_condition"] if stop_frame > image_condition.shape[2]: raise ValueError("LingBot chunk_index exceeds the configured causal image condition horizon.") condition = image_condition[:, :, start_frame:stop_frame] previous = extra.get("camera_tail") if extra.get("camera_action_script") is not None: chunk_actions = extra["camera_action_script"][state.chunk_index] action_trajectory, camera_pitch = integrate_lingbot_camera_actions( chunk_actions, width=inputs.width, height=inputs.height, initial_pose=(previous.poses[-1] if previous is not None else None), initial_pitch=float(extra.get("camera_pitch", 0.0)), ) extra["camera_pitch"] = camera_pitch chunk_inputs = replace( inputs, camera_trajectory=action_trajectory, camera_actions=chunk_actions, num_frames=(block_frames - 1) * self.vae_scale_factor_temporal + 1, num_latent_frames=block_frames, ) camera, camera_tail = self._prepare_camera( chunk_inputs, dtype=extra["dtype"], previous=previous, ) elif extra.get("camera_embedding_cache") is not None: # Same slice request mode takes from its one full-trajectory # embedding, so both paths condition a block identically. trajectory = extra["camera_trajectory_cache"] camera = extra["camera_embedding_cache"][:, :, start_frame:stop_frame] camera_tail = CameraTrajectory( poses=trajectory.poses[stop_frame - 1 : stop_frame].clone(), intrinsics=trajectory.intrinsics[stop_frame - 1 : stop_frame].clone(), ) else: raise RuntimeError("LingBot step execution is missing a camera trajectory or action script.") if camera.shape[2:] != condition.shape[2:]: raise RuntimeError( "folded camera and image condition must share latent frame/height/width geometry; " f"got camera={tuple(camera.shape)}, condition={tuple(condition.shape)}." ) generator = getattr(state.sampling, "generator", None) if not isinstance(generator, torch.Generator): raise ValueError("LingBot World requires the runner-provided torch.Generator.") extra["condition"] = condition extra["camera"] = camera extra["camera_tail"] = camera_tail extra["start_frame"] = start_frame extra["ar_cross_attention"] = self._ar_text_caches( cast(torch.Tensor, state.prompt_embeds), invalidate=False, ) state.latents = randn_tensor( ( 1, self.transformer.config.out_channels, condition.shape[2], condition.shape[3], condition.shape[4], ), generator=generator, device=self.device, dtype=torch.float32, ) state.timesteps = torch.tensor( [timestep for timestep, _ in extra["schedule"]], device=self.device, dtype=torch.float32, ) state.step_in_chunk = 0 state.step_index = 0 def denoise_step( self, input_batch: InputBatch, *, states: Sequence[StepRequestState] | None = None, **kwargs: Any, ) -> torch.Tensor | None: del input_batch, kwargs if states is None or len(states) != 1: raise ValueError("LingBot step execution supports a single request, not a batched request.") self._require_bound_ar_state() state = states[0] extra = state.extra latents = state.latents if latents is None: raise RuntimeError("LingBot step execution requires latents before denoise_step.") timestep = state.current_timestep if timestep is None: raise RuntimeError("LingBot step execution requires a current timestep before denoise_step.") schedule = extra["schedule"] step_in_chunk = state.step_in_chunk return self._dmd_blocks.probe_step( current_latents=latents, condition=extra["condition"], camera=extra["camera"], prompt_embeds=cast(torch.Tensor, state.prompt_embeds), cache=None, ar=self._ar_block_context(extra["ar_cross_attention"]), start_frame=int(extra["start_frame"]), timestep_value=float(schedule[step_in_chunk][0]), step_index=step_in_chunk, ) def step_scheduler(self, state: StepRequestState, noise_pred: torch.Tensor, **kwargs: Any) -> None: del kwargs extra = state.extra latents = state.latents if latents is None: raise RuntimeError("LingBot step execution requires latents before step_scheduler.") generator = getattr(state.sampling, "generator", None) if not isinstance(generator, torch.Generator): raise ValueError("LingBot World requires the runner-provided torch.Generator.") schedule = extra["schedule"] step_in_chunk = state.step_in_chunk next_sigma = schedule[step_in_chunk + 1][1] if step_in_chunk + 1 < len(schedule) else None state.latents = self._dmd_blocks.apply_transition( latents, noise_pred, schedule[step_in_chunk][1], next_sigma=next_sigma, generator=generator, ) state.step_in_chunk += 1 state.step_index = state.step_in_chunk def post_decode(self, state: StepRequestState, **kwargs: Any) -> DiffusionOutput: del kwargs self._require_bound_ar_state() extra = state.extra latents = state.latents if latents is None: raise RuntimeError("LingBot step execution requires latents before post_decode.") self._dmd_blocks.commit_block_kv( latents=latents, condition=extra["condition"], camera=extra["camera"], prompt_embeds=cast(torch.Tensor, state.prompt_embeds), cache=None, ar=self._ar_block_context(extra["ar_cross_attention"]), start_frame=int(extra["start_frame"]), ) completed_chunk_index = state.chunk_index inputs: _LingBotRequestInputs = extra["inputs"] # "video" is the payload key the output formatter treats as primary, so # a streaming client receives frames on ``images`` like any other video # pipeline; latent mode keeps the tensor for offline consumers. if inputs.output_type == "latent": payload: dict[str, Any] = {"latents": latents} else: payload = {"video": self._decode_chunk_to_pixels(latents, output_type=inputs.output_type)} output = { "payload": payload, "metadata": { "ar_diffusion": ARDiffusionChunkMetadata( session_id=state.request_id, request_id=state.request_id, chunk_index=completed_chunk_index, applied_event_ids=(), ).to_dict() }, } state.chunk_index += 1 finished = state.request_denoise_completed if not finished: self._prepare_next_chunk(state) return DiffusionOutput( output=output, stage_durations=self.stage_durations if hasattr(self, "stage_durations") else None, chunk_index=completed_chunk_index, total_chunks=state.total_chunks, finished=finished, ) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: loader = AutoWeightsLoader(self) return cast(set[str], loader.load_weights(weights))