diff --git a/CLAUDE.md b/CLAUDE.md index eb13580..859034a 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ TeleFuser is a high-performance framework for efficient multimodal generation mo **Tech Stack:** Python 3.10-3.13, PyTorch 2.6+, CUDA 12.8+, FastAPI, Ray -**Supported Models:** WanVideo (Wan2.1/2.2), Qwen-Image, Z-Image, FlashVSR, HunyuanVideo, Flux2 Klein, LTX Video, LiveAct, LongCat-Video +**Supported Models:** WanVideo (Wan2.1/2.2), Qwen-Image, Z-Image, FlashVSR, HunyuanVideo, Flux2 Klein, LTX Video, LiveAct, LongCat-Video, LingBot-World ## Commands @@ -41,6 +41,8 @@ telefuser/ │ ├── ltx_video/ # LTX Video: I2V + Audio │ ├── liveact/ # LiveAct: S2V (speech-to-video) │ ├── longcat_video/ # LongCat-Video: T2V, I2V +│ ├── lingbot_world_fast/ # LingBot shared causal-fast engine +│ ├── lingbot_world_v2/ # LingBot-World v2 causal-fast facade │ └── common/ # Shared pipeline utilities ├── models/ # Model architectures: DiT, VAE, text encoders ├── ops/ # Custom operations: attention, FFN, normalization @@ -66,6 +68,11 @@ telefuser/ └── client/ # Python SDK ``` +### LingBot Streaming State + +- Long-running LingBot sessions generate noise and VAE condition latents per chunk; do not retain duration-sized chunk lists. +- Incremental VAE encoder and decoder feature caches must be session-owned so concurrent sessions remain isolated. + ### Layer Architecture Principles For Models TeleFuser's model follows a strict layered architecture for operations: diff --git a/README.md b/README.md index 29aa5e8..0c4456f 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,8 @@ TeleFuser is a high-performance runtime for world model inference and multimodal ## News 📰 +- ✨ **2026-07-15**: Added [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) support for offline generation, interactive WebRTC streaming, and multi-GPU inference. + - ✨ **2026-07-06**: Added external **CacheSeek** latent cache integration for service-mode cross-request reuse. Cache hits can skip the first N denoising steps; the Wan2.2 cache-enabled service example snapshots `[5, 10, 15, 20, 25]` by default. See [docs/en/latent_cache.md](docs/en/latent_cache.md). ## Why TeleFuser @@ -79,7 +81,9 @@ video = pipe( ### 2. Real-Time World Model Demo -TeleFuser includes a bidirectional WebRTC demo for `LingBot-World-Fast`. +TeleFuser includes a bidirectional WebRTC demo for `LingBot-World v2`. +LingBot-World v2 uses camera control and its v2 PPL defaults; its streaming example caps a session at two minutes. + For a laptop browser connected through VS Code Remote SSH, coturn is the only additional system package required; no extra Python package is needed. On Debian or Ubuntu, install it with: @@ -98,15 +102,13 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 \ TELEFUSER_TURN_SERVER='turn:127.0.0.1:3478?transport=tcp' \ TELEFUSER_TURN_USERNAME=telefuser \ TELEFUSER_TURN_CREDENTIAL=telefuser-turn \ -telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ +telefuser stream-serve examples/lingbot/stream_lingbot_world_v2.py \ --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation python examples/stream_server/webrtc_bidirectional_demo.py \ --server-url http://127.0.0.1:8088 \ --port 8091 \ --image-path examples/data/lingbot_world_fast/image.jpg \ - --intrinsics-path examples/data/lingbot_world_fast/intrinsics.npy \ - --frame-num 321 --chunk-size 3 --sample-shift 10.0 --fps 16 \ --turn-url 'turn:localhost:3478?transport=tcp' \ --turn-username telefuser --turn-credential telefuser-turn \ --force-turn-relay --ice-gather-timeout-ms 30000 --no-open @@ -126,7 +128,7 @@ ssh -N -o ExitOnForwardFailure=yes -o ServerAliveInterval=30 \ ``` Then open `http://localhost:8091`. The TURN command and credentials above are development examples. See the -[stream server guide](docs/en/stream_server.md#lingbot-world-fast-streaming) and the +[stream server guide](docs/en/stream_server.md) and the [LingBot example README](examples/lingbot/README.md) for coturn startup and the tested four-H100 setup. If the browser runs on the same physical machine as TeleFuser, no SSH tunnel or TURN server is needed. Unset all @@ -178,7 +180,7 @@ telefuser/ | Pipeline | Task | Notes | |----------|------|-------| -| `LingBot-World-Fast` | Bidirectional world-model streaming | Interactive WebRTC control loop via [examples/lingbot/stream_lingbot_world_fast.py](examples/lingbot/stream_lingbot_world_fast.py) | +| `LingBot-World v2` | Bidirectional world-model streaming | Interactive WebRTC control loop via [examples/lingbot/stream_lingbot_world_v2.py](examples/lingbot/stream_lingbot_world_v2.py) | | `LiveAct` | S2V | Speech-driven talking head generation via [examples/liveact/liveact_s2v_h100.py](examples/liveact/liveact_s2v_h100.py) | | `FlashVSR` | VSR | Streaming video super-resolution via [examples/flashvsr/README.md](examples/flashvsr/README.md) | @@ -219,7 +221,7 @@ See [examples/README.md](examples/README.md) for the example runner and baseline - `AdaTaylorCache` is only calibrated for selected model families. - `torch.compile` support is still experimental in parts of the stack. - Some optimized paths require specific GPU architectures and CUDA versions. -- World-model examples such as `LingBot-World-Fast` require external checkpoints and environment setup. +- World-model examples such as `LingBot-World v2` require external checkpoints and environment setup. - Multi-machine deployment exists in the architecture but may require project-specific integration and validation. ## Development diff --git a/README_zh.md b/README_zh.md index 50c0a10..ee982dd 100644 --- a/README_zh.md +++ b/README_zh.md @@ -17,6 +17,8 @@ TeleFuser 是一个面向世界模型推理与多模态生成的高性能运行 ## News 📰 +- ✨ **2026-07-15**：新增 [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) 支持，支持离线生成、交互式 WebRTC 流和多卡推理。 + - ✨ **2026-07-06**：新增外部 **CacheSeek** latent cache 集成，支持服务模式下跨请求复用；命中后可跳过前 N 步去噪。Wan2.2 服务示例默认快照 `[5, 10, 15, 20, 25]`。配置和安装方式见 [docs/zh/latent_cache.md](docs/zh/latent_cache.md)。 ## 为什么是 TeleFuser @@ -79,7 +81,9 @@ video = pipe( ### 2. 实时世界模型 Demo -TeleFuser 当前提供了 `LingBot-World-Fast` 的双向 WebRTC Demo。 +TeleFuser 当前提供了 `LingBot-World v2` 的双向 WebRTC Demo。 +LingBot-World v2 使用相机控制和 v2 PPL 默认值；其流式示例将单个会话上限设为两分钟。 + 通过 VS Code Remote SSH 从笔记本浏览器访问时，coturn 是唯一需要额外安装的系统软件，不需要增加 Python 包。在 Debian 或 Ubuntu 上执行： @@ -98,15 +102,13 @@ CUDA_VISIBLE_DEVICES=0,1,2,3 \ TELEFUSER_TURN_SERVER='turn:127.0.0.1:3478?transport=tcp' \ TELEFUSER_TURN_USERNAME=telefuser \ TELEFUSER_TURN_CREDENTIAL=telefuser-turn \ -telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ +telefuser stream-serve examples/lingbot/stream_lingbot_world_v2.py \ --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation python examples/stream_server/webrtc_bidirectional_demo.py \ --server-url http://127.0.0.1:8088 \ --port 8091 \ --image-path examples/data/lingbot_world_fast/image.jpg \ - --intrinsics-path examples/data/lingbot_world_fast/intrinsics.npy \ - --frame-num 321 --chunk-size 3 --sample-shift 10.0 --fps 16 \ --turn-url 'turn:localhost:3478?transport=tcp' \ --turn-username telefuser --turn-credential telefuser-turn \ --force-turn-relay --ice-gather-timeout-ms 30000 --no-open @@ -126,7 +128,7 @@ ssh -N -o ExitOnForwardFailure=yes -o ServerAliveInterval=30 \ ``` 然后打开 `http://localhost:8091`。上面的 TURN 账号密码仅作为开发配置示例。coturn 启动方式、生产环境 -注意事项及四张 H100 的实测配置见 [流服务文档](docs/zh/stream_server.md#lingbot-world-fast-流式生成) +注意事项及四张 H100 的实测配置见 [流服务文档](docs/zh/stream_server.md) 和 [LingBot example README](examples/lingbot/README.md)。 如果浏览器和 TeleFuser 服务运行在同一台物理机器上，则不需要 SSH 隧道或 TURN 服务。清除所有 @@ -179,7 +181,7 @@ telefuser/ | Pipeline | 任务 | 说明 | |----------|------|------| -| `LingBot-World-Fast` | 双向世界模型流式推理 | 交互式 WebRTC 控制闭环，见 [examples/lingbot/stream_lingbot_world_fast.py](examples/lingbot/stream_lingbot_world_fast.py) | +| `LingBot-World v2` | 双向世界模型流式推理 | 交互式 WebRTC 控制闭环，见 [examples/lingbot/stream_lingbot_world_v2.py](examples/lingbot/stream_lingbot_world_v2.py) | | `LiveAct` | S2V | 语音驱动数字人视频生成，见 [examples/liveact/liveact_s2v_h100.py](examples/liveact/liveact_s2v_h100.py) | | `FlashVSR` | VSR | 流式视频超分，见 [examples/flashvsr/README.md](examples/flashvsr/README.md) | @@ -220,7 +222,7 @@ telefuser/ - `AdaTaylorCache` 目前只对部分模型家族提供了校准参数。 - `torch.compile` 在部分路径上仍处于实验阶段。 - 一些优化能力依赖特定 GPU 架构和 CUDA 环境。 -- `LingBot-World-Fast` 这类世界模型示例依赖外部权重和额外环境配置。 +- `LingBot-World v2` 这类世界模型示例依赖外部权重和额外环境配置。 - 多机部署在架构上已有支持，但实际落地通常还需要项目级集成与验证。 ## 开发 diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md index 62852d3..35a7995 100644 --- a/examples/lingbot/README.md +++ b/examples/lingbot/README.md @@ -1,18 +1,18 @@ -# LingBot-World-Fast Examples +# LingBot-World Examples -Offline image-to-video and interactive WebRTC streaming generation with LingBot-World-Fast. This page describes -the H100 offline example, the LingBot stream service, and the browser camera controller. +Offline image-to-video and interactive WebRTC streaming generation for LingBot-World-Fast (v1) and LingBot-World +v2. Both variants use the shared causal-fast streaming engine; their checkpoint layout and PPL defaults differ. ## Model Directory -The offline example requires the Wan2.2 I2V base model and the LingBot-World-Fast model. The recommended directory -layout is: +Both offline examples require the Wan2.2 I2V base model. v1 and v2 use separate LingBot checkpoint directories: ```text ${TF_MODEL_ZOO_PATH}/ ├── Wan2.2-I2V-A14B/ └── lingbot/ - └── lingbot-world-fast/ + ├── lingbot-world-fast/ + └── lingbot-world-v2-14b-causal-fast/ ``` Set the model root before running the example: @@ -52,6 +52,20 @@ Default configuration: - Default control directory: `examples/data/lingbot_world_fast/` - Default output: `work_dirs/lingbot_world_fast_i2v_gpu.mp4` +### lingbot_world_v2_image_to_video_h100.py + +Offline camera-controlled v2 generation. The default is 77 frames at 16 FPS: 20 latent frames, exactly five +complete chunks of four. With complete chunk streaming, 81 output frames cannot be represented by `chunk_size=4`. +The v2 checkpoint only supports camera control and uses its PPL-configured Torch SDPA, local attention, sink size, +and timesteps. + +```bash +python examples/lingbot/lingbot_world_v2_image_to_video_h100.py \ + --gpu_num 4 \ + --model_root "${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B" \ + --v2_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-v2-14b-causal-fast" +``` + ## Usage ### Four H100 GPUs @@ -88,7 +102,6 @@ python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ --action_path /path/to/camera_control \ --prompt "A cinematic scene with smooth camera motion" \ --resolution 720p \ - --frame_num 81 \ --fps 16 \ --seed 42 \ --output work_dirs/lingbot_world_fast_custom.mp4 @@ -117,7 +130,6 @@ python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ | `--action_path` | Bundled control directory | Directory containing `poses.npy` and `intrinsics.npy` | | `--prompt` | Bundled English prompt | Positive guidance prompt | | `--resolution` | `480p` | Output resolution; available values are `480p` and `720p` | -| `--frame_num` | `81` | Number of output video frames | | `--fps` | `16` | Output video frame rate | | `--seed` | `42` | Random seed | | `--output` | `work_dirs/lingbot_world_fast_i2v_gpu.mp4` | Output MP4 path | @@ -130,8 +142,8 @@ python examples/lingbot/lingbot_world_fast_image_to_video_h100.py --help ## Notes -- `frame_num` must satisfy the pipeline's complete latent-chunk constraint. The default 81 frames correspond to - 21 latent frames, which are split into seven complete chunks. +- Offline frame counts are fixed in each example's `PPL_CONFIG`: v1 uses 81 frames (seven complete chunks of + three latent frames) and v2 uses 77 frames (five complete chunks of four latent frames). - `--gpu_num` must not exceed the number of GPUs visible to the process. For example, set `CUDA_VISIBLE_DEVICES=0,1,2,3` to select four devices. - The example explicitly closes its parallel workers on exit. If the process is forcibly terminated, check for diff --git a/examples/lingbot/lingbot_world_fast_image_to_video_h100.py b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py index bcb1695..7bed463 100644 --- a/examples/lingbot/lingbot_world_fast_image_to_video_h100.py +++ b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py @@ -47,10 +47,12 @@ RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} PPL_CONFIG = dict( + parallelism=1, control_mode="cam", resolution="480p", frame_num=81, chunk_size=3, + frame_policy="truncate", sample_shift=10.0, seed=42, target_fps=16, @@ -58,13 +60,15 @@ enable_fsdp=False, local_attn_size=-1, sink_size=0, + timestep_indices=(0, 179, 358, 679), max_attention_size=None, + vae_torch_dtype=torch.float32, torch_dtype=torch.bfloat16, ) def get_pipeline( - parallelism: int = 1, + parallelism: int = PPL_CONFIG["parallelism"], model_root: str | None = None, fast_model_root: str | None = None, ) -> LingBotWorldFastPipeline: @@ -82,13 +86,14 @@ def get_pipeline( LingBotWorldFastPipelineConfig( checkpoint_dir=model_root or default_model_root, fast_checkpoint_path=fast_model_root or default_fast_model_root, - vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), + vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=PPL_CONFIG["vae_torch_dtype"]), text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), dit_torch_dtype=dtype, control_type=PPL_CONFIG["control_mode"], max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]], local_attn_size=PPL_CONFIG["local_attn_size"], sink_size=PPL_CONFIG["sink_size"], + timestep_indices=PPL_CONFIG["timestep_indices"], attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]), parallel_config=ParallelConfig( device_ids=list(range(parallelism)) if parallelism > 1 else None, @@ -107,14 +112,12 @@ def run( seed: int = PPL_CONFIG["seed"], resolution: str = PPL_CONFIG["resolution"], action_path: str = DEFAULT_ACTION_PATH, - frame_num: int | None = None, fps: int | None = None, ) -> list[Image.Image]: """Generate a complete offline video through the pipeline core API.""" if resolution not in RESOLUTION_AREAS: raise ValueError(f"Unsupported resolution: {resolution}") pipeline.config.max_area = RESOLUTION_AREAS[resolution] - frame_num = PPL_CONFIG["frame_num"] if frame_num is None else frame_num fps = PPL_CONFIG["target_fps"] if fps is None else fps session_config = LingBotWorldFastSessionConfig( @@ -123,7 +126,8 @@ def run( control_mode=PPL_CONFIG["control_mode"], fps=fps, chunk_size=PPL_CONFIG["chunk_size"], - frame_num=frame_num, + frame_policy=PPL_CONFIG["frame_policy"], + frame_num=PPL_CONFIG["frame_num"], sample_shift=PPL_CONFIG["sample_shift"], seed=seed, max_attention_size=PPL_CONFIG["max_attention_size"], @@ -157,7 +161,7 @@ def run( @click.command() @click.option( "--gpu_num", - default=1, + default=PPL_CONFIG["parallelism"], type=int, help="Number of GPUs used for Ulysses sequence parallelism", ) @@ -166,7 +170,6 @@ def run( @click.option("--prompt", default=DEFAULT_PROMPT, help="Positive guidance prompt") @click.option("--seed", default=PPL_CONFIG["seed"], type=int) @click.option("--resolution", default=PPL_CONFIG["resolution"], type=click.Choice(list(RESOLUTION_AREAS))) -@click.option("--frame_num", default=PPL_CONFIG["frame_num"], type=int, help="Number of output frames") @click.option("--fps", default=PPL_CONFIG["target_fps"], type=int, help="Output video frame rate") @click.option("--model_root", default=None, type=click.Path(exists=True, file_okay=False)) @click.option( @@ -182,7 +185,6 @@ def main( prompt: str, seed: int, resolution: str, - frame_num: int, fps: int, model_root: str, fast_model_root: str, @@ -201,7 +203,6 @@ def main( seed=seed, resolution=resolution, action_path=action_path, - frame_num=frame_num, fps=fps, ) elapsed = time.perf_counter() - start diff --git a/examples/lingbot/lingbot_world_v2_image_to_video_h100.py b/examples/lingbot/lingbot_world_v2_image_to_video_h100.py new file mode 100644 index 0000000..5505d8b --- /dev/null +++ b/examples/lingbot/lingbot_world_v2_image_to_video_h100.py @@ -0,0 +1,237 @@ +"""LingBot-World v2 offline image-to-video example. + +Single GPU: + python examples/lingbot/lingbot_world_v2_image_to_video_h100.py + +Four GPUs with Ulysses sequence parallelism: + python examples/lingbot/lingbot_world_v2_image_to_video_h100.py --gpu_num 4 +""" + +from __future__ import annotations + +import os +import time +from pathlib import Path + +import click +import torch +from PIL import Image + +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.lingbot_world_fast.control import ( + LingBotWorldFastControlBuilder, + LingBotWorldFastOfflineControlSource, + load_camera_control_inputs, + truncate_control_sequence, +) +from telefuser.pipelines.lingbot_world_fast.session import ( + LingBotWorldFastChunkRequest, + LingBotWorldFastGenerationSession, + LingBotWorldFastSessionConfig, + resolve_lingbot_frame_count, +) +from telefuser.pipelines.lingbot_world_v2 import ( + LingBotWorldV2Pipeline, + LingBotWorldV2PipelineConfig, +) +from telefuser.utils.video import save_video + +TF_MODEL_ZOO_PATH = Path(os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo")).expanduser() + +_PROJECT_ROOT = Path(__file__).resolve().parents[2] +_DATA_ROOT = _PROJECT_ROOT / "examples" / "data" / "lingbot_world_fast" +DEFAULT_IMAGE_PATH = str(_DATA_ROOT / "image.jpg") +DEFAULT_ACTION_PATH = str(_DATA_ROOT) +DEFAULT_OUTPUT_DIR = _PROJECT_ROOT / "work_dirs" +DEFAULT_PROMPT = ( + "A serene lakeside scene with a lone tree standing in calm water, surrounded by distant snow-capped " + "mountains under a bright blue sky with drifting white clouds. Gentle ripples reflect the tree and sky." +) +RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} + +PPL_CONFIG = dict( + parallelism=1, + control_mode="cam", + resolution="480p", + frame_num=77, + chunk_size=4, + frame_policy="truncate", + sample_shift=10.0, + seed=42, + target_fps=16, + attn_impl=AttnImplType.TORCH_SDPA, + enable_fsdp=False, + local_attn_size=18, + sink_size=6, + timestep_indices=(0, 250, 500, 750), + max_attention_size=None, + vae_torch_dtype=torch.float32, + torch_dtype=torch.bfloat16, +) + + +def get_pipeline( + parallelism: int = PPL_CONFIG["parallelism"], + model_root: str | None = None, + v2_model_root: str | None = None, +) -> LingBotWorldV2Pipeline: + """Load LingBot-World v2 for offline chunked generation.""" + if model_root is None or v2_model_root is None: + default_model_root = str(TF_MODEL_ZOO_PATH / "Wan2.2-I2V-A14B") + default_v2_model_root = str(TF_MODEL_ZOO_PATH / "lingbot" / "lingbot-world-v2-14b-causal-fast") + else: + default_model_root, default_v2_model_root = model_root, v2_model_root + if parallelism < 1: + raise ValueError(f"parallelism must be positive, got {parallelism}") + dtype = PPL_CONFIG["torch_dtype"] + pipeline = LingBotWorldV2Pipeline(device="cuda", torch_dtype=dtype) + pipeline.init( + LingBotWorldV2PipelineConfig( + checkpoint_dir=model_root or default_model_root, + fast_checkpoint_path=v2_model_root or default_v2_model_root, + vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=PPL_CONFIG["vae_torch_dtype"]), + text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), + dit_torch_dtype=dtype, + control_type=PPL_CONFIG["control_mode"], + max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]], + local_attn_size=PPL_CONFIG["local_attn_size"], + sink_size=PPL_CONFIG["sink_size"], + timestep_indices=PPL_CONFIG["timestep_indices"], + attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]), + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)) if parallelism > 1 else None, + sp_ulysses_degree=parallelism, + enable_fsdp=PPL_CONFIG["enable_fsdp"], + ), + ), + ) + return pipeline + + +def run( + pipeline: LingBotWorldV2Pipeline, + image: Image.Image, + prompt: str, + seed: int = PPL_CONFIG["seed"], + resolution: str = PPL_CONFIG["resolution"], + action_path: str = DEFAULT_ACTION_PATH, + frame_policy: str = PPL_CONFIG["frame_policy"], + fps: int | None = None, +) -> list[Image.Image]: + """Generate a complete offline video through the pipeline core API.""" + if resolution not in RESOLUTION_AREAS: + raise ValueError(f"Unsupported resolution: {resolution}") + pipeline.config.max_area = RESOLUTION_AREAS[resolution] + fps = PPL_CONFIG["target_fps"] if fps is None else fps + + session_config = LingBotWorldFastSessionConfig( + prompt=prompt, + image=image, + control_mode=PPL_CONFIG["control_mode"], + fps=fps, + chunk_size=PPL_CONFIG["chunk_size"], + frame_num=PPL_CONFIG["frame_num"], + sample_shift=PPL_CONFIG["sample_shift"], + seed=seed, + max_attention_size=PPL_CONFIG["max_attention_size"], + frame_policy=frame_policy, + ) + control_context = pipeline.control_context(session_config) + control_builder = LingBotWorldFastControlBuilder(control_context) + poses, intrinsics = load_camera_control_inputs(action_path) + action = None + poses, intrinsics, action = truncate_control_sequence(poses, intrinsics, action, session_config.frame_num) + control_source = LingBotWorldFastOfflineControlSource(control_builder, poses, intrinsics, action) + session = LingBotWorldFastGenerationSession(config=session_config) + frames: list[Image.Image] = [] + try: + for chunk_index in range(control_context.latent_frames // control_context.chunk_size): + result = pipeline( + session, + LingBotWorldFastChunkRequest( + chunk_index=chunk_index, + control=control_source.control_at(chunk_index), + ), + ) + frames.extend(result.frames) + finally: + pipeline.release_session(session) + return frames + + +@click.command() +@click.option( + "--gpu_num", + default=PPL_CONFIG["parallelism"], + type=int, + help="Number of GPUs used for Ulysses sequence parallelism", +) +@click.option("--image_path", default=DEFAULT_IMAGE_PATH, type=click.Path(exists=True)) +@click.option("--action_path", default=DEFAULT_ACTION_PATH, type=click.Path(exists=True, file_okay=False)) +@click.option("--prompt", default=DEFAULT_PROMPT, help="Positive guidance prompt") +@click.option("--seed", default=PPL_CONFIG["seed"], type=int) +@click.option("--resolution", default=PPL_CONFIG["resolution"], type=click.Choice(list(RESOLUTION_AREAS))) +@click.option( + "--frame_policy", + default=PPL_CONFIG["frame_policy"], + type=click.Choice(["strict", "truncate"]), + help="How to handle frame counts that do not align with the latent chunk size", +) +@click.option("--fps", default=PPL_CONFIG["target_fps"], type=int, help="Output video frame rate") +@click.option("--model_root", default=None, type=click.Path(exists=True, file_okay=False)) +@click.option( + "--v2_model_root", + default=None, + type=click.Path(exists=True, file_okay=False), +) +@click.option("--output", default=None, type=click.Path(dir_okay=False), help="Output video path") +def main( + gpu_num: int, + image_path: str, + action_path: str, + prompt: str, + seed: int, + resolution: str, + frame_policy: str, + fps: int, + model_root: str, + v2_model_root: str, + output: str | None, +) -> None: + """Generate an offline video with LingBot-World v2.""" + effective_frame_num, _ = resolve_lingbot_frame_count( + PPL_CONFIG["frame_num"], + PPL_CONFIG["chunk_size"], + frame_policy, + ) + print(f"Requested frames: {PPL_CONFIG['frame_num']}") + print(f"Effective frames: {effective_frame_num}") + pipeline = get_pipeline(gpu_num, model_root, v2_model_root) + try: + image = Image.open(image_path).convert("RGB") + + start = time.perf_counter() + + frames = run( + pipeline, + image, + prompt, + seed=seed, + resolution=resolution, + action_path=action_path, + frame_policy=frame_policy, + fps=fps, + ) + elapsed = time.perf_counter() - start + + output_path = Path(output) if output else DEFAULT_OUTPUT_DIR / f"lingbot_world_v2_i2v_{gpu_num}gpu.mp4" + output_path.parent.mkdir(parents=True, exist_ok=True) + save_video(frames, str(output_path), fps=fps, quality=8) + print(f"Video generation time: {elapsed:.2f} seconds") + print(f"Video saved to: {output_path}") + finally: + pipeline.close() + + +if __name__ == "__main__": + main() diff --git a/examples/lingbot/stream_lingbot_world_fast.py b/examples/lingbot/stream_lingbot_world_fast.py index 2d7667b..b3f1073 100644 --- a/examples/lingbot/stream_lingbot_world_fast.py +++ b/examples/lingbot/stream_lingbot_world_fast.py @@ -27,10 +27,17 @@ control_mode="cam", resolution="480p", target_fps=16, + max_duration_seconds=5.0, + chunk_size=3, + frame_policy="truncate", + sample_shift=10.0, + max_attention_size=None, attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90, enable_fsdp=True, local_attn_size=-1, sink_size=0, + timestep_indices=(0, 179, 358, 679), + vae_torch_dtype=torch.float32, torch_dtype=torch.bfloat16, ) @@ -55,13 +62,14 @@ def get_pipeline( LingBotWorldFastPipelineConfig( checkpoint_dir=model_root or default_model_root, fast_checkpoint_path=fast_model_root or default_fast_model_root, - vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), + vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=PPL_CONFIG["vae_torch_dtype"]), text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), dit_torch_dtype=dtype, control_type=PPL_CONFIG["control_mode"], max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]], local_attn_size=PPL_CONFIG["local_attn_size"], sink_size=PPL_CONFIG["sink_size"], + timestep_indices=PPL_CONFIG["timestep_indices"], attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]), parallel_config=ParallelConfig( device_ids=list(range(parallelism)) if parallelism > 1 else None, @@ -76,4 +84,15 @@ def get_pipeline( def get_service(gpu_num: int = PPL_CONFIG["parallelism"]) -> LingBotWorldFastService: """Build the service loaded by the TeleFuser stream server.""" pipeline = get_pipeline(parallelism=gpu_num) - return LingBotWorldFastService(pipeline, default_fps=PPL_CONFIG["target_fps"]) + return LingBotWorldFastService( + pipeline, + default_fps=PPL_CONFIG["target_fps"], + default_session_config={ + "control_mode": PPL_CONFIG["control_mode"], + "max_duration_seconds": PPL_CONFIG["max_duration_seconds"], + "chunk_size": PPL_CONFIG["chunk_size"], + "frame_policy": PPL_CONFIG["frame_policy"], + "sample_shift": PPL_CONFIG["sample_shift"], + "max_attention_size": PPL_CONFIG["max_attention_size"], + }, + ) diff --git a/examples/lingbot/stream_lingbot_world_v2.py b/examples/lingbot/stream_lingbot_world_v2.py new file mode 100644 index 0000000..0933e76 --- /dev/null +++ b/examples/lingbot/stream_lingbot_world_v2.py @@ -0,0 +1,99 @@ +"""LingBot-World v2 bidirectional streaming service example. + +Run the four-GPU stream service: + TF_MODEL_ZOO_PATH=/path/to/model_zoo \ + telefuser stream-serve examples/lingbot/stream_lingbot_world_v2.py \ + --gpu-num 4 -p 8088 --skip-validation + +Select physical GPUs with CUDA_VISIBLE_DEVICES. For example, use physical GPUs +2 and 3 with --gpu-num 2 and CUDA_VISIBLE_DEVICES=2,3. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.lingbot_world_fast.service import LingBotWorldFastService +from telefuser.pipelines.lingbot_world_v2.pipeline import LingBotWorldV2Pipeline, LingBotWorldV2PipelineConfig + +RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} + +PPL_CONFIG = dict( + parallelism=4, + control_mode="cam", + resolution="480p", + target_fps=16, + max_duration_seconds=120.0, + chunk_size=4, + frame_policy="truncate", + sample_shift=10.0, + max_attention_size=None, + attn_impl=AttnImplType.TORCH_SDPA, + enable_fsdp=False, + local_attn_size=18, + sink_size=6, + timestep_indices=(0, 250, 500, 750), + vae_torch_dtype=torch.float32, + torch_dtype=torch.bfloat16, +) + + +def get_pipeline( + parallelism: int = PPL_CONFIG["parallelism"], + model_root: str | None = None, + v2_model_root: str | None = None, +) -> LingBotWorldV2Pipeline: + """Load LingBot-World v2 with internal multi-GPU workers.""" + if model_root is None or v2_model_root is None: + model_zoo_path = Path(os.environ["TF_MODEL_ZOO_PATH"]).expanduser() + default_model_root = str(model_zoo_path / "Wan2.2-I2V-A14B") + default_v2_model_root = str(model_zoo_path / "lingbot" / "lingbot-world-v2-14b-causal-fast") + else: + default_model_root, default_v2_model_root = model_root, v2_model_root + if parallelism < 1: + raise ValueError(f"parallelism must be positive, got {parallelism}") + dtype = PPL_CONFIG["torch_dtype"] + pipeline = LingBotWorldV2Pipeline(device="cuda", torch_dtype=dtype) + pipeline.init( + LingBotWorldV2PipelineConfig( + checkpoint_dir=model_root or default_model_root, + fast_checkpoint_path=v2_model_root or default_v2_model_root, + vae_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=PPL_CONFIG["vae_torch_dtype"]), + text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), + dit_torch_dtype=dtype, + control_type=PPL_CONFIG["control_mode"], + max_area=RESOLUTION_AREAS[PPL_CONFIG["resolution"]], + local_attn_size=PPL_CONFIG["local_attn_size"], + sink_size=PPL_CONFIG["sink_size"], + timestep_indices=PPL_CONFIG["timestep_indices"], + attention_config=AttentionConfig.dense_attention(PPL_CONFIG["attn_impl"]), + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)) if parallelism > 1 else None, + sp_ulysses_degree=parallelism, + enable_fsdp=PPL_CONFIG["enable_fsdp"], + ), + ), + ) + return pipeline + + +def get_service(gpu_num: int = PPL_CONFIG["parallelism"]) -> LingBotWorldFastService: + """Build the service loaded by the TeleFuser stream server.""" + pipeline = get_pipeline(parallelism=gpu_num) + return LingBotWorldFastService( + pipeline, + default_fps=PPL_CONFIG["target_fps"], + max_generation_seconds=PPL_CONFIG["max_duration_seconds"], + default_session_config={ + "control_mode": PPL_CONFIG["control_mode"], + "max_duration_seconds": PPL_CONFIG["max_duration_seconds"], + "chunk_size": PPL_CONFIG["chunk_size"], + "frame_policy": PPL_CONFIG["frame_policy"], + "sample_shift": PPL_CONFIG["sample_shift"], + "max_attention_size": PPL_CONFIG["max_attention_size"], + }, + ) diff --git a/examples/stream_server/webrtc_bidirectional_demo.py b/examples/stream_server/webrtc_bidirectional_demo.py index 8a8c62f..55f66df 100644 --- a/examples/stream_server/webrtc_bidirectional_demo.py +++ b/examples/stream_server/webrtc_bidirectional_demo.py @@ -606,6 +606,21 @@ #status { text-align: left; } .video-head { align-items: flex-start; flex-direction: column; } } + .field:has(#duration-seconds), + .field:has(#frame-num), + .field:has(#fps), + .field:has(#chunk-size), + .field:has(#seed), + .field:has(#sample-shift), + .field:has(#control-mode), + .field:has(#intrinsics-path), + .field:has(#control-move-step), + .field:has(#control-yaw-step), + .field:has(#control-lateral-step), + .field:has(#show-control-hud), + #duration-help { + display: none; + } @@ -867,35 +882,9 @@ } function requestOptionsFromForm() { - updateFrameNum(); - const options = { - ...DEFAULT_OPTIONS, - fps: numberValue("fps", DEFAULT_OPTIONS.fps ?? 16), - frame_num: numberValue("frame-num", DEFAULT_OPTIONS.frame_num ?? 81), - chunk_size: numberValue("chunk-size", DEFAULT_OPTIONS.chunk_size ?? 3), - sample_shift: numberValue("sample-shift", DEFAULT_OPTIONS.sample_shift ?? 10.0), - seed: numberValue("seed", DEFAULT_OPTIONS.seed ?? 42), - control_mode: $("control-mode").value, - control_move_step: numberValue("control-move-step", DEFAULT_OPTIONS.control_move_step ?? 0.05), - control_yaw_step_degrees: numberValue("control-yaw-step", DEFAULT_OPTIONS.control_yaw_step_degrees ?? 2.0), - control_lateral_step: numberValue("control-lateral-step", DEFAULT_OPTIONS.control_lateral_step ?? 0.05), - control_pitch_step_degrees: DEFAULT_OPTIONS.control_pitch_step_degrees ?? 2.0, - control_pitch_limit_degrees: DEFAULT_OPTIONS.control_pitch_limit_degrees ?? 85.0, - show_control_hud: $("show-control-hud").checked, - }; - const intrinsicsPath = $("intrinsics-path").value.trim(); - if (intrinsicsPath) { - options.intrinsics_path = intrinsicsPath; - } else { - delete options.intrinsics_path; - } - return options; + return { ...DEFAULT_OPTIONS }; } -$("duration-seconds").addEventListener("input", updateFrameNum); -$("fps").addEventListener("input", updateFrameNum); -$("chunk-size").addEventListener("input", updateFrameNum); - function setControlActive(control, active) { const btn = document.querySelector('[data-control="' + control + '"]'); if (btn) btn.classList.toggle("active", active); @@ -1066,11 +1055,13 @@ type: pc.localDescription.type, task: "bidirectional", prompt, - fps: requestOptions.fps, image_path: imagePath, config: { ...requestOptions }, ...requestOptions, }; + if (requestOptions.fps !== undefined) { + requestBody.fps = requestOptions.fps; + } const resp = await fetchJsonWithTimeout( SERVER_URL + "/v1/stream/webrtc/offer", { @@ -1237,25 +1228,7 @@ def main() -> None: if args.force_turn_relay: rtc_config["iceTransportPolicy"] = "relay" - request_options: dict[str, object] = { - "fps": args.fps, - "frame_num": args.frame_num, - "chunk_size": args.chunk_size, - "sample_shift": args.sample_shift, - "seed": args.seed, - "max_sequence_length": args.max_sequence_length, - "control_mode": args.control_mode, - "control_move_step": args.control_move_step, - "control_yaw_step_degrees": args.control_yaw_step_degrees, - "control_lateral_step": args.control_lateral_step, - "control_pitch_step_degrees": args.control_pitch_step_degrees, - "control_pitch_limit_degrees": args.control_pitch_limit_degrees, - "show_control_hud": args.show_control_hud, - } - if args.max_attention_size is not None: - request_options["max_attention_size"] = args.max_attention_size - if args.intrinsics_path: - request_options["intrinsics_path"] = args.intrinsics_path + request_options: dict[str, object] = {} # When proxying, the browser should call the demo origin (no separate port forward needed for --server-url). server_url_for_browser = "" if args.proxy_backend else args.server_url diff --git a/telefuser/models/lingbot_world_fast_dit.py b/telefuser/models/lingbot_world_fast_dit.py index 5ef627c..d7d310c 100644 --- a/telefuser/models/lingbot_world_fast_dit.py +++ b/telefuser/models/lingbot_world_fast_dit.py @@ -383,7 +383,6 @@ def __init__( self.local_attn_size = local_attn_size self.sink_size = sink_size self.control_type = control_type - self.dtype = torch.bfloat16 self.layer_name_list = ["blocks"] control_dim = 6 if control_type == "cam" else 7 @@ -449,19 +448,19 @@ def init_weights(self) -> None: nn.init.zeros_(self.head.head.bias) def _build_timestep_embeddings(self, t: torch.Tensor, seq_len: int) -> tuple[torch.Tensor, torch.Tensor]: - time_dtype = self.time_embedding[0].weight.dtype - if t.dim() == 1: - emb_input = sinusoidal_embedding_1d(self.freq_dim, t).to(dtype=time_dtype) + with torch.amp.autocast(t.device.type, dtype=torch.float32, enabled=t.device.type == "cuda"): + if t.dim() == 1: + emb_input = sinusoidal_embedding_1d(self.freq_dim, t).float() + emb = self.time_embedding(emb_input) + t_mod = self.time_projection(emb).unflatten(1, (6, self.dim)).unsqueeze(1).expand(-1, seq_len, -1, -1) + return emb.unsqueeze(1).float(), t_mod.float() + + flat = t.flatten() + emb_input = sinusoidal_embedding_1d(self.freq_dim, flat).float() emb = self.time_embedding(emb_input) - t_mod = self.time_projection(emb).unflatten(1, (6, self.dim)).unsqueeze(1).expand(-1, seq_len, -1, -1) - return emb.unsqueeze(1).to(self.dtype), t_mod.to(self.dtype) - - flat = t.flatten() - emb_input = sinusoidal_embedding_1d(self.freq_dim, flat).to(dtype=time_dtype) - emb = self.time_embedding(emb_input) - emb = emb.unflatten(0, (t.shape[0], seq_len)) - t_mod = self.time_projection(emb).unflatten(2, (6, self.dim)) - return emb.to(self.dtype), t_mod.to(self.dtype) + emb = emb.unflatten(0, (t.shape[0], seq_len)) + t_mod = self.time_projection(emb).unflatten(2, (6, self.dim)) + return emb.float(), t_mod.float() def _prepare_control_tokens(self, control_tensor: torch.Tensor | None) -> torch.Tensor | None: if control_tensor is None: @@ -513,6 +512,7 @@ def forward( max_attention_size: int = 1_000_000, ) -> torch.Tensor: if y is not None: + y = y.to(device=x.device, dtype=x.dtype) x = torch.cat([x, y], dim=1) x, grid_size = self.patchify(x) @@ -555,7 +555,8 @@ def forward( x = self.head(x, t_head) if self.usp_flag: (x,) = sequence_parallel_unshard(self.device_mesh, [x], [1], [full_seq_len]) - return self.unpatchify(x, grid_size) + output = self.unpatchify(x, grid_size) + return output.float() def enable_usp(self, device_mesh: DeviceMesh | None = None) -> None: """Enable Ulysses sequence parallelism for LingBot-World-Fast DiT.""" @@ -621,11 +622,16 @@ def from_pretrained( if inferred_control_dim not in (6, 7): raise ValueError(f"Unexpected control patch dim: {inferred_control_dim}") inferred_config["control_type"] = "cam" if inferred_control_dim == 6 else "act" + if inferred_config["control_type"] != control_type: + raise ValueError( + f"Checkpoint control projection is {inferred_control_dim}-channel " + f"({inferred_config['control_type']}), not requested {control_type!r}" + ) with init_weights_on_device("meta"): model = cls(**inferred_config) - model.load_state_dict(state_dict, strict=False, assign=True) + model.load_state_dict(state_dict, strict=True, assign=True) model = model.to(dtype=torch_dtype) return model diff --git a/telefuser/models/wan_video_vae.py b/telefuser/models/wan_video_vae.py index 84116e4..915d081 100644 --- a/telefuser/models/wan_video_vae.py +++ b/telefuser/models/wan_video_vae.py @@ -24,6 +24,14 @@ class WanVideoVAEStreamingDecodeState: feat_idx: list[int] = field(default_factory=lambda: [0]) +@dataclass +class WanVideoVAEStreamingEncodeState: + """Session-owned temporal feature cache for incremental VAE encoding.""" + + feat_cache: list[object] = field(default_factory=list) + feat_idx: list[int] = field(default_factory=lambda: [0]) + + def _count_conv3d(model: nn.Module) -> int: """Count Conv3d layers in a model (for feat_cache initialization).""" count = 0 @@ -1338,6 +1346,62 @@ def encode( hidden_states = torch.stack(hidden_states) return hidden_states + def cached_encode_withflag( + self, + video: torch.Tensor, + device: torch.device, + is_first_clip: bool, + is_last_clip: bool, + encode_state: WanVideoVAEStreamingEncodeState, + ) -> torch.Tensor: + """Encode one pixel-space chunk with a session-owned temporal cache.""" + if video.dim() == 4: + video = video.unsqueeze(0) + if video.dim() != 5: + raise ValueError(f"Expected video with four or five dimensions, got shape {tuple(video.shape)}") + + video = video.to(device) + frame_count = video.shape[2] + if is_first_clip: + if frame_count < 1 or (frame_count - 1) % 4: + raise ValueError(f"First VAE encode chunk must contain 4n+1 frames, got {frame_count}") + encode_state.feat_cache = [None] * _count_conv3d(self.model.encoder) + encode_state.feat_idx = [0] + segments = [video[:, :, :1]] + segments.extend(video[:, :, start : start + 4] for start in range(1, frame_count, 4)) + else: + if frame_count < 4 or frame_count % 4: + raise ValueError(f"Subsequent VAE encode chunks must contain 4n frames, got {frame_count}") + if not encode_state.feat_cache: + raise RuntimeError("VAE encode cache must be initialized by the first chunk") + segments = [video[:, :, start : start + 4] for start in range(0, frame_count, 4)] + + encoded_segments = [] + for segment in segments: + encode_state.feat_idx[0] = 0 + encoded_segments.append( + self.model.encoder( + segment, + feat_cache=encode_state.feat_cache, + feat_idx=encode_state.feat_idx, + ) + ) + encoded = torch.cat(encoded_segments, dim=2) + mu, _ = self.model.conv1(encoded).chunk(2, dim=1) + + scale = self.scale + if isinstance(scale[0], torch.Tensor): + scale = [value.to(dtype=mu.dtype, device=mu.device) for value in scale] + mu = (mu - scale[0].view(1, self.z_dim, 1, 1, 1)) * scale[1].view(1, self.z_dim, 1, 1, 1) + else: + scale = scale.to(dtype=mu.dtype, device=mu.device) + mu = (mu - scale[0]) * scale[1] + + if is_last_clip: + encode_state.feat_cache = [] + encode_state.feat_idx = [0] + return mu.squeeze(0) if mu.shape[0] == 1 else mu + def decode( self, hidden_states: torch.Tensor, diff --git a/telefuser/pipelines/lingbot_world_fast/__init__.py b/telefuser/pipelines/lingbot_world_fast/__init__.py index 60d72e0..a633abd 100644 --- a/telefuser/pipelines/lingbot_world_fast/__init__.py +++ b/telefuser/pipelines/lingbot_world_fast/__init__.py @@ -4,7 +4,7 @@ load_action_control_inputs, load_camera_control_inputs, ) -from .denoising import LingBotWorldFastDenoisingStage, LingBotWorldFastTimesteps +from .denoising import LingBotWorldFastDenoisingStage from .pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig from .service import LingBotWorldFastService from .session import ( @@ -27,7 +27,6 @@ "LingBotWorldFastSessionConfig", "LingBotWorldFastSessionState", "LingBotWorldFastSessionStatus", - "LingBotWorldFastTimesteps", "build_action_control_chunk", "build_camera_control_chunk", "load_action_control_inputs", diff --git a/telefuser/pipelines/lingbot_world_fast/control.py b/telefuser/pipelines/lingbot_world_fast/control.py index 531282b..2876b07 100644 --- a/telefuser/pipelines/lingbot_world_fast/control.py +++ b/telefuser/pipelines/lingbot_world_fast/control.py @@ -231,16 +231,6 @@ class LingBotWorldFastControlContext: intrinsics: torch.Tensor -class LingBotWorldFastDeferredControl: - """Lazily materialize a control tensor at the pipeline's legacy execution point.""" - - def __init__(self, factory: Callable[[], torch.Tensor]) -> None: - self._factory = factory - - def __call__(self) -> torch.Tensor: - return self._factory() - - class LingBotWorldFastOfflineControlSource: """Keep offline action data outside the generation session and materialize it once.""" @@ -257,10 +247,10 @@ def __init__( self._action = action self._controls: list[torch.Tensor] | None = None - def control_at(self, chunk_index: int) -> LingBotWorldFastDeferredControl: + def control_at(self, chunk_index: int) -> Callable[[], torch.Tensor]: if chunk_index < 0: raise ValueError(f"chunk_index must be non-negative, got {chunk_index}") - return LingBotWorldFastDeferredControl(lambda: self._materialize()[chunk_index]) + return lambda: self._materialize()[chunk_index] def _materialize(self) -> list[torch.Tensor]: if self._controls is None: @@ -274,9 +264,9 @@ class LingBotWorldFastControlBuilder: def __init__(self, context: LingBotWorldFastControlContext) -> None: self.context = context - def defer(self, action: dict[str, object]) -> LingBotWorldFastDeferredControl: + def defer(self, action: dict[str, object]) -> Callable[[], torch.Tensor]: """Return a control factory for materialization within the pipeline call.""" - return LingBotWorldFastDeferredControl(lambda: self.build(action)) + return lambda: self.build(action) @staticmethod def _align_action_frames(action: torch.Tensor, target_frames: int) -> torch.Tensor: diff --git a/telefuser/pipelines/lingbot_world_fast/denoising.py b/telefuser/pipelines/lingbot_world_fast/denoising.py index 620b851..3d9d93f 100644 --- a/telefuser/pipelines/lingbot_world_fast/denoising.py +++ b/telefuser/pipelines/lingbot_world_fast/denoising.py @@ -13,14 +13,25 @@ from telefuser.utils.logging import logger -@dataclass(frozen=True) -class LingBotWorldFastTimesteps: - indices: tuple[int, ...] = (0, 179, 358, 679) - num_train_timesteps: int = 1000 +def _select_timesteps( + scheduler: FlowUniPCMultistepScheduler, + indices: tuple[int, ...], + shift: float, + num_train_timesteps: int = 1000, +) -> torch.Tensor: + if not indices: + raise ValueError("timestep indices must not be empty") + if any(not isinstance(index, int) or isinstance(index, bool) for index in indices): + raise ValueError(f"timestep indices must be integers, got {indices!r}") + if any(index < 0 or index >= num_train_timesteps for index in indices): + raise ValueError(f"timestep indices must be in [0, {num_train_timesteps}), got {indices!r}") + if tuple(sorted(indices)) != indices or len(set(indices)) != len(indices): + raise ValueError(f"timestep indices must be strictly increasing, got {indices!r}") - def select(self, scheduler: FlowUniPCMultistepScheduler, shift: float) -> torch.Tensor: - scheduler.set_timesteps(self.num_train_timesteps, shift=shift) - return scheduler.timesteps[list(self.indices)].clone() + scheduler.set_timesteps(num_train_timesteps, shift=shift) + if max(indices) >= len(scheduler.timesteps): + raise ValueError(f"timestep index exceeds scheduler output: {indices!r}") + return scheduler.timesteps[list(indices)].clone() @dataclass @@ -112,13 +123,14 @@ def initialize_cache( max_sequence_length: int, sample_shift: float, generator_state: list[int], + timestep_indices: tuple[int, ...] = (0, 179, 358, 679), ) -> bool: """Atomically register session-scoped KV, scheduler, and RNG state.""" if cache_handle in self._cache_registry: raise ValueError(f"Cache handle {cache_handle} is already registered") scheduler = FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=1, use_dynamic_shifting=False) - timesteps = LingBotWorldFastTimesteps().select(scheduler, sample_shift) + timesteps = _select_timesteps(scheduler, tuple(timestep_indices), sample_shift) generator = torch.Generator(device=self.device) generator.set_state(torch.tensor(generator_state, dtype=torch.uint8)) state = _DenoisingCacheState( @@ -168,17 +180,22 @@ def denoise_chunk( for timestep_idx in range(len(timesteps)): schedule_timestep = timesteps[timestep_idx].view(1).to(device=current_latent.device) model_timestep = schedule_timestep.to(dtype=torch.float32) - noise_pred = self.dit( - x=current_latent.to(dtype=self.torch_dtype), - timestep=model_timestep, - context=prompt_emb, - y=condition_chunk, - control_tensor=control_chunk, - kv_cache=self_kv_cache, - crossattn_cache=crossattn_cache, - current_start=current_start, - max_attention_size=max_attention_size, - ) + with torch.amp.autocast( + current_latent.device.type, + dtype=self.torch_dtype, + enabled=current_latent.device.type == "cuda", + ): + noise_pred = self.dit( + x=current_latent.to(dtype=self.torch_dtype), + timestep=model_timestep, + context=prompt_emb, + y=condition_chunk, + control_tensor=control_chunk, + kv_cache=self_kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start, + max_attention_size=max_attention_size, + ) x0 = self._convert_flow_pred_to_x0(noise_pred, current_latent, schedule_timestep[0], scheduler) if timestep_idx < len(timesteps) - 1: next_timestep = timesteps[timestep_idx + 1].view(1).to(device=x0.device) @@ -219,17 +236,22 @@ def denoise_and_update_cache( max_attention_size=max_attention_size, generator=state.generator, ) - self.dit( - x=denoised.to(dtype=self.torch_dtype), - timestep=torch.zeros((1,), dtype=torch.float32, device=self.device), - context=prompt_emb, - y=condition_chunk, - control_tensor=control_chunk, - kv_cache=state.self_kv_cache, - crossattn_cache=state.crossattn_cache, - current_start=current_start, - max_attention_size=max_attention_size, - ) + with torch.amp.autocast( + self.device.type, + dtype=self.torch_dtype, + enabled=self.device.type == "cuda", + ): + self.dit( + x=denoised.to(dtype=self.torch_dtype), + timestep=torch.zeros((1,), dtype=torch.float32, device=self.device), + context=prompt_emb, + y=condition_chunk, + control_tensor=control_chunk, + kv_cache=state.self_kv_cache, + crossattn_cache=state.crossattn_cache, + current_start=current_start, + max_attention_size=max_attention_size, + ) return denoised def has_cache(self, cache_handle: int) -> bool: diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index d2615dc..d224aca 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -28,6 +28,7 @@ LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, LingBotWorldFastSessionStatus, + resolve_lingbot_frame_count, ) @@ -35,7 +36,7 @@ class LingBotWorldFastPipelineConfig: checkpoint_dir: str = "" fast_checkpoint_path: str = "lingbot_world_fast" - vae_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + vae_config: ModelRuntimeConfig = field(default_factory=lambda: ModelRuntimeConfig(torch_dtype=torch.float32)) text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) dit_torch_dtype: torch.dtype = torch.bfloat16 control_type: str = "cam" @@ -44,6 +45,7 @@ class LingBotWorldFastPipelineConfig: max_area: int = 480 * 832 local_attn_size: int = -1 sink_size: int = 0 + timestep_indices: tuple[int, ...] = (0, 179, 358, 679) parallel_config: ParallelConfig = field(default_factory=ParallelConfig) attention_config: AttentionConfig = field(default_factory=AttentionConfig) @@ -92,7 +94,9 @@ def init(self, config: LingBotWorldFastPipelineConfig) -> None: self.text_encoder = WanTextEncoder() self.text_encoder.load_state_dict(load_state_dict(str(checkpoint_root / "models_t5_umt5-xxl-enc-bf16.pth"))) - self.text_encoder = self.text_encoder.to(device=self.text_device, dtype=self.torch_dtype).eval() + self.text_encoder = self.text_encoder.to( + device=self.text_device, dtype=config.text_encoding_config.torch_dtype + ).eval() tokenizer_path = checkpoint_root / "google" / "umt5-xxl" self.tokenizer = HuggingfaceTokenizer(str(tokenizer_path), 512, "whitespace") @@ -101,7 +105,7 @@ def init(self, config: LingBotWorldFastPipelineConfig) -> None: ) self.vae = WanVideoVAE(**vae_cfg) self.vae.load_state_dict(vae_state_dict, strict=False) - self.vae = self.vae.to(device=self.vae_device, dtype=self.torch_dtype).eval() + self.vae = self.vae.to(device=self.vae_device, dtype=config.vae_config.torch_dtype).eval() fast_path = checkpoint_root / config.fast_checkpoint_path dit_device = "cpu" if config.parallel_config.world_size > 1 else self.device @@ -175,25 +179,11 @@ def decode_video_cached( @staticmethod def _best_output_size(w: int, h: int, expected_area: int, dw: int = 16, dh: int = 16) -> tuple[int, int]: - ratio = w / h - ow = math.sqrt(expected_area * ratio) - oh = expected_area / ow - - ow1 = int(ow // dw * dw) - oh1 = int(expected_area / max(ow1, 1) // dh * dh) - oh2 = int(oh // dh * dh) - ow2 = int(expected_area / max(oh2, 1) // dw * dw) - - if ow1 <= 0 or oh1 <= 0: - return max(dw, ow2), max(dh, oh2) - if ow2 <= 0 or oh2 <= 0: - return max(dw, ow1), max(dh, oh1) - - ratio1 = ow1 / oh1 - ratio2 = ow2 / oh2 - if max(ratio / ratio1, ratio1 / ratio) < max(ratio / ratio2, ratio2 / ratio): - return ow1, oh1 - return ow2, oh2 + """Match the LingBot source by independently flooring both dimensions.""" + aspect_ratio = h / w + output_height = int(math.sqrt(expected_area * aspect_ratio) // dh * dh) + output_width = int(math.sqrt(expected_area / aspect_ratio) // dw * dw) + return output_width, output_height def control_context(self, session_config: LingBotWorldFastSessionConfig) -> LingBotWorldFastControlContext: """Return control geometry without allocating a generation runtime.""" @@ -204,8 +194,11 @@ def control_context(self, session_config: LingBotWorldFastSessionConfig) -> Ling self.config.max_area, ) height, width = self.check_resize_height_width(height, width) - frame_num = session_config.frame_num - latent_frames = (frame_num - 1) // 4 + 1 + _, latent_frames = resolve_lingbot_frame_count( + session_config.frame_num, + session_config.chunk_size, + session_config.frame_policy, + ) if session_config.intrinsics is None: focal = float(max(self.config.orig_width, self.config.orig_height)) intrinsics = torch.tensor( @@ -250,17 +243,12 @@ def _validate_session_config(self, session_config: LingBotWorldFastSessionConfig raise ValueError(f"chunk_size must be a positive integer, got {session_config.chunk_size!r}") if session_config.chunk_size < 1: raise ValueError(f"chunk_size must be a positive integer, got {session_config.chunk_size}") - if not isinstance(session_config.frame_num, int) or isinstance(session_config.frame_num, bool): - raise ValueError(f"frame_num must be an integer, got {session_config.frame_num!r}") - frame_num = session_config.frame_num - if frame_num < 1 or (frame_num - 1) % 4: - raise ValueError(f"frame_num must be 4n+1, got {frame_num}") - latent_frames = (frame_num - 1) // 4 + 1 - if latent_frames < session_config.chunk_size or latent_frames % session_config.chunk_size: - raise ValueError( - f"frame_num {frame_num} does not contain a whole number of latent chunks of size " - f"{session_config.chunk_size}" - ) + effective_frame_num, _ = resolve_lingbot_frame_count( + session_config.frame_num, + session_config.chunk_size, + session_config.frame_policy, + ) + session_config.frame_num = effective_frame_num def _validate_control(self, session: LingBotWorldFastGenerationSession, control: torch.Tensor) -> None: expected_device = torch.device(self.device) @@ -279,31 +267,49 @@ def _validate_control(self, session: LingBotWorldFastGenerationSession, control: raise ValueError(f"Control shape must be {expected_shape}, got {tuple(control.shape)}") def _prepare_image_tensor(self, image: Image.Image, height: int, width: int) -> torch.Tensor: - image = image.convert("RGB").resize((width, height), Image.BICUBIC) - array = np.asarray(image, dtype=np.float32) / 255.0 - tensor = ( - torch.from_numpy(array).permute(2, 0, 1).sub_(0.5).div_(0.5).to(self.vae_device, dtype=self.torch_dtype) - ) - return tensor - - def _encode_condition_video(self, image_tensor: torch.Tensor, frame_num: int) -> torch.Tensor: - h, w = image_tensor.shape[1:] - video = torch.cat( - [ - image_tensor.unsqueeze(1), - torch.zeros(3, frame_num - 1, h, w, device=image_tensor.device, dtype=image_tensor.dtype), - ], - dim=1, + """Normalize before applying the source tensor-space bicubic resize.""" + array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + tensor = torch.from_numpy(array).permute(2, 0, 1).sub(0.5).div(0.5) + tensor = torch.nn.functional.interpolate(tensor.unsqueeze(0), size=(height, width), mode="bicubic").squeeze(0) + return tensor.to(self.vae_device, dtype=self.config.vae_config.torch_dtype) + + def _encode_condition_chunk(self, session: LingBotWorldFastGenerationSession) -> torch.Tensor: + """Encode only the current image-conditioning chunk with a persistent VAE cache.""" + chunk_index = session.current_chunk_index + is_first_chunk = chunk_index == 0 + is_last_chunk = chunk_index == session.chunk_count - 1 + pixel_frames = 1 + 4 * (session.chunk_size - 1) if is_first_chunk else 4 * session.chunk_size + video = torch.zeros( + (3, pixel_frames, session.height, session.width), + device=self.vae_device, + dtype=self.config.vae_config.torch_dtype, ) - latent = self.vae.encode([video], device=self.vae_device, tiled=False)[0].to(self.device) + if is_first_chunk: + if session.condition_image is None: + raise RuntimeError("The first condition chunk requires the session image tensor") + video[:, 0] = session.condition_image - latent_h = latent.shape[2] - latent_w = latent.shape[3] - mask = torch.ones(1, frame_num, latent_h, latent_w, device=self.device, dtype=latent.dtype) - mask[:, 1:] = 0 - mask = torch.cat([torch.repeat_interleave(mask[:, 0:1], repeats=4, dim=1), mask[:, 1:]], dim=1) - mask = mask.view(1, mask.shape[1] // 4, 4, latent_h, latent_w).transpose(1, 2)[0] - return torch.cat([mask, latent], dim=0) + latent = self.vae.cached_encode_withflag( + video, + device=self.vae_device, + is_first_clip=is_first_chunk, + is_last_clip=is_last_chunk, + encode_state=session.encoder_state, + ).to(self.device) + if latent.shape[1] != session.chunk_size: + raise RuntimeError( + f"VAE condition chunk has {latent.shape[1]} latent frames, expected {session.chunk_size}" + ) + + mask = torch.zeros( + (4, session.chunk_size, latent.shape[2], latent.shape[3]), + device=latent.device, + dtype=latent.dtype, + ) + if is_first_chunk: + mask[:, 0] = 1 + session.condition_image = None + return torch.cat([mask, latent], dim=0).unsqueeze(0) def _release_session_cache(self, session: LingBotWorldFastGenerationSession) -> bool: cache_handle = session.cache_handle @@ -326,7 +332,10 @@ def release_session(self, session: LingBotWorldFastGenerationSession) -> None: cache_released = self._release_session_cache(session) session.decoder_state.feat_cache = [] session.decoder_state.feat_idx = [0] - session.active = False + session.encoder_state.feat_cache = [] + session.encoder_state.feat_idx = [0] + session.condition_image = None + session.noise_generator = None if not cache_released: session.status = LingBotWorldFastSessionStatus.POISONED session.poisoned_reason = f"Failed to release cache handle {session.cache_handle}" @@ -395,7 +404,7 @@ def _validate_chunk_request( raise RuntimeError(f"Cannot continue poisoned LingBot session: {session.poisoned_reason}") if session.status == LingBotWorldFastSessionStatus.RUNNING: raise RuntimeError("LingBot session already has a chunk in progress") - if not session.active or session.status == LingBotWorldFastSessionStatus.RELEASED: + if session.status == LingBotWorldFastSessionStatus.RELEASED: raise RuntimeError("Cannot generate a chunk from an inactive LingBot session") if request.chunk_index != session.current_chunk_index: raise ValueError( @@ -427,23 +436,23 @@ def _generate_session_chunk( except Exception as exc: session.status = LingBotWorldFastSessionStatus.POISONED session.poisoned_reason = f"{type(exc).__name__}: {exc}" - session.active = False self.release_session(session) raise session.status = LingBotWorldFastSessionStatus.COMMITTED - if not session.active: + done = session.current_chunk_index >= session.chunk_count + if done: self.release_session(session) if session.status == LingBotWorldFastSessionStatus.POISONED: raise RuntimeError(session.poisoned_reason or "Final chunk cleanup failed") logger.info( - f"Generated LingBot chunk {request.chunk_index + 1}/{len(session.noise_chunks)}: {len(chunk_frames)} frames" + f"Generated LingBot chunk {request.chunk_index + 1}/{session.chunk_count}: {len(chunk_frames)} frames" ) return LingBotWorldFastChunkResult( chunk_index=request.chunk_index, frames=chunk_frames, emitted_frames=session.emitted_frames, - done=not session.active, + done=done, session_id=request.session_id, ) @@ -457,6 +466,7 @@ def generate_video( """Drain externally prepared controls through the single-chunk API.""" session = LingBotWorldFastGenerationSession(config=session_config) frames: list[Image.Image] = [] + completed = False try: for chunk_index, control in enumerate(controls): result = self( @@ -468,7 +478,10 @@ def generate_video( progress_callback=progress_callback, ) frames.extend(result.frames) - if session.active: + if result.done: + completed = True + break + if not completed and session.status != LingBotWorldFastSessionStatus.RELEASED: raise ValueError("Control sequence ended before the generation session completed") finally: self.release_session(session) @@ -494,16 +507,6 @@ def _create_initialized_session( self._notify_progress(progress_callback, "preparing_image", width=width, height=height) image_tensor = self._prepare_image_tensor(session_config.image, height, width) - frame_num = session_config.frame_num - self._notify_progress( - progress_callback, - "encoding_condition_video", - frame_num=frame_num, - device=str(self.vae_device), - ) - latent_condition = self._encode_condition_video(image_tensor, frame_num) - self._notify_progress(progress_callback, "condition_video_encoded") - lat_h = control_context.latent_h lat_w = control_context.latent_w lat_f = control_context.latent_frames @@ -524,18 +527,13 @@ def _create_initialized_session( latent_frames=lat_f, latent_height=lat_h, latent_width=lat_w, - total_chunks=math.ceil(lat_f / session_config.chunk_size), - ) - generator = torch.Generator(device=self.device) - generator.manual_seed(int(session_config.seed)) - noise = torch.randn( - (1, 16, lat_f, lat_h, lat_w), - generator=generator, - device=self.device, - dtype=torch.float32, + total_chunks=lat_f // session_config.chunk_size, ) - noise_chunks = list(noise.split(session_config.chunk_size, dim=2)) - condition_chunks = list(latent_condition.unsqueeze(0).split(session_config.chunk_size, dim=2)) + noise_generator = torch.Generator(device=self.device) + noise_generator.manual_seed(int(session_config.seed)) + denoise_generator = torch.Generator(device=self.device) + denoise_seed = (int(session_config.seed) ^ 0x51A7E5EED) & 0x7FFF_FFFF_FFFF_FFFF + denoise_generator.manual_seed(denoise_seed) if before_cache is not None: before_cache() cache_handle = self._next_cache_handle @@ -543,8 +541,8 @@ def _create_initialized_session( session = LingBotWorldFastGenerationSession( prompt_emb=prompt_emb, config=session_config, - noise_chunks=noise_chunks, - condition_chunks=condition_chunks, + condition_image=image_tensor, + noise_generator=noise_generator, latent_h=lat_h, latent_w=lat_w, latent_f=lat_f, @@ -562,7 +560,8 @@ def _create_initialized_session( kv_size=kv_size, max_sequence_length=session_config.max_sequence_length, sample_shift=session_config.sample_shift, - generator_state=generator.get_state().tolist(), + generator_state=denoise_generator.get_state().tolist(), + timestep_indices=getattr(self.config, "timestep_indices", (0, 179, 358, 679)), ) if isinstance(self.denoise_stage, ParallelWorker): self.denoise_stage.initialize_cache(**initialize_cache_kwargs, sync=True) @@ -585,6 +584,16 @@ def _create_initialized_session( session.status = LingBotWorldFastSessionStatus.READY return session + def _next_noise_chunk(self, session: LingBotWorldFastGenerationSession) -> torch.Tensor: + if session.noise_generator is None: + raise RuntimeError("The session noise generator has already been released") + return torch.randn( + (1, 16, session.chunk_size, session.latent_h, session.latent_w), + generator=session.noise_generator, + device=self.device, + dtype=torch.float32, + ) + @torch.inference_mode() def generate_next_chunk( self, @@ -592,14 +601,15 @@ def generate_next_chunk( control: torch.Tensor, progress_callback: Callable[..., None] | None = None, ) -> list[Image.Image]: - if runtime.current_chunk_index >= len(runtime.noise_chunks): - runtime.active = False + if runtime.current_chunk_index >= runtime.chunk_count: return [] with ProfilingContext4Debug("generate_next_chunk"): idx = runtime.current_chunk_index - latent_chunk = runtime.noise_chunks[idx] - condition_chunk = runtime.condition_chunks[idx] + latent_chunk = self._next_noise_chunk(runtime) + self._notify_progress(progress_callback, "encoding_condition_chunk", index=idx) + condition_chunk = self._encode_condition_chunk(runtime) + self._notify_progress(progress_callback, "condition_chunk_encoded", index=idx) control_chunk = control self._notify_progress(progress_callback, "denoising_chunk", index=idx) @@ -638,12 +648,10 @@ def generate_next_chunk( runtime, denoised, is_first_clip=(idx == 0), - is_last_clip=(idx == len(runtime.noise_chunks) - 1), + is_last_clip=(idx == runtime.chunk_count - 1), ) images = self.tensor2video(frames) self._notify_progress(progress_callback, "chunk_decoded", index=idx, frames=len(images)) runtime.current_chunk_index += 1 runtime.emitted_frames += len(images) - if runtime.current_chunk_index >= len(runtime.noise_chunks): - runtime.active = False return images diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index c070f29..f08a188 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -8,7 +8,7 @@ import threading import time import uuid -from collections.abc import AsyncGenerator, Callable +from collections.abc import AsyncGenerator, Callable, Mapping from pathlib import Path import cv2 @@ -26,6 +26,7 @@ LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, LingBotWorldFastSessionState, + resolve_lingbot_frame_count, ) _DIRECTION_ALIASES = { @@ -63,9 +64,19 @@ class LingBotWorldFastService: """Bidirectional WebRTC service for LingBot-World-Fast.""" - def __init__(self, pipeline: LingBotWorldFastPipeline, default_fps: int = 16) -> None: + def __init__( + self, + pipeline: LingBotWorldFastPipeline, + default_fps: int = 16, + default_session_config: Mapping[str, object] | None = None, + max_generation_seconds: float = MAX_GENERATION_SECONDS, + ) -> None: self.pipeline = pipeline self.default_fps = default_fps + self.default_session_config = dict(default_session_config or {}) + if max_generation_seconds <= 0: + raise ValueError(f"max_generation_seconds must be positive, got {max_generation_seconds}") + self.max_generation_seconds = float(max_generation_seconds) self._sessions: dict[str, LingBotWorldFastSessionState] = {} def start(self) -> None: @@ -91,6 +102,23 @@ def _load_image(config: dict) -> Image.Image: return Image.open(image_path).convert("RGB") raise ValueError("LingBotWorldFastService requires 'image' or 'image_path'") + @staticmethod + def _frame_num_for_duration(max_duration_seconds: float, fps: int, chunk_size: int) -> int: + """Return the longest 4n+1 output length with complete latent chunks.""" + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if max_duration_seconds <= 0: + raise ValueError(f"max_duration_seconds must be positive, got {max_duration_seconds}") + max_latent_frames = int(math.floor(max_duration_seconds * fps / 4)) + 1 + latent_frames = (max_latent_frames // chunk_size) * chunk_size + if latent_frames < chunk_size: + raise ValueError( + f"max_duration_seconds={max_duration_seconds} is too short for chunk_size={chunk_size} at fps={fps}" + ) + return 4 * (latent_frames - 1) + 1 + def create_session(self, config: dict) -> str: for stale_session_id, stale_state in list(self._sessions.items()): if not stale_state.active: @@ -99,6 +127,7 @@ def create_session(self, config: dict) -> str: raise RuntimeError( "LingBotWorldFastService supports one active session at a time; stop it before reconnecting" ) + defaults = self.default_session_config session_id = config.get("session_id") or str(uuid.uuid4()) image = self._load_image(config) @@ -106,42 +135,86 @@ def create_session(self, config: dict) -> str: if intrinsics is None and config.get("intrinsics_path"): intrinsics = np.load(Path(config["intrinsics_path"])) - fps = int(config.get("fps") or self.default_fps) - frame_num = int(config.get("frame_num", 81)) + fps_value = config.get("fps", defaults.get("fps", self.default_fps)) + if fps_value is None: + fps_value = self.default_fps + chunk_size_value = config.get("chunk_size", defaults.get("chunk_size", 3)) + if chunk_size_value is None: + chunk_size_value = 3 + max_duration_value = config.get( + "max_duration_seconds", + defaults.get("max_duration_seconds", self.max_generation_seconds), + ) + if max_duration_value is None: + max_duration_value = self.max_generation_seconds + + fps = int(fps_value) + chunk_size = int(chunk_size_value) + max_duration_seconds = float(max_duration_value) if fps <= 0: raise ValueError(f"fps must be positive, got {fps}") + if chunk_size <= 0: + raise ValueError(f"chunk_size must be positive, got {chunk_size}") + if max_duration_seconds <= 0: + raise ValueError(f"max_duration_seconds must be positive, got {max_duration_seconds}") + if max_duration_seconds > self.max_generation_seconds: + raise ValueError(f"max_duration_seconds must not exceed {self.max_generation_seconds:g}") + + frame_policy = str(config.get("frame_policy", defaults.get("frame_policy", "truncate"))) + requested_frame_num = config.get("frame_num") + frame_num = ( + int(requested_frame_num) + if requested_frame_num is not None + else self._frame_num_for_duration(max_duration_seconds, fps, chunk_size) + ) + frame_num, _ = resolve_lingbot_frame_count(frame_num, chunk_size, frame_policy) duration_seconds = (frame_num - 1) / fps - if duration_seconds > MAX_GENERATION_SECONDS: + if duration_seconds > max_duration_seconds: raise ValueError( - f"LingBot streaming duration must not exceed {MAX_GENERATION_SECONDS:g} seconds, " + f"LingBot streaming duration must not exceed {max_duration_seconds:g} seconds, " f"got {duration_seconds:g} seconds" ) session_config = LingBotWorldFastSessionConfig( - prompt=config.get("prompt", ""), + prompt=config.get("prompt", defaults.get("prompt", "")), image=image, - control_mode=config.get("control_mode", "cam"), + control_mode=config.get("control_mode", defaults.get("control_mode", "cam")), fps=fps, - chunk_size=int(config.get("chunk_size", 3)), + chunk_size=chunk_size, frame_num=frame_num, - sample_shift=float(config.get("sample_shift", 10.0)), - seed=int(config.get("seed", 42)), - max_attention_size=config.get("max_attention_size"), - max_sequence_length=int(config.get("max_sequence_length", 512)), + frame_policy=frame_policy, + sample_shift=float(config.get("sample_shift", defaults.get("sample_shift", 10.0))), + seed=int(config.get("seed", defaults.get("seed", 42))), + max_attention_size=config.get("max_attention_size", defaults.get("max_attention_size")), + max_sequence_length=int(config.get("max_sequence_length", defaults.get("max_sequence_length", 512))), intrinsics=intrinsics, - control_move_step=float(config.get("control_move_step", 0.05)), - control_yaw_step_degrees=float(config.get("control_yaw_step_degrees", 2.0)), - control_lateral_step=float(config.get("control_lateral_step", 0.05)), - control_pitch_step_degrees=float(config.get("control_pitch_step_degrees", 2.0)), - control_pitch_limit_degrees=float(config.get("control_pitch_limit_degrees", 85.0)), - show_control_hud=bool(config.get("show_control_hud", True)), + control_move_step=float(config.get("control_move_step", defaults.get("control_move_step", 0.05))), + control_yaw_step_degrees=float( + config.get( + "control_yaw_step_degrees", + defaults.get("control_yaw_step_degrees", 2.0), + ) + ), + control_lateral_step=float(config.get("control_lateral_step", defaults.get("control_lateral_step", 0.05))), + control_pitch_step_degrees=float( + config.get( + "control_pitch_step_degrees", + defaults.get("control_pitch_step_degrees", 2.0), + ) + ), + control_pitch_limit_degrees=float( + config.get( + "control_pitch_limit_degrees", + defaults.get("control_pitch_limit_degrees", 85.0), + ) + ), + show_control_hud=bool(config.get("show_control_hud", defaults.get("show_control_hud", True))), ) control_context = self.pipeline.control_context(session_config) state = LingBotWorldFastSessionState( config=session_config, control_context=control_context, output_queue=asyncio.Queue(), - loop=None, ) self._sessions[session_id] = state logger.info(f"LingBotWorld session created: {session_id}") @@ -491,7 +564,7 @@ def _run_worker_loop( total_chunks=control_context.latent_frames // control_context.chunk_size, ) chunk_index = 0 - while state.active and runtime.active: + while state.active: with state.control_lock: controls_held = bool(state.pressed_controls) if controls_held: @@ -562,6 +635,9 @@ def _run_worker_loop( self._put_output(state, payload) emit_status("chunk_sent", index=chunk_index, frames=len(frames)) chunk_index += 1 + if result.done: + break + except Exception as exc: logger.exception(f"LingBotWorld worker failed: session={session_id}, error={exc}") self._put_output( diff --git a/telefuser/pipelines/lingbot_world_fast/session.py b/telefuser/pipelines/lingbot_world_fast/session.py index 3a25c5a..f9db47b 100644 --- a/telefuser/pipelines/lingbot_world_fast/session.py +++ b/telefuser/pipelines/lingbot_world_fast/session.py @@ -10,11 +10,34 @@ import torch from PIL import Image -from telefuser.models.wan_video_vae import WanVideoVAEStreamingDecodeState +from telefuser.models.wan_video_vae import ( + WanVideoVAEStreamingDecodeState, + WanVideoVAEStreamingEncodeState, +) from .control import LingBotWorldFastControlContext +def resolve_lingbot_frame_count(frame_num: int, chunk_size: int, frame_policy: str) -> tuple[int, int]: + """Return the effective video and latent frame counts for a session.""" + if not isinstance(chunk_size, int) or isinstance(chunk_size, bool) or chunk_size < 1: + raise ValueError(f"chunk_size must be a positive integer, got {chunk_size!r}") + if not isinstance(frame_num, int) or isinstance(frame_num, bool): + raise ValueError(f"frame_num must be an integer, got {frame_num!r}") + if frame_num < 1 or (frame_num - 1) % 4: + raise ValueError(f"frame_num must be 4n+1, got {frame_num}") + if frame_policy not in {"strict", "truncate"}: + raise ValueError(f"frame_policy must be 'strict' or 'truncate', got {frame_policy!r}") + + latent_frames = (frame_num - 1) // 4 + 1 + if frame_policy == "truncate": + latent_frames -= latent_frames % chunk_size + if latent_frames < chunk_size or latent_frames % chunk_size: + raise ValueError(f"frame_num {frame_num} does not contain a whole number of latent chunks of size {chunk_size}") + effective_frame_num = 4 * (latent_frames - 1) + 1 + return effective_frame_num, latent_frames + + @dataclass class LingBotWorldFastSessionConfig: prompt: str @@ -23,6 +46,7 @@ class LingBotWorldFastSessionConfig: fps: int = 16 chunk_size: int = 3 frame_num: int = 81 + frame_policy: str = "truncate" sample_shift: float = 10.0 seed: int = 42 max_attention_size: int | None = None @@ -81,8 +105,9 @@ class LingBotWorldFastGenerationSession: config: LingBotWorldFastSessionConfig prompt_emb: torch.Tensor | None = field(default=None, repr=False) - noise_chunks: list[torch.Tensor] = field(default_factory=list, repr=False) - condition_chunks: list[torch.Tensor] = field(default_factory=list, repr=False) + condition_image: torch.Tensor | None = field(default=None, repr=False) + noise_generator: torch.Generator | None = field(default=None, repr=False) + encoder_state: WanVideoVAEStreamingEncodeState = field(default_factory=WanVideoVAEStreamingEncodeState) latent_h: int = 0 latent_w: int = 0 latent_f: int = 0 @@ -95,7 +120,6 @@ class LingBotWorldFastGenerationSession: decoder_state: WanVideoVAEStreamingDecodeState = field(default_factory=WanVideoVAEStreamingDecodeState) current_chunk_index: int = 0 emitted_frames: int = 0 - active: bool = True status: LingBotWorldFastSessionStatus = LingBotWorldFastSessionStatus.NEW poisoned_reason: str | None = None transaction_lock: object = field(default_factory=threading.RLock, repr=False) @@ -103,6 +127,12 @@ class LingBotWorldFastGenerationSession: world_kv_binding: object | None = None world_kv_cached_latents: dict[int, torch.Tensor] = field(default_factory=dict) + @property + def chunk_count(self) -> int: + if self.chunk_size < 1: + raise RuntimeError("LingBot generation session has not been initialized") + return self.latent_f // self.chunk_size + @dataclass class LingBotWorldFastSessionState: diff --git a/telefuser/pipelines/lingbot_world_v2/__init__.py b/telefuser/pipelines/lingbot_world_v2/__init__.py new file mode 100644 index 0000000..fe7972e --- /dev/null +++ b/telefuser/pipelines/lingbot_world_v2/__init__.py @@ -0,0 +1,13 @@ +"""LingBot-World v2 causal-fast pipeline.""" + +from .pipeline import ( + LingBotWorldV2Pipeline, + LingBotWorldV2PipelineConfig, + resolve_lingbot_world_v2_transformers, +) + +__all__ = [ + "LingBotWorldV2Pipeline", + "LingBotWorldV2PipelineConfig", + "resolve_lingbot_world_v2_transformers", +] diff --git a/telefuser/pipelines/lingbot_world_v2/pipeline.py b/telefuser/pipelines/lingbot_world_v2/pipeline.py new file mode 100644 index 0000000..837ec01 --- /dev/null +++ b/telefuser/pipelines/lingbot_world_v2/pipeline.py @@ -0,0 +1,77 @@ +"""LingBot-World v2 causal-fast facade and checkpoint validation.""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +from safetensors import safe_open + +from telefuser.pipelines.lingbot_world_fast.pipeline import ( + LingBotWorldFastPipeline, + LingBotWorldFastPipelineConfig, +) +from telefuser.pipelines.lingbot_world_fast.session import LingBotWorldFastSessionConfig + + +def resolve_lingbot_world_v2_transformers(checkpoint_path: str | Path) -> Path: + """Validate and return the v2 transformer directory for either accepted layout.""" + candidate = Path(checkpoint_path).expanduser().resolve() + transformers_dir = candidate if candidate.name == "transformers" else candidate / "transformers" + checkpoint_root = transformers_dir.parent + config_path = checkpoint_root / "config.json" + index_path = transformers_dir / "diffusion_pytorch_model.safetensors.index.json" + if not config_path.is_file(): + raise FileNotFoundError(f"LingBot v2 config.json not found: {config_path}") + if not index_path.is_file(): + raise FileNotFoundError(f"LingBot v2 weight index not found: {index_path}") + + index = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = index.get("weight_map") + if not isinstance(weight_map, dict) or not weight_map: + raise ValueError(f"LingBot v2 weight index has no weight_map: {index_path}") + shard_names = set(weight_map.values()) + if not all(isinstance(name, str) for name in shard_names): + raise ValueError(f"LingBot v2 weight index has invalid shard names: {index_path}") + missing = sorted(name for name in shard_names if not (transformers_dir / name).is_file()) + if missing: + raise FileNotFoundError(f"LingBot v2 is missing checkpoint shards: {missing}") + + control_key = "patch_embedding_wancamctrl.weight" + if control_key not in weight_map: + raise ValueError(f"LingBot v2 checkpoint is missing {control_key}") + with safe_open(transformers_dir / weight_map[control_key], framework="pt", device="cpu") as tensors: + control_shape = tuple(tensors.get_tensor(control_key).shape) + if len(control_shape) != 2 or control_shape[1] % (64 * 2 * 2): + raise ValueError(f"LingBot v2 control projection has invalid shape: {control_shape}") + control_channels = control_shape[1] // (64 * 2 * 2) + if control_channels != 6: + raise ValueError(f"LingBot v2 requires a six-channel camera checkpoint, got {control_channels} channels") + return transformers_dir + + +@dataclass +class LingBotWorldV2PipelineConfig(LingBotWorldFastPipelineConfig): + """Configuration for the public LingBot-World v2 causal-fast checkpoint.""" + + fast_checkpoint_path: str = "transformers" + control_type: str = "cam" + local_attn_size: int = 18 + sink_size: int = 6 + timestep_indices: tuple[int, ...] = (0, 250, 500, 750) + + +class LingBotWorldV2Pipeline(LingBotWorldFastPipeline): + """Thin v2 facade over the shared LingBot causal-fast engine.""" + + def init(self, config: LingBotWorldV2PipelineConfig) -> None: + if config.control_type != "cam": + raise ValueError("LingBot-World v2 causal-fast supports camera control only") + transformers_dir = resolve_lingbot_world_v2_transformers( + Path(config.checkpoint_dir) / config.fast_checkpoint_path + if not Path(config.fast_checkpoint_path).is_absolute() + else config.fast_checkpoint_path + ) + config.fast_checkpoint_path = str(transformers_dir) + super().init(config) diff --git a/telefuser/service/api/routers/webrtc.py b/telefuser/service/api/routers/webrtc.py index 6700f85..fe041e5 100644 --- a/telefuser/service/api/routers/webrtc.py +++ b/telefuser/service/api/routers/webrtc.py @@ -54,6 +54,18 @@ def __init__(self, api_server: ApiServer) -> None: configuration=configuration, ) + @staticmethod + def _resolve_fps(svc: object, requested_fps: int | None) -> int: + """Use an explicit request FPS or the underlying service default.""" + if requested_fps is not None: + return requested_fps + + service = getattr(svc, "service", svc) + default_fps = getattr(service, "default_fps", None) + if default_fps is not None: + return int(default_fps) + return 24 + async def handle_offer(self, request: WebRTCOfferRequest) -> WebRTCOfferResponse: svc = self.api.stream_service if svc is None or not svc.is_running: @@ -68,7 +80,7 @@ async def handle_offer(self, request: WebRTCOfferRequest) -> WebRTCOfferResponse async def _handle_server_push(self, svc, request: WebRTCOfferRequest) -> WebRTCOfferResponse: session_id = request.session_id - task_data = request.model_dump(exclude={"sdp", "type"}) + task_data = request.model_dump(exclude={"sdp", "type"}, exclude_none=True) task_data["task_id"] = session_id generator = svc.stream_task(task_data) @@ -79,7 +91,7 @@ async def _handle_server_push(self, svc, request: WebRTCOfferRequest) -> WebRTCO offer_sdp=request.sdp, offer_type=request.type, generator=generator, - fps=request.fps or 24, + fps=self._resolve_fps(svc, request.fps), ) except RuntimeError as exc: raise HTTPException(status_code=503, detail=str(exc)) @@ -117,7 +129,7 @@ async def _handle_bidirectional(self, svc, request: WebRTCOfferRequest) -> WebRT output_generator=output_gen, on_input=lambda sid, chunk: svc.push_chunk(sid, chunk), on_close=lambda sid: svc.close_session(sid), - fps=int(config.get("fps") or 24), + fps=self._resolve_fps(svc, request.fps if request.fps is not None else config.get("fps")), ) except RuntimeError as exc: try: diff --git a/telefuser/service/api/stream_schema.py b/telefuser/service/api/stream_schema.py index 57bc621..d43b902 100644 --- a/telefuser/service/api/stream_schema.py +++ b/telefuser/service/api/stream_schema.py @@ -44,7 +44,7 @@ class WebRTCOfferRequest(BaseModel): type: str = Field(default="offer", description="SDP type") task: str = Field(description="Task type, e.g. t2v, i2v") prompt: str | None = None - fps: int | None = Field(default=24, description="Target video FPS") + fps: int | None = Field(default=None, description="Target video FPS") config: dict = Field(default_factory=dict, description="Session configuration (bidirectional mode)") model_config = {"extra": "allow"} diff --git a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py index a97e73b..960d03d 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py @@ -18,17 +18,15 @@ def _session( *, - active: bool = True, + status: LingBotWorldFastSessionStatus = LingBotWorldFastSessionStatus.READY, current_chunk_index: int = 0, chunk_count: int = 2, ) -> LingBotWorldFastGenerationSession: empty = torch.empty(0) return LingBotWorldFastGenerationSession( config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), - status=LingBotWorldFastSessionStatus.READY, + status=status, prompt_emb=empty, - noise_chunks=[empty for _ in range(chunk_count)], - condition_chunks=[empty for _ in range(chunk_count)], latent_h=1, latent_w=1, latent_f=chunk_count, @@ -38,7 +36,6 @@ def _session( chunk_size=1, max_attention_size=1, cache_handle=7, - active=active, current_chunk_index=current_chunk_index, ) @@ -99,9 +96,9 @@ def test_chunk_request_rejects_invalid_inputs() -> None: LingBotWorldFastChunkRequest(chunk_index=0, control=object()) -def test_pipeline_call_rejects_inactive_runtime() -> None: +def test_pipeline_call_rejects_released_runtime() -> None: pipeline = _pipeline() - runtime = _session(active=False) + runtime = _session(status=LingBotWorldFastSessionStatus.RELEASED) pipeline.generate_next_chunk = MagicMock() with pytest.raises(RuntimeError, match="inactive"): @@ -177,7 +174,6 @@ def test_pipeline_call_releases_runtime_when_generation_fails() -> None: with pytest.raises(RuntimeError, match="generation failed"): pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=_control())) - assert runtime.active is False assert runtime.status == LingBotWorldFastSessionStatus.POISONED assert runtime.poisoned_reason == "RuntimeError: generation failed" pipeline.denoise_stage.release_cache.assert_called_once_with(7) @@ -206,7 +202,6 @@ def test_final_chunk_releases_decoder_state_and_cache() -> None: def generate_next_chunk(session, control, progress_callback=None): session.current_chunk_index = 1 - session.active = False return [] pipeline.generate_next_chunk = MagicMock(side_effect=generate_next_chunk) @@ -268,15 +263,13 @@ def hold_transaction() -> None: def test_generate_video_drains_runtime_and_releases_it() -> None: pipeline = LingBotWorldFastPipeline(device="cpu") - pipeline.release_session = MagicMock(side_effect=lambda state: setattr(state, "active", False)) + pipeline.release_session = MagicMock() frame = Image.new("RGB", (8, 8)) def generate(runtime_state, request, progress_callback=None): runtime_state.current_chunk_index += 1 runtime_state.emitted_frames += 1 - if runtime_state.current_chunk_index == 2: - runtime_state.active = False - return SimpleNamespace(frames=[frame]) + return SimpleNamespace(frames=[frame], done=runtime_state.current_chunk_index == 2) config = LingBotWorldFastSessionConfig(prompt="test", image=frame) diff --git a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py index aafe92b..0b71a41 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py @@ -1,13 +1,18 @@ from types import SimpleNamespace from unittest.mock import MagicMock +import numpy as np import pytest import torch from PIL import Image -from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage, LingBotWorldFastTimesteps -from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline +from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage +from telefuser.pipelines.lingbot_world_fast.pipeline import ( + LingBotWorldFastPipeline, + LingBotWorldFastPipelineConfig, +) from telefuser.pipelines.lingbot_world_fast.session import LingBotWorldFastSessionConfig +from telefuser.pipelines.lingbot_world_v2 import LingBotWorldV2Pipeline, LingBotWorldV2PipelineConfig def _build_runtime_pipeline() -> LingBotWorldFastPipeline: @@ -21,6 +26,7 @@ def _build_runtime_pipeline() -> LingBotWorldFastPipeline: orig_width=16, local_attn_size=-1, sink_size=0, + vae_config=SimpleNamespace(torch_dtype=torch.float32), ) pipeline.dit = SimpleNamespace( patch_size=(1, 2, 2), @@ -30,12 +36,8 @@ def _build_runtime_pipeline() -> LingBotWorldFastPipeline: ) pipeline.denoise_stage = MagicMock() pipeline._next_cache_handle = 0 - pipeline.timesteps = LingBotWorldFastTimesteps() pipeline.encode_prompt = MagicMock(return_value=torch.zeros(1, 4, 8)) pipeline._prepare_image_tensor = MagicMock(return_value=torch.zeros(3, 16, 16)) - pipeline._encode_condition_video = MagicMock( - side_effect=lambda _image, frame_num: torch.zeros(17, (frame_num - 1) // 4 + 1, 2, 2) - ) return pipeline @@ -53,18 +55,46 @@ def _create_runtime(frame_num: int, seed: int = 42): return pipeline, runtime +def test_v1_and_v2_defaults_match_the_shared_source_contract() -> None: + image = Image.new("RGB", (16, 16)) + + assert LingBotWorldFastPipelineConfig().vae_config.torch_dtype == torch.float32 + assert LingBotWorldV2PipelineConfig().vae_config.torch_dtype == torch.float32 + assert LingBotWorldFastSessionConfig(prompt="v1", image=image).frame_policy == "truncate" + + +def test_v1_and_v2_share_source_image_geometry_and_preprocessing() -> None: + pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32) + pipeline.vae_device = torch.device("cpu") + pipeline.config = SimpleNamespace(vae_config=SimpleNamespace(torch_dtype=torch.float32)) + image = Image.fromarray(np.arange(5 * 7 * 3, dtype=np.uint8).reshape(5, 7, 3), mode="RGB") + + actual = pipeline._prepare_image_tensor(image, height=6, width=8) + source_tensor = torch.from_numpy(np.asarray(image, dtype=np.float32) / 255.0).permute(2, 0, 1) + expected = torch.nn.functional.interpolate( + source_tensor.sub(0.5).div(0.5).unsqueeze(0), + size=(6, 8), + mode="bicubic", + ).squeeze(0) + + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + assert LingBotWorldFastPipeline._best_output_size(1024, 768, 480 * 832) == (720, 544) + assert LingBotWorldV2Pipeline._best_output_size is LingBotWorldFastPipeline._best_output_size + assert LingBotWorldV2Pipeline._prepare_image_tensor is LingBotWorldFastPipeline._prepare_image_tensor + + def test_aligned_81_frame_runtime_has_seven_complete_latent_chunks() -> None: pipeline, runtime = _create_runtime(frame_num=81) assert runtime.latent_f == 21 - assert len(runtime.noise_chunks) == 7 - assert len(runtime.condition_chunks) == 7 - assert all(chunk.shape[2] == 3 for chunk in runtime.noise_chunks) - assert all(chunk.shape[2] == 3 for chunk in runtime.condition_chunks) + assert runtime.chunk_count == 7 + assert runtime.noise_generator is not None + assert runtime.condition_image is not None + assert not hasattr(runtime, "noise_chunks") + assert not hasattr(runtime, "condition_chunks") assert runtime.cache_handle == 0 assert not hasattr(runtime, "self_kv_cache") pipeline.denoise_stage.initialize_cache.assert_called_once() - pipeline._encode_condition_video.assert_called_once() def test_generation_sessions_receive_isolated_cache_and_decoder_handles() -> None: @@ -102,13 +132,17 @@ def test_cache_initialization_failure_triggers_global_cleanup() -> None: def test_runtime_noise_is_reproducible_and_seed_dependent() -> None: - _, first = _create_runtime(frame_num=21, seed=7) - _, repeated = _create_runtime(frame_num=21, seed=7) - _, different = _create_runtime(frame_num=21, seed=8) + first_pipeline, first = _create_runtime(frame_num=21, seed=7) + repeated_pipeline, repeated = _create_runtime(frame_num=21, seed=7) + different_pipeline, different = _create_runtime(frame_num=21, seed=8) - first_noise = torch.cat(first.noise_chunks, dim=2) - repeated_noise = torch.cat(repeated.noise_chunks, dim=2) - different_noise = torch.cat(different.noise_chunks, dim=2) + first_noise = torch.cat([first_pipeline._next_noise_chunk(first) for _ in range(first.chunk_count)], dim=2) + repeated_noise = torch.cat( + [repeated_pipeline._next_noise_chunk(repeated) for _ in range(repeated.chunk_count)], dim=2 + ) + different_noise = torch.cat( + [different_pipeline._next_noise_chunk(different) for _ in range(different.chunk_count)], dim=2 + ) torch.testing.assert_close(first_noise, repeated_noise) assert not torch.equal(first_noise, different_noise) @@ -153,7 +187,7 @@ def denoise(active_generator: torch.Generator) -> torch.Tensor: assert not torch.equal(first, second) -def test_final_chunk_marks_runtime_inactive_and_releases_denoise_state() -> None: +def test_final_chunk_reaches_derived_chunk_boundary() -> None: pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32) pipeline.vae_device = torch.device("cpu") pipeline.denoise_stage = object() @@ -161,16 +195,18 @@ def test_final_chunk_marks_runtime_inactive_and_releases_denoise_state() -> None expected_frames = [Image.new("RGB", (8, 8)) for _ in range(9)] pipeline.tensor2video = MagicMock(return_value=expected_frames) cached_latent = torch.zeros(1, 1, 3, 2, 2) + pipeline._encode_condition_chunk = MagicMock(return_value=torch.zeros_like(cached_latent)) runtime = SimpleNamespace( current_chunk_index=0, - noise_chunks=[torch.zeros_like(cached_latent)], - condition_chunks=[torch.zeros_like(cached_latent)], + chunk_count=1, + noise_generator=torch.Generator(device="cpu").manual_seed(42), + latent_h=2, + latent_w=2, config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))), chunk_size=3, frame_tokens=1, world_kv_cached_latents={0: cached_latent}, world_kv_binding=None, - active=True, emitted_frames=0, ) @@ -179,12 +215,29 @@ def test_final_chunk_marks_runtime_inactive_and_releases_denoise_state() -> None assert frames == expected_frames assert runtime.current_chunk_index == 1 assert runtime.emitted_frames == 9 - assert runtime.active is False + assert runtime.current_chunk_index == runtime.chunk_count + + +def test_runtime_truncates_non_aligned_latent_frame_count() -> None: + _, runtime = _create_runtime(frame_num=13) + assert runtime.latent_f == 3 + assert runtime.config.frame_num == 9 + assert runtime.chunk_count == 1 -def test_runtime_rejects_non_aligned_frame_count() -> None: + +def test_strict_frame_policy_rejects_non_aligned_latent_frame_count() -> None: + pipeline = _build_runtime_pipeline() with pytest.raises(ValueError, match="frame_num"): - _create_runtime(frame_num=13) + pipeline._create_initialized_session( + LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + frame_num=13, + chunk_size=3, + frame_policy="strict", + ) + ) def test_runtime_rejects_frame_count_smaller_than_first_chunk() -> None: diff --git a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py index 6ea0de7..02e8968 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -46,8 +46,7 @@ def control_context(*args, **kwargs): def generate_chunk(runtime, request, progress_callback=None): assert request.control is deferred_control - runtime.active = False - return SimpleNamespace(frames=[Image.new("RGB", (8, 8))]) + return SimpleNamespace(frames=[Image.new("RGB", (8, 8))], done=True) pipeline.control_context.side_effect = control_context pipeline.side_effect = generate_chunk @@ -230,6 +229,45 @@ def test_create_session_limits_stream_generation_to_20_seconds() -> None: ) +def test_create_session_uses_truncated_frame_count_for_duration_validation() -> None: + pipeline = MagicMock() + service = LingBotWorldFastService(pipeline) + + session_id = service.create_session( + { + "image": Image.new("RGB", (8, 8)), + "fps": 16, + "chunk_size": 3, + "frame_num": 13, + "max_duration_seconds": 0.5, + } + ) + + assert service._sessions[session_id].config.frame_num == 9 + assert pipeline.control_context.call_args.args[0].frame_num == 9 + service.close_session(session_id) + + +@pytest.mark.parametrize( + ("field", "value", "message"), + [ + ("fps", 0, "fps must be positive"), + ("chunk_size", 0, "chunk_size must be positive"), + ("chunk_size", -1, "chunk_size must be positive"), + ("max_duration_seconds", 0, "max_duration_seconds must be positive"), + ("max_duration_seconds", -1, "max_duration_seconds must be positive"), + ], +) +def test_create_session_rejects_non_positive_stream_parameters(field: str, value: int, message: str) -> None: + service = LingBotWorldFastService(MagicMock()) + + request = {"image": Image.new("RGB", (8, 8)), field: value} + if field == "max_duration_seconds": + request["frame_num"] = 9 + with pytest.raises(ValueError, match=message): + service.create_session(request) + + def test_create_session_initializes_fixed_intrinsics_from_intrinsics_path() -> None: pipeline = MagicMock() service = LingBotWorldFastService(pipeline) diff --git a/tests/unit/pipelines/lingbot_world_v2/test_service.py b/tests/unit/pipelines/lingbot_world_v2/test_service.py new file mode 100644 index 0000000..280c34c --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_v2/test_service.py @@ -0,0 +1,81 @@ +from unittest.mock import MagicMock, patch + +import torch +from PIL import Image +from click.testing import CliRunner + +from examples.lingbot import lingbot_world_v2_image_to_video_h100 as offline_example +from examples.lingbot import stream_lingbot_world_v2 as stream_example +from telefuser.core.config import AttnImplType +from telefuser.pipelines.lingbot_world_fast.session import LingBotWorldFastSessionConfig + + +def test_v2_stream_get_pipeline_maps_ppl_config_to_internal_workers() -> None: + pipeline = MagicMock() + + with patch.object(stream_example, "LingBotWorldV2Pipeline", return_value=pipeline) as pipeline_cls: + result = stream_example.get_pipeline( + parallelism=4, + model_root="/models/Wan2.2-I2V-A14B", + v2_model_root="/models/lingbot-world-v2-14b-causal-fast", + ) + + assert result is pipeline + pipeline_cls.assert_called_once_with(device="cuda", torch_dtype=torch.bfloat16) + config = pipeline.init.call_args.args[0] + assert config.checkpoint_dir == "/models/Wan2.2-I2V-A14B" + assert config.fast_checkpoint_path == "/models/lingbot-world-v2-14b-causal-fast" + assert config.local_attn_size == 18 + assert config.sink_size == 6 + assert config.timestep_indices == (0, 250, 500, 750) + assert config.attention_config.attn_impl == AttnImplType.TORCH_SDPA + assert config.parallel_config.device_ids == [0, 1, 2, 3] + + +def test_v2_stream_service_constructs_v2_session_from_ppl_config() -> None: + pipeline = MagicMock() + + with patch.object(stream_example, "get_pipeline", return_value=pipeline) as get_pipeline: + service = stream_example.get_service(gpu_num=4) + + get_pipeline.assert_called_once_with(parallelism=4) + session_id = service.create_session({"image": Image.new("RGB", (8, 8))}) + session_config = pipeline.control_context.call_args.args[0] + + assert isinstance(session_config, LingBotWorldFastSessionConfig) + assert session_config.frame_num == 1917 + assert session_config.chunk_size == 4 + assert session_config.frame_policy == "truncate" + assert session_config.sample_shift == 10.0 + assert service.default_fps == 16 + assert service.max_generation_seconds == 120.0 + assert session_id in service._sessions + + +def test_v2_offline_cli_forwards_only_supported_run_arguments(tmp_path) -> None: + image_path = tmp_path / "input.png" + Image.new("RGB", (8, 8)).save(image_path) + output_path = tmp_path / "output.mp4" + pipeline = MagicMock() + + with ( + patch.object(offline_example, "get_pipeline", return_value=pipeline), + patch.object(offline_example, "run", return_value=[]) as run, + patch.object(offline_example, "save_video"), + ): + result = CliRunner().invoke( + offline_example.main, + [ + "--image_path", + str(image_path), + "--action_path", + str(tmp_path), + "--output", + str(output_path), + ], + ) + + assert result.exit_code == 0, result.output + assert "frame_num" not in run.call_args.kwargs + assert "max_sequence_length" not in run.call_args.kwargs + pipeline.close.assert_called_once_with() diff --git a/tests/unit/service/test_service_routes.py b/tests/unit/service/test_service_routes.py index 58b523a..cf46f8e 100644 --- a/tests/unit/service/test_service_routes.py +++ b/tests/unit/service/test_service_routes.py @@ -319,6 +319,75 @@ async def create_bidirectional_session(self, **kwargs: object) -> tuple[str, str assert webrtc_manager.fps == 16 +def test_bidirectional_webrtc_offer_uses_service_default_fps_when_omitted() -> None: + class RunningStreamService: + def __init__(self) -> None: + self.config: dict | None = None + self.service = types.SimpleNamespace(default_fps=16) + + def create_session(self, config: dict) -> str: + self.config = config + return "pipeline-session" + + async def pull_chunks(self, session_id: str): + if False: + yield session_id + + def push_chunk(self, session_id: str, chunk: dict) -> None: + pass + + def close_session(self, session_id: str) -> None: + pass + + class WebRTCSessionManager: + def __init__(self) -> None: + self.fps: int | None = None + + async def create_bidirectional_session(self, **kwargs: object) -> tuple[str, str]: + self.fps = int(kwargs["fps"]) + return "answer-sdp", "answer" + + stream_service = RunningStreamService() + webrtc_manager = WebRTCSessionManager() + routes = WebRTCRoutes.__new__(WebRTCRoutes) + routes.api = Mock(stream_service=stream_service) + routes._session_manager = webrtc_manager + request = WebRTCOfferRequest(session_id="request-session", sdp="offer-sdp", task="i2v") + + asyncio.run(routes._handle_bidirectional(stream_service, request)) + + assert stream_service.config == {"session_id": "request-session", "task": "i2v"} + assert webrtc_manager.fps == 16 + + +def test_server_push_webrtc_offer_uses_service_default_fps_when_omitted() -> None: + class RunningStreamService: + service = types.SimpleNamespace(default_fps=16) + + async def stream_task(self, task_data: dict): + if False: + yield task_data + + class WebRTCSessionManager: + def __init__(self) -> None: + self.fps: int | None = None + + async def create_session(self, **kwargs: object) -> tuple[str, str]: + self.fps = int(kwargs["fps"]) + return "answer-sdp", "answer" + + stream_service = RunningStreamService() + webrtc_manager = WebRTCSessionManager() + routes = WebRTCRoutes.__new__(WebRTCRoutes) + routes.api = Mock(stream_service=stream_service) + routes._session_manager = webrtc_manager + request = WebRTCOfferRequest(session_id="request-session", sdp="offer-sdp", task="i2v") + + asyncio.run(routes._handle_server_push(stream_service, request)) + + assert webrtc_manager.fps == 16 + + def test_bidirectional_webrtc_offer_prefers_explicit_top_level_config() -> None: class RunningStreamService: def __init__(self) -> None: