3875f8ded4
Signed-off-by: zhangbin.327 <zhangbin.327@bytedance.com> Signed-off-by: Jack47 <scsvip@gmail.com> Co-authored-by: zhangbin.327 <zhangbin.327@bytedance.com> Co-authored-by: Hongsheng Liu <liuhongsheng4@huawei.com>
1293 lines
57 KiB
Python
1293 lines
57 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
"""Request-scoped LingBot-World v2 causal DMD pipeline."""
|
|
|
|
from __future__ import annotations
|
|
|
|
import math
|
|
import os
|
|
from collections.abc import Callable, Iterable, Iterator
|
|
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.forward_context import set_forward_context_denoise_step_idx
|
|
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
|
|
from vllm_omni.diffusion.models.lingbot_world.actions import (
|
|
LINGBOT_CAMERA_ACTION_SCHEMA,
|
|
LINGBOT_CAMERA_TRAJECTORY_SCHEMA,
|
|
integrate_lingbot_camera_actions,
|
|
parse_lingbot_camera_action_frames,
|
|
)
|
|
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.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 tqdm.std import tqdm as TqdmProgressBar
|
|
|
|
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"
|
|
_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."""
|
|
|
|
prompt: str
|
|
image: PIL.Image.Image | torch.Tensor
|
|
camera_trajectory: CameraTrajectory | None
|
|
camera_actions: tuple[tuple[str, ...], ...] | 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'."
|
|
)
|
|
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
|
|
|
|
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,
|
|
}
|
|
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,
|
|
ProgressBarMixin,
|
|
DiffusionPipelineProfilerMixin,
|
|
):
|
|
"""LingBot-World v2 I2V generation with a request-local causal cache."""
|
|
|
|
_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)
|
|
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 = tuple(tuple(actions) for actions in camera_actions)
|
|
if (camera_trajectory is None) == (camera_actions is None):
|
|
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,
|
|
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,
|
|
) -> 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:
|
|
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,
|
|
)
|
|
]
|
|
|
|
def _ar_transformer_cache(
|
|
self,
|
|
*,
|
|
condition: torch.Tensor,
|
|
cross_attention: list[LingBotAttentionCache],
|
|
commit_current: bool,
|
|
) -> LingBotTransformerCache:
|
|
state = self._ar_diffusion_kv_state
|
|
if state is None:
|
|
raise RuntimeError("LingBot AR cache requested without a bound state.")
|
|
patch_frames, patch_height, patch_width = self.transformer.config.patch_size
|
|
frames = condition.shape[2] // patch_frames
|
|
height = condition.shape[3] // patch_height
|
|
width = condition.shape[4] // patch_width
|
|
seq_len = frames * height * width
|
|
return LingBotTransformerCache(
|
|
self_attention=state.get_kv_caches(
|
|
self._AR_BRANCH,
|
|
seq_len=seq_len,
|
|
commit_current=commit_current,
|
|
),
|
|
cross_attention=cross_attention,
|
|
)
|
|
|
|
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:
|
|
block_shape = (
|
|
1,
|
|
self.transformer.config.out_channels,
|
|
condition.shape[2],
|
|
condition.shape[3],
|
|
condition.shape[4],
|
|
)
|
|
current_latents = randn_tensor(
|
|
block_shape,
|
|
generator=generator,
|
|
device=self.device,
|
|
dtype=torch.float32,
|
|
)
|
|
for step_index, (timestep_value, sigma) in enumerate(schedule):
|
|
if not self.od_config.enforce_eager:
|
|
torch.compiler.cudagraph_mark_step_begin()
|
|
set_forward_context_denoise_step_idx(step_index)
|
|
timestep = torch.full((1,), float(timestep_value), device=self.device, dtype=torch.float32)
|
|
# Checkpoint channel contract:
|
|
# [noise/x_t(16), temporal_mask(4), image_latent(16)] -> 36.
|
|
model_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1)
|
|
current_cache = (
|
|
self._ar_transformer_cache(
|
|
condition=condition,
|
|
cross_attention=cast(
|
|
list[LingBotAttentionCache],
|
|
ar_cross_attention,
|
|
),
|
|
commit_current=False,
|
|
)
|
|
if ar_cross_attention is not None
|
|
else cast(LingBotTransformerCache, cache)
|
|
)
|
|
flow_prediction = self.transformer(
|
|
hidden_states=model_input,
|
|
timestep=timestep,
|
|
encoder_hidden_states=prompt_embeds,
|
|
camera_hidden_states=camera,
|
|
cache=current_cache,
|
|
start_frame=start_frame,
|
|
update_cache=False,
|
|
)
|
|
if flow_prediction.shape != current_latents.shape:
|
|
raise RuntimeError(
|
|
"transformer flow prediction shape must match the 16-channel noise latent, "
|
|
f"got {tuple(flow_prediction.shape)} and {tuple(current_latents.shape)}."
|
|
)
|
|
# The checkpoint's flow parameterization is inverted by
|
|
# x0 = x_t - sigma * flow. Intermediate steps re-noise that x0
|
|
# estimate at the next sigma; the final step keeps x0 as this
|
|
# block's generated latent.
|
|
x0 = current_latents - sigma * flow_prediction.float()
|
|
if step_index + 1 < len(schedule):
|
|
next_sigma = schedule[step_index + 1][1]
|
|
noise = randn_tensor(
|
|
current_latents.shape,
|
|
generator=generator,
|
|
device=self.device,
|
|
dtype=torch.float32,
|
|
)
|
|
current_latents = (1.0 - next_sigma) * x0 + next_sigma * noise
|
|
else:
|
|
current_latents = x0
|
|
progress_bar.update()
|
|
|
|
# Commit K/V only for the final clean block, never for noisy probes.
|
|
cache_input = torch.cat((current_latents.to(dtype=condition.dtype), condition), dim=1)
|
|
commit_cache = (
|
|
self._ar_transformer_cache(
|
|
condition=condition,
|
|
cross_attention=cast(
|
|
list[LingBotAttentionCache],
|
|
ar_cross_attention,
|
|
),
|
|
commit_current=True,
|
|
)
|
|
if ar_cross_attention is not None
|
|
else cast(LingBotTransformerCache, cache)
|
|
)
|
|
if not self.od_config.enforce_eager:
|
|
torch.compiler.cudagraph_mark_step_begin()
|
|
self.transformer(
|
|
hidden_states=cache_input,
|
|
timestep=torch.zeros(1, device=self.device, dtype=torch.float32),
|
|
encoder_hidden_states=prompt_embeds,
|
|
camera_hidden_states=camera,
|
|
cache=commit_cache,
|
|
start_frame=start_frame,
|
|
update_cache=True,
|
|
)
|
|
if ar_cross_attention is not None:
|
|
state = self._ar_diffusion_kv_state
|
|
if state is None:
|
|
raise RuntimeError("LingBot AR state disappeared before KV commit.")
|
|
state.commit_paged_context(self._AR_BRANCH)
|
|
return current_latents
|
|
|
|
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)
|
|
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 load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
|
|
loader = AutoWeightsLoader(self)
|
|
return cast(set[str], loader.load_weights(weights))
|