diff --git a/.gitignore b/.gitignore index b6dfd4c..194bb09 100755 --- a/.gitignore +++ b/.gitignore @@ -151,3 +151,7 @@ INSTALL_HYS.md AGENTS.md _version.py.mcp.json telefuser/_version.py +# LingBot regression example assets +!examples/data/lingbot_world_fast/image.jpg +!examples/data/lingbot_world_fast/poses.npy +!examples/data/lingbot_world_fast/intrinsics.npy diff --git a/README.md b/README.md index b6186f4..9effb0a 100644 --- a/README.md +++ b/README.md @@ -82,19 +82,46 @@ video = pipe( TeleFuser includes a bidirectional WebRTC demo for `LingBot-World-Fast`. ```bash -export TF_MODEL_ZOO_PATH=/path/to/model_zoo -# Expected subdirectories: Wan2.2-I2V-A14B and lingbot-world-fast - +TF_MODEL_ZOO_PATH=/path/to/model_zoo \ +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 \ - -p 8088 \ - --skip-validation + --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation python examples/stream_server/webrtc_bidirectional_demo.py \ - --server-url http://localhost:8088 \ - --image-path /path/to/input.png + --server-url http://127.0.0.1:8088 \ + --port 8091 \ + --image-path examples/data/lingbot_world_fast/image.jpg \ + --action-path examples/data/lingbot_world_fast \ + --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 +``` + +This starts a continuous session where the client sends control messages over a WebRTC DataChannel and receives +generated video frames over media tracks. When the browser runs on a laptop through VS Code Remote SSH, configure +TURN over TCP and forward ports `8091` and `3478`; port `8088` does not need forwarding because the demo proxies +signaling requests. Keep local port `3478` equal to remote port `3478`; the forwarded 8091 port may use any available +local port. Without VS Code, run the equivalent tunnel from a terminal on the laptop: + +```bash +ssh -N -o ExitOnForwardFailure=yes -o ServerAliveInterval=30 \ + -L 8091:127.0.0.1:8091 \ + -L 3478:127.0.0.1:3478 \ + USER@SERVER_HOST ``` -This starts a continuous session where the client sends control messages over a WebRTC DataChannel and receives generated video frames over media tracks. +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 +[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 +`TELEFUSER_TURN_*` variables, start the service on `127.0.0.1:8088`, run the demo without any `--turn-*` or +`--force-turn-relay` arguments, and open `http://localhost:8091`. This does not apply when only the shell is on the +server through SSH but the browser still runs on a laptop. ### 3. Batch Service Mode diff --git a/docs/en/service.md b/docs/en/service.md index fe2a4bb..fd7daf9 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -45,8 +45,8 @@ telefuser serve \ --parallelism 1 # For real-time world model streaming (WebRTC support is included in the default install) -export TF_MODEL_ZOO_PATH=/path/to/model_zoo -# Expected subdirectories: Wan2.2-I2V-A14B and lingbot-world-fast +# Set TF_MODEL_ZOO_PATH and PPL_CONFIG["parallelism"] in +# examples/lingbot/stream_lingbot_world_fast.py telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py -p 8088 --skip-validation ``` diff --git a/docs/en/stream_server.md b/docs/en/stream_server.md index bcc1393..08bd3db 100644 --- a/docs/en/stream_server.md +++ b/docs/en/stream_server.md @@ -109,6 +109,7 @@ telefuser stream-serve [OPTIONS] |------|---------|-------------| | `--port`, `-p` | `8088` | Server port | | `--host` | `0.0.0.0` | Bind address | +| `--gpu-num`, `-g` | `1` | GPU count passed to `get_service(gpu_num=...)` when supported | | `--security-level` | `strict` | Pipeline validation level (`none`, `basic`, `strict`, `sandbox`). `sandbox` is a best-effort restricted-load check, not runtime isolation. | | `--skip-validation` | `false` | Skip pipeline file security checks | @@ -129,7 +130,8 @@ telefuser stream-serve my_pipeline.py -p 8088 --skip-validation ## Creating a Stream Pipeline -A stream pipeline is a Python file that defines a `get_service()` function returning a service object. The service must implement one of two protocols. +A stream pipeline is a Python file that defines a `get_service()` function returning a service object. It may accept +an optional `gpu_num` parameter supplied by `stream-serve --gpu-num`. The service must implement one of two protocols. ### ServerPushService Protocol @@ -557,23 +559,61 @@ LingBot-World-Fast requires two sets of weights: | Setting | Example | Description | |---------|---------|-------------| -| `TF_MODEL_ZOO_PATH` | `/storage/model_zoo` | Root model directory read by the example script | +| `TF_MODEL_ZOO_PATH` | `/storage/model_zoo` | Required environment variable that locates both model trees | | Base model subdirectory | `/storage/model_zoo/Wan2.2-I2V-A14B` | Base Wan2.2 I2V weights containing VAE, T5 text encoder, and tokenizer | -| Fast model subdirectory | `/storage/model_zoo/lingbot-world-fast` | LingBot-World-Fast DiT weights | +| Fast model subdirectory | `/storage/model_zoo/lingbot/lingbot-world-fast` | LingBot-World-Fast DiT weights | + +#### VS Code Remote SSH: TURN over TCP + +When the page is opened by a laptop browser through VS Code Remote SSH, the browser and TeleFuser are not on the +same network interface even though both URLs use `localhost`. SDP signaling can use the demo's HTTP proxy, but RTP +media still needs an ICE route. VS Code forwards TCP rather than arbitrary UDP relay ports, so use TURN over TCP. + +For development on a trusted host, install coturn and run a loopback-only server in a separate terminal: + +```bash +sudo apt-get install -y coturn + +turnserver -n -m 1 \ + --listening-ip=127.0.0.1 \ + --relay-ip=127.0.0.1 \ + --listening-port=3478 \ + --min-port=49160 --max-port=49200 \ + --user=telefuser:telefuser-turn \ + --realm=telefuser.local \ + --fingerprint --lt-cred-mech \ + --no-tls --no-dtls --no-cli \ + --allow-loopback-peers \ + --simple-log --log-file=/tmp/telefuser-turn.log +``` + +`--allow-loopback-peers` is needed only because coturn and TeleFuser are peers on the same host. This command is for +local development and must not be exposed directly to the internet. Verify the TCP allocation and credentials: + +```bash +turnutils_uclient -t -y -c \ + -u telefuser -w telefuser-turn -p 3478 127.0.0.1 +``` #### Start the Server -Recommended launch with two GPUs: DiT on GPU 0, text encoder and VAE on GPU 1. If GPU 1 VRAM is tight, move VAE back to CPU. +Set `TF_MODEL_ZOO_PATH` before starting the service and pass the worker count with `--gpu-num`. The service creates +Ulysses workers internally, so the server command does not need `torchrun` or distributed environment variables. +Use `CUDA_VISIBLE_DEVICES` to select the physical GPUs. ```bash -TELEFUSER_TURN_SERVER='turn:127.0.0.1:3478' \ -TELEFUSER_TURN_USERNAME=telefuser \ -TELEFUSER_TURN_CREDENTIAL=your-turn-password \ TF_MODEL_ZOO_PATH=/storage/model_zoo \ +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 \ - -p 8088 --host 0.0.0.0 --skip-validation + --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation ``` +The streaming example enables FSDP and Ulysses sequence parallelism. All four worker logs should show a device mesh +with `dims=[4]` and `Enabling FSDP for lingbot_world_fast_denoise`. + Wait for the following log line before connecting the browser demo: ```text @@ -588,23 +628,24 @@ curl --noproxy '*' http://127.0.0.1:8088/v1/service/health #### Start the Browser Demo -When using VS Code Remote SSH to access the server from a laptop browser, browser JavaScript runs locally. In this case, use TURN and forward the remote `3478` port to local `3478`. +When using VS Code Remote SSH, run the demo on the remote TeleFuser host. The demo proxies signaling and session +requests to 8088 by default. ```bash python examples/stream_server/webrtc_bidirectional_demo.py \ - --server-url http://localhost:8088 \ + --server-url http://127.0.0.1:8088 \ --port 8091 \ - --image-path /tmp/lingbot_test_input.png \ - --frame-num 81 \ + --image-path examples/data/lingbot_world_fast/image.jpg \ + --action-path examples/data/lingbot_world_fast \ + --frame-num 321 \ --chunk-size 3 \ + --sample-shift 10.0 \ --fps 16 \ --turn-url 'turn:localhost:3478?transport=tcp' \ --turn-username telefuser \ - --turn-credential your-turn-password \ + --turn-credential telefuser-turn \ --force-turn-relay \ --ice-gather-timeout-ms 30000 \ - --control-lateral-step 0.25 \ - --control-yaw-step-degrees 12 \ --no-open ``` @@ -614,59 +655,112 @@ Open in browser: http://localhost:8091 ``` -`--image-path` is a server-side file path, not the laptop local path. The demo enables proxying by default; the browser only needs to access the demo port. Requests to `/v1/stream/webrtc/*` are forwarded by the demo process to `--server-url`. +In the VS Code **Ports** panel, forward the following remote TCP ports: -#### Without TURN: View on Server Browser +| Remote port | Local port | Purpose | +|-------------|------------|---------| +| `8091` | Any available port | Demo HTML and proxied `/v1/stream/webrtc/*` requests | +| `3478` | `3478` | TURN-over-TCP connection referenced by `turn:localhost:3478` | -If the browser actually runs on the server side (e.g., Chrome in remote desktop, VNC, or noVNC), you can skip TURN. Do not set `TELEFUSER_TURN_*` and do not pass `--turn-url` to the demo. +Port 8088 does not need forwarding while the demo proxy is enabled. Port 3478 must retain local port 3478 unless +the demo's `--turn-url` is changed to match another local port. Open the exact 8091 URL displayed by VS Code using +`http://`; `localhost` is a browser secure context, so HTTPS is not needed for this development workflow. + +Do not forward the coturn relay range (`49160-49200`) through VS Code. Browser relay traffic stays inside the +forwarded TURN TCP connection; the relay range is used on the remote host between coturn and the WebRTC peer. +Internet-facing production TURN is different and must expose its configured relay range through the firewall. + +If you use a normal SSH client instead of the VS Code Ports panel, start the tunnel in a separate laptop terminal. +Replace `USER` and `SERVER_HOST` with the server login: + +```bash +ssh -N \ + -o ExitOnForwardFailure=yes \ + -o ServerAliveInterval=30 \ + -L 8091:127.0.0.1:8091 \ + -L 3478:127.0.0.1:3478 \ + USER@SERVER_HOST +``` + +Keep the tunnel running and open `http://localhost:8091`. Add `-p SSH_PORT` for a non-default SSH port or +`-i /path/to/private_key` for a specific identity file. Omit `-N` if the same connection should also open an +interactive shell. + +If local port 8091 is occupied, use `-L 18091:127.0.0.1:8091` and open `http://localhost:18091`. If local port 3478 +is occupied, use `-L 13478:127.0.0.1:3478` and pass +`--turn-url 'turn:localhost:13478?transport=tcp'` to the browser demo. Do not change the server-side +`TELEFUSER_TURN_SERVER`, which still connects to coturn at remote port 3478. + +`--image-path` and `--action-path` are server-side paths, not laptop-local paths. For real-time keyboard control, +the service loads only `intrinsics.npy` from `--action-path` and keeps its first row fixed for the session. The demo +enables proxying by default; the browser only needs to access the demo port. Requests to `/v1/stream/webrtc/*` are +forwarded by the demo process to `--server-url`. + +#### Run the Browser and Service on the Same Machine + +If the browser and GPU service run on the same physical machine, you can skip coturn and SSH forwarding. Examples +include Chrome opened directly on the workstation or through remote desktop, VNC, or noVNC. An SSH login alone does +not make the laptop browser local: if the browser still runs on the laptop, use the TURN and port-forwarding setup +above. + +Do not set `TELEFUSER_TURN_*`, and do not pass `--turn-url`, TURN credentials, or `--force-turn-relay` to the demo. Server: ```bash env -u TELEFUSER_TURN_SERVER \ --u TELEFUSER_TURN_USERNAME \ --u TELEFUSER_TURN_CREDENTIAL \ -TF_MODEL_ZOO_PATH=/storage/model_zoo \ -telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ - -p 8088 --host 0.0.0.0 --skip-validation + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + TF_MODEL_ZOO_PATH=/path/to/model_zoo \ + CUDA_VISIBLE_DEVICES=0,1,2,3 \ + telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ + --gpu-num 4 -p 8088 --host 127.0.0.1 --skip-validation ``` Demo: ```bash env -u TELEFUSER_TURN_SERVER \ --u TELEFUSER_TURN_USERNAME \ --u TELEFUSER_TURN_CREDENTIAL \ -python examples/stream_server/webrtc_bidirectional_demo.py \ + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + python examples/stream_server/webrtc_bidirectional_demo.py \ --server-url http://127.0.0.1:8088 \ --port 8091 \ - --image-path /tmp/lingbot_test_input.png \ - --frame-num 81 \ + --image-path examples/data/lingbot_world_fast/image.jpg \ + --action-path examples/data/lingbot_world_fast \ + --frame-num 321 \ --chunk-size 3 \ + --sample-shift 10.0 \ --fps 16 \ - --control-lateral-step 0.25 \ - --control-yaw-step-degrees 12 \ --no-open ``` -Open on server browser: +Open in the browser on the same machine: ```text -http://127.0.0.1:8091 +http://localhost:8091 ``` #### Direction Control -The demo supports keyboard arrow keys and a page D-pad: +The demo supports the page D-pad, arrow keys, and the source-compatible `WASD`/`IJKL` keyboard controls: | Input | cam mode meaning | -|-------|-----------------| -| `↑` | Camera forward | -| `↓` | Camera backward | -| `←` | Turn left and strafe left | -| `→` | Turn right and strafe right | - -Controls only affect chunks that have not yet started generating; chunks already denoising or decoding are not immediately changed. Therefore, it is recommended to hold a direction key early after connecting, rather than clicking near the end of the video. +|-------|------------------| +| `↑` / `W` | Move forward | +| `↓` / `S` | Move backward | +| `←` / `A` | Strafe left | +| `→` / `D` | Strafe right | +| `J` / `L` | Yaw left / right | +| `I` / `K` | Pitch up / down | + +Controls only affect chunks that have not yet started generating; chunks already denoising or decoding are not +immediately changed. Holding a key continuously generates controlled chunks. After all keys are released, the +service stops requesting new chunks and WebRTC repeats the most recently emitted frame. + +The camera pose starts at identity. Each latent interval integrates four video-frame control steps, matching the +LingBot source trajectory generator. Camera intrinsics are fixed when the session is created. Absolute pose and +pitch state are accumulated across chunks, including the relative-pose delta across each chunk boundary. The following states in the DataChannel log indicate that direction control has been consumed by the server and applied to a generation chunk: @@ -681,13 +775,13 @@ Common control strength parameters: | Parameter | Default | Description | |-----------|---------|-------------| -| `--control-move-step` | `0.18` | Forward/backward displacement step size | -| `--control-yaw-step-degrees` | `10.0` | Yaw angle per latent frame | -| `--control-lateral-step` | `0.12` | Left/right lateral step size | +| `--control-move-step` | `0.05` | Forward/backward displacement per video frame | +| `--control-yaw-step-degrees` | `2.0` | Yaw angle per video frame | +| `--control-lateral-step` | `0.05` | Lateral displacement per video frame | +| `--control-pitch-step-degrees` | `2.0` | Pitch angle per video frame | +| `--control-pitch-limit-degrees` | `85.0` | Absolute pitch limit | | `--show-control-hud / --no-show-control-hud` | `true` | Whether to overlay direction HUD on controlled chunks | -If left/right effects are subtle, increase `--control-lateral-step`, e.g., `0.25` or `0.3`. - #### `cam` vs `act` Control Modes | Mode | Input | Description | @@ -707,28 +801,56 @@ Only switch to `act` when using action-control weights: --control-mode act ``` -The service example currently sets `control_type="cam"` inside -`examples/lingbot/stream_lingbot_world_fast.py`. To use action-control -weights, edit `PPL_CONFIG["control_type"]` in that script as well as passing -`--control-mode act` to the demo. +The service example sets `PPL_CONFIG["control_mode"]="cam"` inside +`examples/lingbot/stream_lingbot_world_fast.py`. A session's control mode must +match the mode used to initialize the pipeline. To use action-control weights, +set that value to `"act"` before starting the service and pass `--control-mode +act` to the demo. + +#### Frame and Control Contract + +LingBot accepts complete latent chunks only. `frame_num` must be `4n + 1`, and +the resulting latent-frame count must divide evenly by `chunk_size`. With the +default chunk size of 3, valid output lengths are 9, 21, 33, ..., 321 frames. +The browser accepts a duration and rounds down to the largest complete chunk: +10 seconds becomes 153 frames and 20 seconds becomes 321 frames at 16 FPS. +Direct API requests with invalid values are rejected. Offline control files +must provide enough poses and action samples for the requested video window. +Intrinsics may be either one static `(4,)` value or one `(frames, 4)` value per +video frame. Explicit control tensors must match the pipeline's selected mode, +device, dtype, and full latent-chunk shape. #### VRAM and Resolution -LingBot's KV cache grows with `frame_num` and output resolution. 832×480 at 81 frames approaches the 80 GB H100 VRAM limit. Start with a lower resolution to verify the pipeline: +LingBot's global KV cache grows with `frame_num` and output resolution. FSDP shards model parameters, while Ulysses +shards KV heads across workers. More GPUs therefore allow longer sessions, but the scaling is not perfectly linear. +The following 832x480 configurations were tested with FSDP, `chunk_size=3`, and `sample_shift=10.0`: + +| GPUs | Duration | Frames | Result | +|------|----------|--------|--------| +| 2 H100 80 GB | 10 seconds | 153 | Passed | +| 2 H100 80 GB | 20 seconds | 321 | KV-cache CUDA OOM | +| 4 H100 80 GB | 20 seconds | 321 | Passed, 27/27 chunks | + +The four-GPU run peaked at approximately 58.6 GiB on GPU 0 and 41.6 GiB on GPUs 1-3. Start with five seconds to +verify the pipeline before selecting the maximum duration: ```bash -# In examples/lingbot/stream_lingbot_world_fast.py: -PPL_CONFIG["max_area"] = 99840 # approximately 416×240 +# In the browser demo or session configuration: +--frame-num 81 ``` Common tuning parameters: | Parameter | Effect | |-----------|--------| -| `PPL_CONFIG["max_area"]` | Reduce output area, significantly lower VRAM usage | +| `PPL_CONFIG["resolution"]` | Select either the 480p or 720p output area | | `--frame-num` | Reduce total generated frames and latent chunk count | | `--chunk-size` | Affect the latent chunk size per generation step | +An OOM marks the parallel worker as failed. Closing the browser session is not enough; restart the stream service +before trying another duration. + The service currently allows only one LingBot active session. Before reconnecting, click **Stop** on the demo, or call: ```bash diff --git a/docs/zh/service.md b/docs/zh/service.md index c4adbcd..51de42c 100644 --- a/docs/zh/service.md +++ b/docs/zh/service.md @@ -45,8 +45,8 @@ telefuser serve \ --parallelism 1 # 实时世界模型流式推理(默认安装已包含 WebRTC 支持) -export TF_MODEL_ZOO_PATH=/path/to/model_zoo -# 预期子目录:Wan2.2-I2V-A14B 和 lingbot-world-fast +# 在 examples/lingbot/stream_lingbot_world_fast.py 中设置 +# TF_MODEL_ZOO_PATH 和 PPL_CONFIG["parallelism"] telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py -p 8088 --skip-validation ``` diff --git a/docs/zh/stream_server.md b/docs/zh/stream_server.md index 3469b50..5c1b5ba 100644 --- a/docs/zh/stream_server.md +++ b/docs/zh/stream_server.md @@ -109,6 +109,7 @@ telefuser stream-serve [选项] |------|--------|------| | `--port`, `-p` | `8088` | 服务端口 | | `--host` | `0.0.0.0` | 绑定地址 | +| `--gpu-num`, `-g` | `1` | 管线支持时传给 `get_service(gpu_num=...)` 的 GPU 数量 | | `--security-level` | `strict` | 管线验证级别(`none`、`basic`、`strict`、`sandbox`)。`sandbox` 是 best-effort 受限加载检查,不是运行时隔离。 | | `--skip-validation` | `false` | 跳过管线文件安全检查 | @@ -129,7 +130,8 @@ telefuser stream-serve my_pipeline.py -p 8088 --skip-validation ## 创建流式管线 -流式管线是一个 Python 文件,定义一个 `get_service()` 函数,返回一个服务对象。该服务必须实现以下两种协议之一。 +流式管线是一个 Python 文件,定义一个返回服务对象的 `get_service()` 函数。该函数可以接收由 +`stream-serve --gpu-num` 提供的可选 `gpu_num` 参数。服务必须实现以下两种协议之一。 ### ServerPushService 协议 @@ -557,23 +559,60 @@ LingBot-World-Fast 需要两类权重: | 配置 | 示例 | 说明 | |----------|------|------| -| `TF_MODEL_ZOO_PATH` | `/storage/model_zoo` | 示例脚本读取的模型根目录 | +| `TF_MODEL_ZOO_PATH` | `/storage/model_zoo` | 必填环境变量,用于定位两套模型目录 | | 基础模型子目录 | `/storage/model_zoo/Wan2.2-I2V-A14B` | 基础 Wan2.2 I2V 权重,包含 VAE、T5 文本编码器和 tokenizer | -| Fast 模型子目录 | `/storage/model_zoo/lingbot-world-fast` | LingBot-World-Fast DiT 权重 | +| Fast 模型子目录 | `/storage/model_zoo/lingbot/lingbot-world-fast` | LingBot-World-Fast DiT 权重 | + +#### VS Code Remote SSH:使用 TCP TURN + +通过 VS Code Remote SSH 在笔记本浏览器中打开远端页面时,浏览器和 TeleFuser 实际不在同一个网络接口。 +SDP 信令可以由 demo 的 HTTP 代理转发,但 RTP 媒体仍需要可用的 ICE 路径。VS Code 转发的是 TCP 端口, +无法直接转发任意 UDP relay 端口,因此该场景应使用 TCP TURN。 + +在可信服务器上进行开发测试时,可以安装 coturn,并在单独终端中启动仅监听 loopback 的临时服务: + +```bash +sudo apt-get install -y coturn + +turnserver -n -m 1 \ + --listening-ip=127.0.0.1 \ + --relay-ip=127.0.0.1 \ + --listening-port=3478 \ + --min-port=49160 --max-port=49200 \ + --user=telefuser:telefuser-turn \ + --realm=telefuser.local \ + --fingerprint --lt-cred-mech \ + --no-tls --no-dtls --no-cli \ + --allow-loopback-peers \ + --simple-log --log-file=/tmp/telefuser-turn.log +``` + +由于 coturn 和 TeleFuser 位于同一台服务器,测试配置需要 `--allow-loopback-peers`。该命令仅用于本地开发, +不要直接暴露到公网。可以使用以下命令验证 TCP allocation 和账号密码: + +```bash +turnutils_uclient -t -y -c \ + -u telefuser -w telefuser-turn -p 3478 127.0.0.1 +``` #### 启动服务端 -两张 GPU 的推荐启动方式如下。DiT 放在 GPU0,文本编码器和 VAE 放在 GPU1;如果 GPU1 显存紧张,可以把 VAE 改回 CPU。 +启动前设置 `TF_MODEL_ZOO_PATH`,并通过 `--gpu-num` 传入 worker 数量。TeleFuser 会在服务内部创建 +Ulysses worker,因此启动命令不需要 `torchrun` 或分布式环境变量。使用 `CUDA_VISIBLE_DEVICES` 选择物理 GPU。 ```bash -TELEFUSER_TURN_SERVER='turn:127.0.0.1:3478' \ -TELEFUSER_TURN_USERNAME=telefuser \ -TELEFUSER_TURN_CREDENTIAL=your-turn-password \ TF_MODEL_ZOO_PATH=/storage/model_zoo \ +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 \ - -p 8088 --host 0.0.0.0 --skip-validation + --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation ``` +streaming example 默认同时启用 FSDP 和 Ulysses sequence parallel。四个 worker 的日志都应出现 +`dims=[4]` 的 DeviceMesh,以及 `Enabling FSDP for lingbot_world_fast_denoise`。 + 等待日志出现以下内容后再连接浏览器 demo: ```text @@ -588,23 +627,24 @@ curl --noproxy '*' http://127.0.0.1:8088/v1/service/health #### 启动浏览器 demo -通过 VS Code Remote SSH 使用笔记本浏览器访问远端 demo 时,浏览器 JavaScript 实际运行在笔记本本地。此时建议使用 TURN,并将远端 `3478` 端口转发到本地 `3478`。 +通过 VS Code Remote SSH 使用笔记本浏览器时,在远端 TeleFuser 主机上运行 demo。demo 默认会代理到 +8088 的信令和 session 请求。 ```bash python examples/stream_server/webrtc_bidirectional_demo.py \ - --server-url http://localhost:8088 \ + --server-url http://127.0.0.1:8088 \ --port 8091 \ - --image-path /tmp/lingbot_test_input.png \ - --frame-num 81 \ + --image-path examples/data/lingbot_world_fast/image.jpg \ + --action-path examples/data/lingbot_world_fast \ + --frame-num 321 \ --chunk-size 3 \ + --sample-shift 10.0 \ --fps 16 \ --turn-url 'turn:localhost:3478?transport=tcp' \ --turn-username telefuser \ - --turn-credential your-turn-password \ + --turn-credential telefuser-turn \ --force-turn-relay \ --ice-gather-timeout-ms 30000 \ - --control-lateral-step 0.25 \ - --control-yaw-step-degrees 12 \ --no-open ``` @@ -614,59 +654,108 @@ python examples/stream_server/webrtc_bidirectional_demo.py \ http://localhost:8091 ``` -`--image-path` 是服务端文件路径,不是笔记本本地路径。demo 默认开启代理,浏览器只需要访问 demo 端口;`/v1/stream/webrtc/*` 请求会由 demo 进程转发到 `--server-url`。 +在 VS Code 的 **端口(Ports)** 面板中转发以下远端 TCP 端口: -#### 不使用 TURN:在服务器浏览器查看 +| 远端端口 | 本地端口 | 用途 | +|----------|----------|------| +| `8091` | 任意可用端口 | demo HTML,以及代理后的 `/v1/stream/webrtc/*` 请求 | +| `3478` | `3478` | `turn:localhost:3478` 使用的 TCP TURN 连接 | -如果浏览器真的运行在服务器侧,例如远程桌面、VNC 或 noVNC 中的 Chrome,可以不使用 TURN。此时不要设置 `TELEFUSER_TURN_*`,demo 也不要传 `--turn-url`。 +demo 代理开启时不需要转发 8088。3478 必须保持本地端口也是 3478,除非同步修改 demo 的 +`--turn-url`。请用 `http://` 打开 VS Code 显示的准确 8091 转发地址;开发场景中的 `localhost` +属于浏览器 secure context,不要求额外配置 HTTPS。 + +不要通过 VS Code 转发 coturn 的 relay 范围(`49160-49200`)。浏览器 relay 流量封装在已转发的 TURN +TCP 连接中;relay 范围只在远端主机上的 coturn 与 WebRTC peer 之间使用。公网生产 TURN 与此不同, +必须在防火墙中开放所配置的 relay 端口范围。 + +如果不使用 VS Code 的端口面板,而是通过普通 SSH 客户端登录服务器,请在笔记本的另一个终端中建立 +隧道。将 `USER` 和 `SERVER_HOST` 替换为实际的服务器登录信息: + +```bash +ssh -N \ + -o ExitOnForwardFailure=yes \ + -o ServerAliveInterval=30 \ + -L 8091:127.0.0.1:8091 \ + -L 3478:127.0.0.1:3478 \ + USER@SERVER_HOST +``` + +保持隧道命令运行,然后打开 `http://localhost:8091`。SSH 使用非默认端口时添加 `-p SSH_PORT`,指定 +私钥时添加 `-i /path/to/private_key`。如果希望同一连接同时进入交互式 shell,可以去掉 `-N`。 + +如果笔记本的 8091 已被占用,可以改用 `-L 18091:127.0.0.1:8091`,并打开 +`http://localhost:18091`。如果本地 3478 已被占用,可以使用 `-L 13478:127.0.0.1:3478`,同时把 +浏览器 demo 参数改为 `--turn-url 'turn:localhost:13478?transport=tcp'`。服务端的 +`TELEFUSER_TURN_SERVER` 不需要修改,它仍然连接远端服务器上的 coturn 3478 端口。 + +`--image-path` 和 `--action-path` 都是服务端路径,不是笔记本本地路径。实时键盘控制只会从 +`--action-path` 加载 `intrinsics.npy`,并在整个会话中固定使用第一行内参。demo 默认开启代理,浏览器只需 +访问 demo 端口;`/v1/stream/webrtc/*` 请求会由 demo 进程转发到 `--server-url`。 + +#### 浏览器和服务运行在同一台机器 + +如果浏览器和 GPU 服务运行在同一台物理机器上,可以不使用 coturn,也不需要 SSH 端口转发。例如直接在 +工作站上打开 Chrome,或者使用远程桌面、VNC、noVNC 中的服务器侧浏览器。仅通过 SSH 登录服务器并不代表 +笔记本浏览器也在服务器本机;如果浏览器仍运行在笔记本上,必须使用前面的 TURN 和端口转发方案。 + +该模式不要设置 `TELEFUSER_TURN_*`,demo 也不要传入 `--turn-url`、TURN 账号密码或 +`--force-turn-relay`。 服务端: ```bash env -u TELEFUSER_TURN_SERVER \ --u TELEFUSER_TURN_USERNAME \ --u TELEFUSER_TURN_CREDENTIAL \ -TF_MODEL_ZOO_PATH=/storage/model_zoo \ -telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ - -p 8088 --host 0.0.0.0 --skip-validation + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + TF_MODEL_ZOO_PATH=/path/to/model_zoo \ + CUDA_VISIBLE_DEVICES=0,1,2,3 \ + telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ + --gpu-num 4 -p 8088 --host 127.0.0.1 --skip-validation ``` demo: ```bash env -u TELEFUSER_TURN_SERVER \ --u TELEFUSER_TURN_USERNAME \ --u TELEFUSER_TURN_CREDENTIAL \ -python examples/stream_server/webrtc_bidirectional_demo.py \ + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + python examples/stream_server/webrtc_bidirectional_demo.py \ --server-url http://127.0.0.1:8088 \ --port 8091 \ - --image-path /tmp/lingbot_test_input.png \ - --frame-num 81 \ + --image-path examples/data/lingbot_world_fast/image.jpg \ + --action-path examples/data/lingbot_world_fast \ + --frame-num 321 \ --chunk-size 3 \ + --sample-shift 10.0 \ --fps 16 \ - --control-lateral-step 0.25 \ - --control-yaw-step-degrees 12 \ --no-open ``` -在服务器浏览器打开: +在同一台机器的浏览器中打开: ```text -http://127.0.0.1:8091 +http://localhost:8091 ``` #### 方向控制 -demo 支持键盘方向键和页面 D-pad: +demo 支持页面 D-pad、方向键以及与 LingBot 源码一致的 `WASD`/`IJKL` 键盘控制: | 输入 | cam 模式含义 | |------|--------------| -| `↑` | 相机前进 | -| `↓` | 相机后退 | -| `←` | 左转并向左横移 | -| `→` | 右转并向右横移 | +| `↑` / `W` | 向前移动 | +| `↓` / `S` | 向后移动 | +| `←` / `A` | 向左横移 | +| `→` / `D` | 向右横移 | +| `J` / `L` | 向左/向右偏航 | +| `I` / `K` | 向上/向下俯仰 | + +控制只会作用于“尚未开始生成”的后续 chunk;已经在 denoising 或 decoding 的 chunk 不会被即时改变。 +长按按键会持续生成受控 chunk;松开所有按键后,服务端停止请求新 chunk,WebRTC 会重复显示最近输出的一帧。 -控制只会作用于“尚未开始生成”的后续 chunk;已经在 denoising 或 decoding 的 chunk 不会被即时改变。因此建议在连接成功后尽早长按方向键,而不是等视频快结束时再点击。 +相机位姿从单位矩阵开始。每个 latent 间隔累计 4 次视频帧控制步,与 LingBot 源码中的轨迹生成方式一致。 +相机内参在创建会话时固定,绝对位姿与 pitch 状态跨 chunk 累计,并保留 chunk 边界两侧的相对位姿变化。 DataChannel 日志中出现以下状态时,表示方向控制已被服务端消费并应用到某个生成 chunk: @@ -681,13 +770,13 @@ demo 默认开启 `Control HUD`。该 HUD 会叠加在使用了方向控制的 | 参数 | 默认值 | 说明 | |------|--------|------| -| `--control-move-step` | `0.18` | 前进/后退位移步长 | -| `--control-yaw-step-degrees` | `10.0` | 每个 latent frame 的转向角度 | -| `--control-lateral-step` | `0.12` | 左右横向平移步长 | +| `--control-move-step` | `0.05` | 每个视频帧的前进/后退位移 | +| `--control-yaw-step-degrees` | `2.0` | 每个视频帧的偏航角度 | +| `--control-lateral-step` | `0.05` | 每个视频帧的横向位移 | +| `--control-pitch-step-degrees` | `2.0` | 每个视频帧的俯仰角度 | +| `--control-pitch-limit-degrees` | `85.0` | 绝对俯仰角限制 | | `--show-control-hud / --no-show-control-hud` | `true` | 是否在受控 chunk 上叠加方向 HUD | -如果左右效果不明显,可以增大 `--control-lateral-step`,例如 `0.25` 或 `0.3`。 - #### `cam` 与 `act` 控制模式 | 模式 | 输入 | 说明 | @@ -707,27 +796,52 @@ demo 默认开启 `Control HUD`。该 HUD 会叠加在使用了方向控制的 --control-mode act ``` -服务示例当前在 `examples/lingbot/stream_lingbot_world_fast.py` 中把 -`control_type` 固定为 `"cam"`。如果要使用 action-control 权重,除了给 demo -传 `--control-mode act`,还需要修改该脚本里的 `PPL_CONFIG["control_type"]`。 +服务示例在 `examples/lingbot/stream_lingbot_world_fast.py` 中设置 +`PPL_CONFIG["control_mode"]="cam"`。session 的控制模式必须与 pipeline 初始化 +时的模式一致。若要使用 action-control 权重,请在启动服务前将该值改为 `"act"`, +并给 demo 传入 `--control-mode act`。 + +#### 帧数与控制输入契约 + +LingBot 只接受完整 latent chunk。`frame_num` 必须为 `4n + 1`,并且得到的 +latent frame 数必须能被 `chunk_size` 整除。默认 `chunk_size=3` 时,合法输出帧数为 +9、21、33、…、321。浏览器接收时长输入,并向下换算为最大合法完整 chunk:16 FPS 下, +10 秒对应 153 帧,20 秒对应 321 帧。直接 API 请求中的非法值会报错。离线 control 文件必须提供覆盖 +目标视频窗口的 pose 和 action。intrinsics 可以是单个静态 `(4,)` 值,也可以是逐视频帧 +的 `(frames, 4)` 值。显式 control tensor 必须与 pipeline 的控制模式、device、dtype +和完整 latent-chunk shape 一致。 #### 显存与分辨率 -LingBot 的 KV cache 会随 `frame_num` 和输出分辨率增长。832x480、81 帧接近 80GB H100 的显存上限,推荐先使用较低分辨率验证链路: +LingBot 的全局 KV cache 会随 `frame_num` 和输出分辨率增长。FSDP 负责分片模型参数,Ulysses 会在 +worker 之间分摊 KV heads,因此更多 GPU 可以支持更长 session,但不会完全线性扩展。以下配置已在 +832x480、`chunk_size=3`、`sample_shift=10.0` 下完成实测: + +| GPU | 时长 | 帧数 | 结果 | +|-----|------|------|------| +| 2 张 H100 80 GB | 10 秒 | 153 | 通过 | +| 2 张 H100 80 GB | 20 秒 | 321 | KV cache CUDA OOM | +| 4 张 H100 80 GB | 20 秒 | 321 | 通过,27/27 chunks | + +四卡 20 秒测试峰值约为 GPU0 58.6 GiB,GPU1-3 各 41.6 GiB。建议先使用 5 秒验证完整链路,再选择 +最大时长: ```bash -# 在 examples/lingbot/stream_lingbot_world_fast.py 中: -PPL_CONFIG["max_area"] = 99840 # 约 416x240 +# 在浏览器 demo 或 session 配置中: +--frame-num 81 ``` 常用调参方向: | 参数 | 影响 | |------|------| -| `PPL_CONFIG["max_area"]` | 降低输出面积,显著降低显存占用 | +| `PPL_CONFIG["resolution"]` | 选择 480p 或 720p 输出面积 | | `--frame-num` | 降低总生成帧数和 latent chunk 数 | | `--chunk-size` | 影响每次生成的 latent chunk 大小 | +发生 OOM 后 parallel worker 会被标记为 failed。此时仅关闭浏览器 session 不足以恢复,必须重启 +stream service 后才能继续测试。 + 服务当前只允许一个 LingBot active session。重新连接前请点击 Stop,或调用: ```bash diff --git a/examples/data/lingbot_world_fast/image.jpg b/examples/data/lingbot_world_fast/image.jpg new file mode 100644 index 0000000..98e4431 Binary files /dev/null and b/examples/data/lingbot_world_fast/image.jpg differ diff --git a/examples/data/lingbot_world_fast/intrinsics.npy b/examples/data/lingbot_world_fast/intrinsics.npy new file mode 100644 index 0000000..3c65d2d Binary files /dev/null and b/examples/data/lingbot_world_fast/intrinsics.npy differ diff --git a/examples/data/lingbot_world_fast/poses.npy b/examples/data/lingbot_world_fast/poses.npy new file mode 100644 index 0000000..194fdb5 Binary files /dev/null and b/examples/data/lingbot_world_fast/poses.npy differ diff --git a/examples/example_config.yaml b/examples/example_config.yaml index 44c6dae..a0a05d7 100644 --- a/examples/example_config.yaml +++ b/examples/example_config.yaml @@ -101,6 +101,21 @@ pipelines: ppl_config_overrides: attn_impl: TORCH_SDPA + # ============================================================ + # lingbot_world_fast - Image to Video + # ============================================================ + + lingbot_world_fast_i2v: + script: lingbot/lingbot_world_fast_image_to_video_h100.py + gpu_count: 1 + output_type: video + input_image_path: examples/data/lingbot_world_fast/image.jpg + timeout_seconds: 3600 + resolution: "480p" + ppl_config_overrides: + attn_impl: TORCH_SDPA + frame_num: 13 + # ============================================================ # hunyuan_video # ============================================================ diff --git a/examples/lingbot/README.md b/examples/lingbot/README.md new file mode 100644 index 0000000..995d366 --- /dev/null +++ b/examples/lingbot/README.md @@ -0,0 +1,343 @@ +# LingBot-World-Fast 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. + +## Model Directory + +The offline example requires the Wan2.2 I2V base model and the LingBot-World-Fast model. The recommended directory +layout is: + +```text +${TF_MODEL_ZOO_PATH}/ +├── Wan2.2-I2V-A14B/ +└── lingbot/ + └── lingbot-world-fast/ +``` + +Set the model root before running the example: + +```bash +export TF_MODEL_ZOO_PATH=/path/to/model_zoo +``` + +## Feature Support + +| Feature | Support | +| --- | --- | +| Offline image-to-video | ✔️ | +| Camera control | ✔️ | +| Continuous chunked generation | ✔️ | +| Single-GPU inference | ✔️ | +| Ulysses Sequence Parallel | ✔️ | +| FSDP | Enabled by default for the streaming example; disabled for the offline example | +| H100 Sage Attention | ✔️ | + +## Files + +### lingbot_world_fast_image_to_video_h100.py + +Offline video generation on H100 GPUs using LingBot-World-Fast with camera control. + +Default configuration: + +- Resolution: `480p` +- Output frames: `81` +- Frame rate: `16 FPS` +- Chunk size: `3` latent frames +- Seed: `42` +- Control mode: `cam` +- Attention backend: `SAGE_ATTN_2_8_8_SM90` +- Default input: `examples/data/lingbot_world_fast/image.jpg` +- Default control directory: `examples/data/lingbot_world_fast/` +- Default output: `work_dirs/lingbot_world_fast_i2v_gpu.mp4` + +## Usage + +### Four H100 GPUs + +The recommended configuration uses four H100 GPUs with Ulysses sequence parallelism for DiT inference: + +```bash +python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ + --gpu_num 4 \ + --model_root "${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B" \ + --fast_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast" +``` + +### Single H100 GPU + +```bash +python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ + --gpu_num 1 \ + --model_root "${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B" \ + --fast_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast" +``` + +### Custom Input and Output + +`--action_path` accepts a directory. In camera-control mode, the directory must contain `poses.npy` and +`intrinsics.npy`. + +```bash +python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ + --gpu_num 4 \ + --model_root "${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B" \ + --fast_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast" \ + --image_path /path/to/input.jpg \ + --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 +``` + +### Repository-Provided Input + +The repository includes an image, poses, and intrinsics that can be used directly. Omit the input, control, and +output options to use these defaults: + +```bash +python examples/lingbot/lingbot_world_fast_image_to_video_h100.py \ + --gpu_num 4 \ + --model_root "${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B" \ + --fast_model_root "${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast" +``` + +## Options + +| Option | Default | Description | +| --- | --- | --- | +| `--gpu_num` | `1` | Number of GPUs used for Ulysses sequence parallelism | +| `--model_root` | `${TF_MODEL_ZOO_PATH}/Wan2.2-I2V-A14B` | Wan2.2 I2V base model directory | +| `--fast_model_root` | `${TF_MODEL_ZOO_PATH}/lingbot/lingbot-world-fast` | LingBot-World-Fast model directory | +| `--image_path` | Bundled `image.jpg` | Input image path | +| `--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 | + +Display the command-line help: + +```bash +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. +- `--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 + residual `spawn_main` child processes. + +## Real-Time Streaming + +The streaming example defaults to camera control, four Ulysses workers, and FSDP. `stream-serve --gpu-num` is +passed to `get_service(gpu_num=...)`; use `CUDA_VISIBLE_DEVICES` to select the physical devices. Do not use +`torchrun` because TeleFuser creates the workers internally. + +### Tested GPU and Duration Limits + +The global KV cache grows with the requested frame count even when FSDP is enabled. The following 832x480 limits +were verified on H100 80 GB GPUs with `chunk_size=3`, `16 FPS`, and `sample_shift=10.0`: + +| GPUs | Duration selected in the page | Frame count | Result | +| --- | --- | --- | --- | +| 2 H100 | 10 seconds | 153 | Passed | +| 2 H100 | 20 seconds | 321 | CUDA OOM while allocating KV cache | +| 4 H100 | 20 seconds | 321 | Passed, 27/27 chunks | + +The four-GPU 20-second test used FSDP and Ulysses degree 4. Peak memory was approximately 58.6 GiB on GPU 0 and +41.6 GiB on GPUs 1-3. These are tested values, not universal limits; other resolutions and concurrent GPU users +change the available capacity. + +### Start a Local TURN Server for VS Code Remote SSH + +VS Code forwards TCP ports. For a laptop browser accessing a remote host, use TURN over TCP and force relay mode. +The following development-only coturn command binds to loopback and uses a small relay range: + +```bash +sudo apt-get install -y coturn + +turnserver -n -m 1 \ + --listening-ip=127.0.0.1 \ + --relay-ip=127.0.0.1 \ + --listening-port=3478 \ + --min-port=49160 --max-port=49200 \ + --user=telefuser:telefuser-turn \ + --realm=telefuser.local \ + --fingerprint --lt-cred-mech \ + --no-tls --no-dtls --no-cli \ + --allow-loopback-peers \ + --simple-log --log-file=/tmp/telefuser-turn.log +``` + +Keep this command running in its own terminal. `--allow-loopback-peers` is required here because both TeleFuser and +coturn run on the same remote host; do not copy this loopback configuration to an internet-facing production TURN +server. Verify the credentials and TCP allocation locally: + +```bash +turnutils_uclient -t -y -c \ + -u telefuser -w telefuser-turn -p 3478 127.0.0.1 +``` + +### Start the Four-GPU LingBot Service + +```bash +TF_MODEL_ZOO_PATH=/path/to/model_zoo \ +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 \ + --gpu-num 4 -p 8088 --host 0.0.0.0 --skip-validation +``` + +Wait for `Starting stream server on 0.0.0.0:8088`, then verify readiness: + +```bash +curl --noproxy '*' http://127.0.0.1:8088/v1/service/health +``` + +### Start the Browser Demo + +Run the demo on the remote TeleFuser host. It proxies signaling requests to port 8088, so the laptop browser does +not need direct access to 8088. + +```bash +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 \ + --action-path examples/data/lingbot_world_fast \ + --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 +``` + +In the VS Code **Ports** panel, forward these remote TCP ports: + +| Remote port | Required local port | Purpose | +| --- | --- | --- | +| `8091` | Any available port | Demo page and proxied SDP/session HTTP requests | +| `3478` | `3478` | Browser TURN-over-TCP connection used by `turn:localhost:3478` | + +Do not forward 8088 when proxying is enabled. Port 3478 must keep local port 3478 unless `--turn-url` is changed to +the alternative local port. Open the forwarded 8091 URL shown by VS Code using `http://`, then hard-refresh the +page after restarting the demo. Do not forward coturn's `49160-49200` relay range through VS Code: browser relay +traffic remains inside the forwarded TURN TCP connection, while that range is used remotely between coturn and the +WebRTC peer. + +If you connect with a normal SSH client instead of VS Code Remote SSH, run this command in a separate terminal on +the laptop. Replace `USER` and `SERVER_HOST` with the SSH login used for the server: + +```bash +ssh -N \ + -o ExitOnForwardFailure=yes \ + -o ServerAliveInterval=30 \ + -L 8091:127.0.0.1:8091 \ + -L 3478:127.0.0.1:3478 \ + USER@SERVER_HOST +``` + +Keep the command running and open `http://localhost:8091`. Add `-p SSH_PORT` or `-i /path/to/private_key` when the +server requires a non-default SSH port or identity file. To obtain an interactive shell from the same connection, +omit `-N`. If local port 8091 is occupied, change only the first mapping, for example +`-L 18091:127.0.0.1:8091`, and open `http://localhost:18091`. + +If local port 3478 is occupied, map another local port such as `-L 13478:127.0.0.1:3478` and change the browser demo +argument to `--turn-url 'turn:localhost:13478?transport=tcp'`. The server-side +`TELEFUSER_TURN_SERVER='turn:127.0.0.1:3478?transport=tcp'` remains unchanged. + +### Run the Demo Entirely on One Machine + +If both the browser and the GPU service run on the same physical machine, neither coturn nor SSH forwarding is +needed. This includes a browser opened directly on the workstation or through its remote desktop, VNC, or noVNC +session. It does not include an SSH shell on the server with the browser still running on a laptop. + +Start the service in the first terminal with TURN variables explicitly removed: + +```bash +env -u TELEFUSER_TURN_SERVER \ + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + TF_MODEL_ZOO_PATH=/path/to/model_zoo \ + CUDA_VISIBLE_DEVICES=0,1,2,3 \ + telefuser stream-serve examples/lingbot/stream_lingbot_world_fast.py \ + --gpu-num 4 -p 8088 --host 127.0.0.1 --skip-validation +``` + +Start the demo in a second terminal without `--turn-url`, TURN credentials, or `--force-turn-relay`: + +```bash +env -u TELEFUSER_TURN_SERVER \ + -u TELEFUSER_TURN_USERNAME \ + -u TELEFUSER_TURN_CREDENTIAL \ + 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 \ + --action-path examples/data/lingbot_world_fast \ + --frame-num 321 \ + --chunk-size 3 \ + --sample-shift 10.0 \ + --fps 16 \ + --no-open +``` + +Open `http://localhost:8091` in the browser on that machine. + +`--image-path` and `--action-path` are paths on the remote host. In real-time `cam` mode, `image.jpg` supplies the +initial frame, only the first row of `intrinsics.npy` is used as fixed intrinsics, and `poses.npy` is not replayed; +camera poses are generated from live controls. + +### Camera Controls + +The page has separate translation and rotation pads: + +| Input | Camera operation | +| --- | --- | +| `W` or `↑` | Move forward | +| `S` or `↓` | Move backward | +| `A` or `←` | Strafe left | +| `D` or `→` | Strafe right | +| `J` | Yaw left | +| `L` | Yaw right | +| `I` | Pitch up | +| `K` | Pitch down | + +Multiple keys can be held together, for example `W+J`. Move/strafe steps default to `0.05` per video frame; +yaw/pitch steps default to `2°` per video frame, with pitch limited to `±85°`. Camera pose and pitch accumulate +across chunks. Controls affect only chunks that have not started inference. Releasing all keys stops requesting new +chunks, and WebRTC repeats the most recent frame while idle. **Reset Control** returns the accumulated pose to the +identity pose. + +### Troubleshooting + +- **Blank 8091 page:** confirm the demo process is listening, open the exact forwarded URL from the VS Code Ports + panel using `http://`, and hard-refresh. The demo uses a threaded HTTP server so VS Code probe connections do not + block page requests. +- **No relay candidate:** confirm local port 3478 maps to remote 3478, the TURN credentials match, and the demo log + shows `iceTransportPolicy=relay`. Browser logs should report at least one `typ=relay` candidate. +- **Static preview:** inspect the DataChannel log for `control_state` and `applying_direction_control`. If the server + logged CUDA OOM, restart the service because a failed parallel worker cannot process another session. +- **Session already active:** click **Stop** or delete it with + `curl -X DELETE http://127.0.0.1:8088/v1/stream/webrtc/`. +- **Residual workers after a forced exit:** terminate stale `spawn_main` processes before restarting. + +For a production/public TURN deployment, TLS, firewall, and relay-port requirements, see the +[stream server guide](../../docs/en/stream_server.md#public-network-deployment). diff --git a/examples/lingbot/lingbot_world_fast_image_to_video_h100.py b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py new file mode 100644 index 0000000..bcb1695 --- /dev/null +++ b/examples/lingbot/lingbot_world_fast_image_to_video_h100.py @@ -0,0 +1,219 @@ +"""LingBot-World-Fast offline image-to-video example. + +Single GPU: + python examples/lingbot/lingbot_world_fast_image_to_video_h100.py + +Four GPUs with Ulysses sequence parallelism: + python examples/lingbot/lingbot_world_fast_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_action_control_inputs, + load_camera_control_inputs, + truncate_control_sequence, +) +from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig +from telefuser.pipelines.lingbot_world_fast.session import ( + LingBotWorldFastChunkRequest, + LingBotWorldFastGenerationSession, + LingBotWorldFastSessionConfig, +) +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( + control_mode="cam", + resolution="480p", + frame_num=81, + chunk_size=3, + sample_shift=10.0, + seed=42, + target_fps=16, + attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90, + enable_fsdp=False, + local_attn_size=-1, + sink_size=0, + max_attention_size=None, + torch_dtype=torch.bfloat16, +) + + +def get_pipeline( + parallelism: int = 1, + model_root: str | None = None, + fast_model_root: str | None = None, +) -> LingBotWorldFastPipeline: + """Load LingBot-World-Fast for offline chunked generation.""" + if model_root is None or fast_model_root is None: + default_model_root = str(TF_MODEL_ZOO_PATH / "Wan2.2-I2V-A14B") + default_fast_model_root = str(TF_MODEL_ZOO_PATH / "lingbot" / "lingbot-world-fast") + else: + default_model_root, default_fast_model_root = model_root, fast_model_root + if parallelism < 1: + raise ValueError(f"parallelism must be positive, got {parallelism}") + dtype = PPL_CONFIG["torch_dtype"] + pipeline = LingBotWorldFastPipeline(device="cuda", torch_dtype=dtype) + pipeline.init( + 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), + 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"], + 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: LingBotWorldFastPipeline, + image: Image.Image, + prompt: str, + 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( + prompt=prompt, + image=image, + control_mode=PPL_CONFIG["control_mode"], + fps=fps, + chunk_size=PPL_CONFIG["chunk_size"], + frame_num=frame_num, + sample_shift=PPL_CONFIG["sample_shift"], + seed=seed, + max_attention_size=PPL_CONFIG["max_attention_size"], + ) + control_context = pipeline.control_context(session_config) + control_builder = LingBotWorldFastControlBuilder(control_context) + if session_config.control_mode == "act": + poses, intrinsics, action = load_action_control_inputs(action_path) + else: + 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=1, + 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_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( + "--fast_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_num: int, + fps: int, + model_root: str, + fast_model_root: str, + output: str | None, +) -> None: + """Generate an offline video with LingBot-World-Fast.""" + pipeline = get_pipeline(gpu_num, model_root, fast_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_num=frame_num, + fps=fps, + ) + elapsed = time.perf_counter() - start + + output_path = Path(output) if output else DEFAULT_OUTPUT_DIR / f"lingbot_world_fast_i2v_{gpu_num}gpu.mp4" + output_path.parent.mkdir(parents=True, exist_ok=True) + save_video(frames, str(output_path), fps=fps, quality=6) + 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 f982bb8..2d7667b 100644 --- a/examples/lingbot/stream_lingbot_world_fast.py +++ b/examples/lingbot/stream_lingbot_world_fast.py @@ -1,83 +1,79 @@ +"""LingBot-World-Fast 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_fast.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 ModelRuntimeConfig -from telefuser.pipelines.lingbot_world_fast.pipeline import ( - LingBotWorldFastPipeline, - LingBotWorldFastPipelineConfig, -) +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig from telefuser.pipelines.lingbot_world_fast.service import LingBotWorldFastService -TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "model_zoo") +RESOLUTION_AREAS = {"480p": 480 * 832, "720p": 720 * 1280} + PPL_CONFIG = dict( - name="lingbot_world_fast_stream", - # LingBot-World-Fast reuses the Wan2.2 base weights (VAE + T5 text encoder + ``google/umt5-xxl`` - # tokenizer), which ship in the shared Wan2.2-I2V-A14B directory. The DiT fast weights live in their - # own ``lingbot-world-fast`` directory, given as an absolute path so the pipeline keeps it standalone - # rather than nesting it under ``checkpoint_dir``. - checkpoint_dir=TF_MODEL_ZOO_PATH + "/Wan2.2-I2V-A14B", - fast_checkpoint_subdir=TF_MODEL_ZOO_PATH + "/lingbot-world-fast", - control_type="cam", - vae_device="cuda", - vae_device_id=0, - text_device="cuda", - text_device_id=0, - dit_device="cuda", - max_area=480 * 832, + parallelism=4, + control_mode="cam", + resolution="480p", + target_fps=16, + attn_impl=AttnImplType.SAGE_ATTN_2_8_8_SM90, + enable_fsdp=True, + local_attn_size=-1, + sink_size=0, torch_dtype=torch.bfloat16, ) -class _LocalModuleManager: - def __init__(self) -> None: - self._modules: list[tuple[str, object, str]] = [] - - def fetch_module( - self, - model_name: str, - file_path: str | None = None, - require_model_path: bool = False, - index: int | None = None, - ): - matches = [(module, path) for name, module, path in self._modules if name == model_name] - if not matches: - return None - module, path = matches[0] - return (module, path) if require_model_path else module - - def add_module(self, module, name: str, path: str = "manual") -> None: - self._modules.append((name, module, path)) - - def get_model_info(self): - return [{"name": name, "path": path} for name, _, path in self._modules] - - -def get_service() -> LingBotWorldFastService: +def get_pipeline( + parallelism: int = PPL_CONFIG["parallelism"], + model_root: str | None = None, + fast_model_root: str | None = None, +) -> LingBotWorldFastPipeline: + """Load LingBot-World-Fast with internal multi-GPU workers.""" + if model_root is None or fast_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_fast_model_root = str(model_zoo_path / "lingbot" / "lingbot-world-fast") + else: + default_model_root, default_fast_model_root = model_root, fast_model_root + if parallelism < 1: + raise ValueError(f"parallelism must be positive, got {parallelism}") dtype = PPL_CONFIG["torch_dtype"] - mm = _LocalModuleManager() - - pipeline = LingBotWorldFastPipeline(device=PPL_CONFIG["dit_device"], torch_dtype=dtype) + pipeline = LingBotWorldFastPipeline(device="cuda", torch_dtype=dtype) pipeline.init( - mm, LingBotWorldFastPipelineConfig( - checkpoint_dir=PPL_CONFIG["checkpoint_dir"], - fast_checkpoint_subdir=PPL_CONFIG["fast_checkpoint_subdir"], - vae_config=ModelRuntimeConfig( - device_type=PPL_CONFIG["vae_device"], - device_id=PPL_CONFIG["vae_device_id"], - torch_dtype=dtype, - ), - text_encoding_config=ModelRuntimeConfig( - device_type=PPL_CONFIG["text_device"], - device_id=PPL_CONFIG["text_device_id"], - torch_dtype=dtype, - ), + 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), + text_encoding_config=ModelRuntimeConfig(device_type="cuda", device_id=0, torch_dtype=dtype), dit_torch_dtype=dtype, - control_type=PPL_CONFIG["control_type"], - max_area=PPL_CONFIG["max_area"], + 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"], + 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 LingBotWorldFastService(pipeline) + 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"]) diff --git a/examples/stream_server/webrtc_bidirectional_demo.py b/examples/stream_server/webrtc_bidirectional_demo.py index a5513f2..0d7ddc7 100644 --- a/examples/stream_server/webrtc_bidirectional_demo.py +++ b/examples/stream_server/webrtc_bidirectional_demo.py @@ -30,6 +30,12 @@ DEFAULT_SERVER_URL = "http://localhost:8088" DEFAULT_PORT = 8091 +DEFAULT_SAMPLE_SHIFT = 10.0 +MAX_GENERATION_SECONDS = 20.0 +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." +) HTML_TEMPLATE = """ @@ -527,16 +533,16 @@ } .dpad { display: grid; - grid-template-columns: 54px 54px 54px; - grid-template-rows: 54px 54px 54px; + grid-template-columns: 44px 44px 44px; + grid-template-rows: 44px 44px 44px; gap: 6px; justify-content: center; - margin: 16px 0 2px; + margin: 8px 0 2px; user-select: none; } .dpad button { - width: 54px; - height: 54px; + width: 44px; + height: 44px; padding: 0; border: 1px solid #cbd5e1; background: #f8fafc; @@ -550,8 +556,31 @@ color: var(--blue); } .dpad .empty { - width: 54px; - height: 54px; + width: 44px; + height: 44px; + } + .control-pads { + display: grid; + grid-template-columns: 1fr 1fr; + gap: 8px; + margin-top: 14px; + } + .control-pad { + padding: 9px 4px 6px; + border: 1px solid #e2e8f0; + border-radius: 6px; + background: #f8fafc; + } + .control-pad h3 { + margin: 0; + color: #475569; + font-size: 11px; + text-align: center; + } + .field-help { + margin-top: 4px; + color: var(--muted); + font-size: 11px; } #messages { height: 210px; @@ -596,7 +625,7 @@

Inputs

- +
@@ -604,22 +633,27 @@
- - + +
- - + +
- + +
+
+ +
+
@@ -661,16 +695,35 @@
-
-
- -
- -
- -
- -
+
+
+

Move · WASD / Arrows

+
+
+ +
+ +
+ +
+ +
+
+
+
+

Rotate · IJKL

+
+
+ +
+ +
+ +
+ +
+
+
@@ -686,8 +739,10 @@ const SERVER_URL = __SERVER_URL__; const RTC_CONFIG = __RTC_CONFIG__; const DEFAULT_IMAGE_PATH = __IMAGE_PATH__; +const DEFAULT_PROMPT = __PROMPT__; const DEFAULT_OPTIONS = __REQUEST_OPTIONS__; const ICE_GATHER_TIMEOUT_MS = __ICE_GATHER_TIMEOUT_MS__; +const MAX_GENERATION_SECONDS = __MAX_GENERATION_SECONDS__; let pc = null; let dc = null; @@ -695,14 +750,18 @@ let cleaning = false; const pressedControls = new Set(); const keyToControl = { - ArrowUp: "up", - ArrowDown: "down", - ArrowLeft: "left", - ArrowRight: "right", - KeyW: "up", - KeyS: "down", - KeyA: "left", - KeyD: "right", + ArrowUp: "w", + ArrowDown: "s", + ArrowLeft: "a", + ArrowRight: "d", + KeyW: "w", + KeyA: "a", + KeyS: "s", + KeyD: "d", + KeyI: "i", + KeyJ: "j", + KeyK: "k", + KeyL: "l", }; function $(id) { @@ -731,19 +790,45 @@ return Number.isFinite(value) ? value : fallback; } +function frameNumForDuration(durationSeconds, fps, chunkSize) { + const duration = Math.min(MAX_GENERATION_SECONDS, Math.max(0.5, durationSeconds)); + const targetFrames = Math.floor(duration * fps) + 1; + const targetLatentFrames = Math.floor((targetFrames - 1) / 4) + 1; + const completeLatentFrames = Math.max(chunkSize, Math.floor(targetLatentFrames / chunkSize) * chunkSize); + return 4 * (completeLatentFrames - 1) + 1; +} + +function updateFrameNum() { + const fps = Math.max(1, numberValue("fps", DEFAULT_OPTIONS.fps ?? 16)); + const chunkSize = Math.max(1, numberValue("chunk-size", DEFAULT_OPTIONS.chunk_size ?? 3)); + const duration = Math.min( + MAX_GENERATION_SECONDS, + Math.max(0.5, numberValue("duration-seconds", 5.0)), + ); + $("duration-seconds").value = duration; + const frameNum = frameNumForDuration(duration, fps, chunkSize); + $("frame-num").value = frameNum; + const actualDuration = (frameNum - 1) / fps; + $("duration-help").textContent = + "Actual duration: " + actualDuration.toFixed(2) + " s · maximum: " + MAX_GENERATION_SECONDS + " s"; +} + function fillDefaults() { + $("prompt").value = DEFAULT_PROMPT; $("image-path").value = DEFAULT_IMAGE_PATH || ""; $("fps").value = DEFAULT_OPTIONS.fps ?? 16; - $("frame-num").value = DEFAULT_OPTIONS.frame_num ?? 81; $("chunk-size").value = DEFAULT_OPTIONS.chunk_size ?? 3; - $("sample-shift").value = DEFAULT_OPTIONS.sample_shift ?? 5.0; + $("duration-seconds").max = MAX_GENERATION_SECONDS; + $("duration-seconds").value = ((DEFAULT_OPTIONS.frame_num ?? 81) - 1) / (DEFAULT_OPTIONS.fps ?? 16); + $("sample-shift").value = DEFAULT_OPTIONS.sample_shift ?? 10.0; $("seed").value = DEFAULT_OPTIONS.seed ?? 42; $("control-mode").value = DEFAULT_OPTIONS.control_mode || "cam"; $("action-path").value = DEFAULT_OPTIONS.action_path || ""; - $("control-move-step").value = DEFAULT_OPTIONS.control_move_step ?? 0.18; - $("control-yaw-step").value = DEFAULT_OPTIONS.control_yaw_step_degrees ?? 10.0; - $("control-lateral-step").value = DEFAULT_OPTIONS.control_lateral_step ?? 0.12; + $("control-move-step").value = DEFAULT_OPTIONS.control_move_step ?? 0.05; + $("control-yaw-step").value = DEFAULT_OPTIONS.control_yaw_step_degrees ?? 2.0; + $("control-lateral-step").value = DEFAULT_OPTIONS.control_lateral_step ?? 0.05; $("show-control-hud").checked = DEFAULT_OPTIONS.show_control_hud ?? true; + updateFrameNum(); } async function fetchJsonWithTimeout(url, options, timeoutMs) { @@ -782,17 +867,20 @@ } 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 ?? 5.0), + 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.18), - control_yaw_step_degrees: numberValue("control-yaw-step", DEFAULT_OPTIONS.control_yaw_step_degrees ?? 10.0), - control_lateral_step: numberValue("control-lateral-step", DEFAULT_OPTIONS.control_lateral_step ?? 0.12), + 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 actionPath = $("action-path").value.trim(); @@ -804,6 +892,10 @@ return 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); @@ -1060,24 +1152,36 @@ def main() -> None: parser.add_argument("--fps", type=int, default=16, help="Output WebRTC FPS and pipeline FPS") parser.add_argument("--frame-num", type=int, default=81, help="Requested LingBot frame count") parser.add_argument("--chunk-size", type=int, default=3, help="LingBot latent chunk size") - parser.add_argument("--sample-shift", type=float, default=5.0, help="LingBot sampler shift") + parser.add_argument("--sample-shift", type=float, default=DEFAULT_SAMPLE_SHIFT, help="LingBot sampler shift") parser.add_argument("--seed", type=int, default=42, help="LingBot random seed") parser.add_argument("--max-attention-size", type=int, default=None, help="Optional LingBot max attention size") parser.add_argument("--max-sequence-length", type=int, default=512, help="LingBot max text sequence length") parser.add_argument("--control-mode", default="cam", choices=("cam", "act"), help="LingBot control mode") - parser.add_argument("--action-path", default="", help="Optional LingBot camera/action control file") - parser.add_argument("--control-move-step", type=float, default=0.18, help="LingBot direction-control move step") + parser.add_argument("--action-path", default="", help="Optional LingBot camera/action control directory") + parser.add_argument("--control-move-step", type=float, default=0.05, help="LingBot video-frame move step") parser.add_argument( "--control-yaw-step-degrees", type=float, - default=10.0, - help="LingBot direction-control yaw step per latent frame", + default=2.0, + help="LingBot yaw step per video frame", ) parser.add_argument( "--control-lateral-step", type=float, - default=0.12, - help="LingBot direction-control lateral strafe step", + default=0.05, + help="LingBot video-frame lateral strafe step", + ) + parser.add_argument( + "--control-pitch-step-degrees", + type=float, + default=2.0, + help="LingBot pitch step per video frame", + ) + parser.add_argument( + "--control-pitch-limit-degrees", + type=float, + default=85.0, + help="LingBot absolute pitch limit", ) parser.add_argument( "--show-control-hud", @@ -1144,6 +1248,8 @@ def main() -> None: "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: @@ -1158,8 +1264,10 @@ def main() -> None: HTML_TEMPLATE.replace("__SERVER_URL__", json.dumps(server_url_for_browser)) .replace("__RTC_CONFIG__", json.dumps(rtc_config)) .replace("__IMAGE_PATH__", json.dumps(args.image_path)) + .replace("__PROMPT__", json.dumps(DEFAULT_PROMPT)) .replace("__REQUEST_OPTIONS__", json.dumps(request_options)) .replace("__ICE_GATHER_TIMEOUT_MS__", str(args.ice_gather_timeout_ms)) + .replace("__MAX_GENERATION_SECONDS__", str(MAX_GENERATION_SECONDS)) ) class Handler(http.server.BaseHTTPRequestHandler): @@ -1208,10 +1316,12 @@ def _proxy(self) -> None: self.wfile.write(resp_body) def do_GET(self) -> None: + body = html.encode("utf-8") self.send_response(200) self.send_header("Content-Type", "text/html; charset=utf-8") + self.send_header("Content-Length", str(len(body))) self.end_headers() - self.wfile.write(html.encode("utf-8")) + self.wfile.write(body) def do_POST(self) -> None: if self._proxy_backend(): @@ -1228,7 +1338,7 @@ def do_DELETE(self) -> None: def log_message(self, format: str, *_args: object) -> None: pass - server = http.server.HTTPServer(("0.0.0.0", args.port), Handler) + server = http.server.ThreadingHTTPServer(("0.0.0.0", args.port), Handler) url = f"http://localhost:{args.port}" print(f"Serving LingBot-World-Fast WebRTC demo at {url}") print(f"Stream server: {args.server_url}") diff --git a/telefuser/core/base_pipeline.py b/telefuser/core/base_pipeline.py index dc07a19..6f72a59 100644 --- a/telefuser/core/base_pipeline.py +++ b/telefuser/core/base_pipeline.py @@ -46,9 +46,6 @@ def wrapped_init(self, *args, **kwargs): @wraps(original_call) def wrapped_call(self, *args, **kwargs): - import gc - - from telefuser.platforms import current_platform from telefuser.utils.profiler import reset_timing_registry reset_timing_registry() @@ -56,9 +53,14 @@ def wrapped_call(self, *args, **kwargs): try: return original_call(self, *args, **kwargs) finally: - # Clear GPU memory after pipeline execution - gc.collect() - current_platform.empty_cache() + if getattr(self, "clear_memory_after_call", True): + # Clear GPU memory after pipeline execution. + import gc + + from telefuser.platforms import current_platform + + gc.collect() + current_platform.empty_cache() cls.__call__ = wrapped_call diff --git a/telefuser/entrypoints/cli/main.py b/telefuser/entrypoints/cli/main.py index 35cda0c..fd1babf 100644 --- a/telefuser/entrypoints/cli/main.py +++ b/telefuser/entrypoints/cli/main.py @@ -136,6 +136,13 @@ def serve( @click.argument("pipe_path") @click.option("--port", "-p", default=8088, type=int, help="Server port") @click.option("--host", default="0.0.0.0", type=str, help="Server host") +@click.option( + "--gpu-num", + "-g", + default=1, + type=click.IntRange(min=1), + help="Number of GPUs passed to a stream pipeline get_service(gpu_num=...) factory", +) @click.option( "--security-level", type=click.Choice(["none", "basic", "strict", "sandbox"], case_sensitive=False), @@ -152,6 +159,7 @@ def stream_serve( pipe_path: str, port: int, host: str, + gpu_num: int, security_level: str, skip_validation: bool, ) -> None: @@ -178,7 +186,14 @@ def stream_serve( from telefuser.service.main import run_stream_server - run_stream_server(pipe_path, port, host, skip_validation=skip_validation, security_level=security_level) + run_stream_server( + pipe_path, + port, + host, + gpu_num=gpu_num, + skip_validation=skip_validation, + security_level=security_level, + ) @main.command(name="validate") diff --git a/telefuser/models/lingbot_world_fast_dit.py b/telefuser/models/lingbot_world_fast_dit.py index c9685cf..5ef627c 100644 --- a/telefuser/models/lingbot_world_fast_dit.py +++ b/telefuser/models/lingbot_world_fast_dit.py @@ -8,8 +8,14 @@ import torch.nn as nn import torch.nn.functional as F from einops import rearrange +from torch.distributed.device_mesh import DeviceMesh from telefuser.core.base_model import BaseModel +from telefuser.core.config import AttentionConfig +from telefuser.distributed.device_mesh import get_ulysses_group, get_ulysses_world_size +from telefuser.distributed.parallel_shard import sequence_parallel_shard, sequence_parallel_unshard +from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads +from telefuser.ops.attention import attention as attn_func from telefuser.ops.normalization import LayerNorm, RMSNorm from telefuser.utils.logging import logger from telefuser.utils.model_weight import init_weights_on_device, load_state_dict @@ -44,6 +50,7 @@ def __init__( self.head_dim = dim // num_heads self.local_attn_size = local_attn_size self.sink_size = sink_size + self.attention_config = AttentionConfig() self.q = nn.Linear(dim, dim) self.k = nn.Linear(dim, dim) @@ -88,7 +95,10 @@ def _apply_causal_rope( dim=-1, ).reshape(seq_len, 1, -1) - return apply_rotary_emb(x, (cos, sin)) + roped = apply_rotary_emb(x[:, :seq_len], (cos, sin)) + if x.shape[1] == seq_len: + return roped + return torch.cat([roped, x[:, seq_len:]], dim=1) def forward( self, @@ -99,16 +109,28 @@ def forward( kv_cache: dict[str, torch.Tensor | int], current_start: int, max_attention_size: int, + device_mesh: DeviceMesh | None = None, ) -> torch.Tensor: q = rearrange(self.norm_q(self.q(x)), "b s (n d) -> b s n d", n=self.num_heads) k = rearrange(self.norm_k(self.k(x)), "b s (n d) -> b s n d", n=self.num_heads) v = rearrange(self.v(x), "b s (n d) -> b s n d", n=self.num_heads) + group = get_ulysses_group(device_mesh) + ulysses_enabled = group is not None and get_ulysses_world_size(device_mesh) > 1 frame_tokens = grid_size[1] * grid_size[2] start_frame = current_start // frame_tokens - - q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame) - k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame) + valid_seq_len = math.prod(grid_size) + if ulysses_enabled: + q = ulysses_scatter_heads(q, group)() + k = ulysses_scatter_heads(k, group)() + v = ulysses_scatter_heads(v, group)() + padded_seq_len = q.shape[1] + q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame)[:, :valid_seq_len] + k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame)[:, :valid_seq_len] + v = v[:, :valid_seq_len] + else: + q = self._apply_causal_rope(q, freqs_cos, freqs_sin, grid_size, start_frame) + k = self._apply_causal_rope(k, freqs_cos, freqs_sin, grid_size, start_frame) num_new_tokens = q.shape[1] current_end = current_start + num_new_tokens @@ -141,15 +163,24 @@ def forward( k_cache = cache_k[:, attn_start:local_end] v_cache = cache_v[:, attn_start:local_end] - q = q.permute(0, 2, 1, 3) - k_cache = k_cache.permute(0, 2, 1, 3) - v_cache = v_cache.permute(0, 2, 1, 3) - out = F.scaled_dot_product_attention(q, k_cache, v_cache, is_causal=False) - out = out.permute(0, 2, 1, 3).contiguous() + out = attn_func( + q, + k_cache, + v_cache, + attention_config=self.attention_config, + input_layout="BSND", + output_layout="BSND", + ) kv_cache["global_end_index"] = current_end kv_cache["local_end_index"] = local_end + if ulysses_enabled: + pad_len = padded_seq_len - num_new_tokens + if pad_len > 0: + out = F.pad(out, (0, 0, 0, 0, 0, pad_len)) + out = ulysses_gather_heads(out, group, num_heads=self.num_heads)() + out = rearrange(out, "b s n d -> b s (n d)") return self.o(out) @@ -162,6 +193,7 @@ def __init__(self, dim: int, num_heads: int, eps: float = 1e-6) -> None: self.dim = dim self.num_heads = num_heads self.head_dim = dim // num_heads + self.attention_config = AttentionConfig() self.q = nn.Linear(dim, dim) self.k = nn.Linear(dim, dim) @@ -176,21 +208,28 @@ def forward( context: torch.Tensor, cache: dict[str, torch.Tensor | bool] | None, ) -> torch.Tensor: - q = rearrange(self.norm_q(self.q(x)), "b s (n d) -> b n s d", n=self.num_heads) + q = rearrange(self.norm_q(self.q(x)), "b s (n d) -> b s n d", n=self.num_heads) if cache is not None and bool(cache.get("is_init", False)): k = cache["k"] v = cache["v"] else: - k = rearrange(self.norm_k(self.k(context)), "b s (n d) -> b n s d", n=self.num_heads) - v = rearrange(self.v(context), "b s (n d) -> b n s d", n=self.num_heads) + k = rearrange(self.norm_k(self.k(context)), "b s (n d) -> b s n d", n=self.num_heads) + v = rearrange(self.v(context), "b s (n d) -> b s n d", n=self.num_heads) if cache is not None: cache["k"] = k cache["v"] = v cache["is_init"] = True - out = F.scaled_dot_product_attention(q, k, v, is_causal=False) - out = rearrange(out, "b n s d -> b s (n d)") + out = attn_func( + q, + k, + v, + attention_config=self.attention_config, + input_layout="BSND", + output_layout="BSND", + ) + out = rearrange(out, "b s n d -> b s (n d)") return self.o(out) @@ -252,6 +291,7 @@ def forward( current_start: int, max_attention_size: int, control_tokens: torch.Tensor | None = None, + device_mesh: DeviceMesh | None = None, ) -> torch.Tensor: modulation = self.modulation.to(dtype=t_mod.dtype, device=t_mod.device) shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = (modulation.unsqueeze(0) + t_mod).chunk( @@ -274,6 +314,7 @@ def forward( kv_cache=kv_cache, current_start=current_start, max_attention_size=max_attention_size, + device_mesh=device_mesh, ) x = self.gate(x, gate_msa, attn_out) @@ -388,6 +429,8 @@ def __init__( freqs = precompute_freqs_cis_3d(head_dim) self.freqs_cos = torch.cat([f.real for f in freqs], dim=-1) self.freqs_sin = torch.cat([f.imag for f in freqs], dim=-1) + self.device_mesh: DeviceMesh | None = None + self.usp_flag = False self.init_weights() @@ -480,6 +523,15 @@ def forward( freqs_cos = self.freqs_cos.to(device=x.device) freqs_sin = self.freqs_sin.to(device=x.device) + full_seq_len = seq_len + + if self.usp_flag: + shard_tensors = [x, t_mod, control_tokens] + shard_dims = [1, 1, 1] + if t_head.shape[1] == seq_len: + shard_tensors.append(t_head) + shard_dims.append(1) + sequence_parallel_shard(self.device_mesh, shard_tensors, shard_dims) if kv_cache is None or crossattn_cache is None: raise ValueError("LingBotWorldFastDiT requires kv_cache and crossattn_cache") @@ -497,11 +549,26 @@ def forward( current_start=current_start, max_attention_size=max_attention_size, control_tokens=control_tokens, + device_mesh=self.device_mesh, ) 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) + def enable_usp(self, device_mesh: DeviceMesh | None = None) -> None: + """Enable Ulysses sequence parallelism for LingBot-World-Fast DiT.""" + self.device_mesh = device_mesh if device_mesh is not None else self.device_mesh + self.usp_flag = get_ulysses_world_size(self.device_mesh) > 1 + + def set_attention_config(self, attention_config: AttentionConfig) -> None: + """Set the unified attention backend for self-attention and cross-attention.""" + logger.info(f"LingBot-World-Fast DiT set attention config to {attention_config.attn_impl}") + for block in self.blocks: + block.self_attn.attention_config = attention_config + block.cross_attn.attention_config = attention_config + @classmethod def from_pretrained( cls, diff --git a/telefuser/models/wan_video_vae.py b/telefuser/models/wan_video_vae.py index 51e9782..84116e4 100644 --- a/telefuser/models/wan_video_vae.py +++ b/telefuser/models/wan_video_vae.py @@ -1,5 +1,7 @@ from __future__ import annotations +from dataclasses import dataclass, field + import torch import torch.distributed as dist import torch.nn as nn @@ -14,6 +16,14 @@ CACHE_T = 2 +@dataclass +class WanVideoVAEStreamingDecodeState: + """Session-owned temporal feature cache for incremental VAE decoding.""" + + 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 @@ -1389,6 +1399,7 @@ def cached_decode_withflag( device: torch.device, is_first_clip: bool, is_last_clip: bool, + decode_state: WanVideoVAEStreamingDecodeState | None = None, ) -> torch.Tensor: """Decode with persistent feature cache for streaming generation. @@ -1400,15 +1411,26 @@ def cached_decode_withflag( device: Target device is_first_clip: If True, clear cache before decoding (first segment) is_last_clip: If True, clear cache after decoding (last segment) + decode_state: Optional session-owned cache. The legacy model-owned cache + is used when omitted. Returns: Decoded video tensor [C, T_out, H_out, W_out] """ + feat_cache = self._feat_cache if decode_state is None else decode_state.feat_cache + feat_idx = self._feat_idx if decode_state is None else decode_state.feat_idx + # Clear cache on first clip if is_first_clip: conv_num = _count_conv3d(self.model.decoder) - self._feat_cache = [None] * conv_num - self._feat_idx = [0] + feat_cache = [None] * conv_num + feat_idx = [0] + if decode_state is None: + self._feat_cache = feat_cache + self._feat_idx = feat_idx + else: + decode_state.feat_cache = feat_cache + decode_state.feat_idx = feat_idx # Add batch dimension if needed if hidden_state.dim() == 4: @@ -1429,25 +1451,29 @@ def cached_decode_withflag( x = self.model.conv2(z) for i in range(iter_): - self._feat_idx[0] = 0 # Reset index for each frame + feat_idx[0] = 0 # Reset index for each frame if i == 0: out = self.model.decoder( x[:, :, i : i + 1, :, :], - feat_cache=self._feat_cache, - feat_idx=self._feat_idx, + feat_cache=feat_cache, + feat_idx=feat_idx, ) else: out_ = self.model.decoder( x[:, :, i : i + 1, :, :], - feat_cache=self._feat_cache, - feat_idx=self._feat_idx, + feat_cache=feat_cache, + feat_idx=feat_idx, ) out = torch.cat([out, out_], 2) # Clear cache on last clip if is_last_clip: - self._feat_cache = [] - self._feat_idx = [0] + if decode_state is None: + self._feat_cache = [] + self._feat_idx = [0] + else: + decode_state.feat_cache = [] + decode_state.feat_idx = [0] video = out.clamp_(-1, 1) diff --git a/telefuser/ops/attention/backends.py b/telefuser/ops/attention/backends.py index e4112d2..1d63f4f 100644 --- a/telefuser/ops/attention/backends.py +++ b/telefuser/ops/attention/backends.py @@ -88,7 +88,7 @@ def _try_import_sage_attn() -> None: """Import Sage Attention.""" global SAGE_ATTN_AVAILABLE, sageattention - for module_name in ["sageattention", "tf_kernel.sageattn2"]: + for module_name in ["tf_kernel.sageattn2", "sageattention"]: try: if importlib.util.find_spec(module_name) is not None: sageattention = importlib.import_module(module_name) diff --git a/telefuser/pipelines/lingbot_world_fast/__init__.py b/telefuser/pipelines/lingbot_world_fast/__init__.py index cbd6796..60d72e0 100644 --- a/telefuser/pipelines/lingbot_world_fast/__init__.py +++ b/telefuser/pipelines/lingbot_world_fast/__init__.py @@ -1,5 +1,4 @@ from .control import ( - CameraControlChunk, build_action_control_chunk, build_camera_control_chunk, load_action_control_inputs, @@ -9,20 +8,25 @@ from .pipeline import LingBotWorldFastPipeline, LingBotWorldFastPipelineConfig from .service import LingBotWorldFastService from .session import ( - LingBotWorldFastRuntimeState, + LingBotWorldFastChunkRequest, + LingBotWorldFastChunkResult, + LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, LingBotWorldFastSessionState, + LingBotWorldFastSessionStatus, ) __all__ = [ - "CameraControlChunk", + "LingBotWorldFastChunkRequest", + "LingBotWorldFastChunkResult", "LingBotWorldFastDenoisingStage", + "LingBotWorldFastGenerationSession", "LingBotWorldFastPipeline", "LingBotWorldFastPipelineConfig", - "LingBotWorldFastRuntimeState", "LingBotWorldFastService", "LingBotWorldFastSessionConfig", "LingBotWorldFastSessionState", + "LingBotWorldFastSessionStatus", "LingBotWorldFastTimesteps", "build_action_control_chunk", "build_camera_control_chunk", diff --git a/telefuser/pipelines/lingbot_world_fast/control.py b/telefuser/pipelines/lingbot_world_fast/control.py index 6ab5862..531282b 100644 --- a/telefuser/pipelines/lingbot_world_fast/control.py +++ b/telefuser/pipelines/lingbot_world_fast/control.py @@ -1,5 +1,6 @@ from __future__ import annotations +from collections.abc import Callable from dataclasses import dataclass from pathlib import Path @@ -143,13 +144,6 @@ def get_plucker_embeddings( return torch.cat([rays_o, rays_d], dim=-1).view(n_frames, height, width, 6) -@dataclass -class CameraControlChunk: - control_tensor: torch.Tensor - num_latent_frames: int - control_type: str - - def build_camera_control_chunk( poses: torch.Tensor, intrinsics: torch.Tensor, @@ -157,7 +151,7 @@ def build_camera_control_chunk( latent_w: int, height: int, width: int, -) -> CameraControlChunk: +) -> torch.Tensor: plucker = get_plucker_embeddings(poses, intrinsics, height, width, only_rays_d=False) control = rearrange( plucker, @@ -167,7 +161,7 @@ def build_camera_control_chunk( h=latent_h, w=latent_w, ).contiguous() - return CameraControlChunk(control_tensor=control, num_latent_frames=control.shape[2], control_type="cam") + return control def build_action_control_chunk( @@ -178,7 +172,7 @@ def build_action_control_chunk( latent_w: int, height: int, width: int, -) -> CameraControlChunk: +) -> torch.Tensor: plucker = get_plucker_embeddings(poses, intrinsics, height, width, only_rays_d=True) action_map = action[:, None, None, :].repeat(1, height, width, 1) combined = torch.cat([plucker, action_map], dim=-1) @@ -190,7 +184,217 @@ def build_action_control_chunk( h=latent_h, w=latent_w, ).contiguous() - return CameraControlChunk(control_tensor=control, num_latent_frames=control.shape[2], control_type="act") + return control + + +def truncate_control_sequence( + poses: object, + intrinsics: object, + action: object | None, + frame_num: int, +) -> tuple[object, object, object | None]: + """Select a video-rate pose/action window and keep the first camera intrinsics fixed.""" + if frame_num < 1 or (frame_num - 1) % 4: + raise ValueError(f"frame_num must be 4n+1, got {frame_num}") + if len(poses) < frame_num: + raise ValueError(f"Pose sequence has {len(poses)} frames but frame_num requires {frame_num}") + + intrinsics_arr = np.asarray(intrinsics) + if intrinsics_arr.ndim not in (1, 2) or intrinsics_arr.shape[-1] != 4: + raise ValueError("Intrinsics must have shape (4,) or (frames, 4)") + if intrinsics_arr.ndim == 2 and len(intrinsics_arr) < 1: + raise ValueError("Intrinsics sequence must contain at least one row") + if action is not None and len(action) < frame_num: + raise ValueError(f"Action sequence has {len(action)} frames but frame_num requires {frame_num}") + + trimmed_poses = poses[:frame_num] + trimmed_intrinsics = intrinsics_arr[0] if intrinsics_arr.ndim == 2 else intrinsics_arr + trimmed_action = action[:frame_num] if action is not None else None + return trimmed_poses, trimmed_intrinsics, trimmed_action + + +@dataclass(frozen=True) +class LingBotWorldFastControlContext: + """Static geometry required to convert external actions into model controls.""" + + control_type: str + device: str | torch.device + control_dtype: torch.dtype + orig_height: int + orig_width: int + height: int + width: int + latent_h: int + latent_w: int + latent_frames: int + chunk_size: int + 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.""" + + def __init__( + self, + builder: LingBotWorldFastControlBuilder, + poses: object, + intrinsics: object, + action: object | None = None, + ) -> None: + self._builder = builder + self._poses = poses + self._intrinsics = intrinsics + self._action = action + self._controls: list[torch.Tensor] | None = None + + def control_at(self, chunk_index: int) -> LingBotWorldFastDeferredControl: + if chunk_index < 0: + raise ValueError(f"chunk_index must be non-negative, got {chunk_index}") + return LingBotWorldFastDeferredControl(lambda: self._materialize()[chunk_index]) + + def _materialize(self) -> list[torch.Tensor]: + if self._controls is None: + self._controls = self._builder.build_sequence(self._poses, self._intrinsics, self._action) + return self._controls + + +class LingBotWorldFastControlBuilder: + """Convert externally owned action data into per-chunk model control tensors.""" + + def __init__(self, context: LingBotWorldFastControlContext) -> None: + self.context = context + + def defer(self, action: dict[str, object]) -> LingBotWorldFastDeferredControl: + """Return a control factory for materialization within the pipeline call.""" + return LingBotWorldFastDeferredControl(lambda: self.build(action)) + + @staticmethod + def _align_action_frames(action: torch.Tensor, target_frames: int) -> torch.Tensor: + if action.ndim == 1: + action = action.unsqueeze(0) + if action.shape[0] == target_frames: + return action + video_rate_frames = (target_frames - 1) * 4 + 1 + if action.shape[0] == video_rate_frames: + return action[::4] + raise ValueError( + f"Action length must be {target_frames} latent frames or {video_rate_frames} video-rate frames, " + f"got {action.shape[0]}" + ) + + @staticmethod + def _validate_poses(poses: torch.Tensor, target_frames: int) -> None: + if poses.shape != (target_frames, 4, 4): + raise ValueError(f"Poses must have shape ({target_frames}, 4, 4), got {tuple(poses.shape)}") + + @staticmethod + def _resample_intrinsics(intrinsics: torch.Tensor, target_frames: int) -> torch.Tensor: + if intrinsics.ndim == 1: + if intrinsics.shape[0] != 4: + raise ValueError(f"Static intrinsics must have shape (4,), got {tuple(intrinsics.shape)}") + return intrinsics.unsqueeze(0).repeat(target_frames, 1) + if intrinsics.ndim != 2 or intrinsics.shape[1] != 4: + raise ValueError(f"Intrinsics must have shape (4,) or (frames, 4), got {tuple(intrinsics.shape)}") + if intrinsics.shape[0] < 1: + raise ValueError("Intrinsics sequence must contain at least one row") + return intrinsics[0:1].repeat(target_frames, 1) + + def build(self, action: dict[str, object]) -> torch.Tensor: + """Build one model control tensor from one external chunk action.""" + if "control_tensor" in action: + return torch.as_tensor( + action["control_tensor"], + device=self.context.device, + dtype=self.context.control_dtype, + ) + poses = action.get("poses") + intrinsics = action.get("intrinsics", self.context.intrinsics) + if poses is None: + raise ValueError("External action requires poses") + poses_t = torch.as_tensor(poses, dtype=torch.float32, device=self.context.device) + intrinsics_t = torch.as_tensor(intrinsics, dtype=torch.float32, device=self.context.device) + self._validate_poses(poses_t, self.context.chunk_size) + intrinsics_t = self._transform_intrinsics(self._resample_intrinsics(intrinsics_t, self.context.chunk_size)) + previous_pose = action.get("previous_pose") + if previous_pose is None: + poses_rel = compute_relative_poses(poses_t, framewise=True) + else: + previous_pose_t = torch.as_tensor(previous_pose, dtype=torch.float32, device=self.context.device) + if previous_pose_t.shape != (4, 4): + raise ValueError(f"Previous pose must have shape (4, 4), got {tuple(previous_pose_t.shape)}") + poses_with_boundary = torch.cat([previous_pose_t.unsqueeze(0), poses_t], dim=0) + poses_rel = compute_relative_poses(poses_with_boundary, framewise=True)[1:] + return self._build_tensor(poses_rel, intrinsics_t, action.get("action")).to(dtype=torch.float32) + + def build_sequence( + self, + poses: object, + intrinsics: object, + action: object | None = None, + ) -> list[torch.Tensor]: + """Build all controls for an offline action sequence without storing it in a session.""" + poses_t = torch.as_tensor(poses, dtype=torch.float32) + source_frames = len(poses_t) + if source_frames < 2: + raise ValueError("Control sequence requires at least two poses") + if poses_t.shape != (source_frames, 4, 4): + raise ValueError(f"Poses must have shape (frames, 4, 4), got {tuple(poses_t.shape)}") + interpolated = interpolate_camera_poses( + src_indices=np.linspace(0, source_frames - 1, source_frames), + src_rot_mat=np.asarray(poses_t[:, :3, :3]), + src_trans_vec=np.asarray(poses_t[:, :3, 3]), + tgt_indices=np.linspace(0, source_frames - 1, self.context.latent_frames), + ) + poses_rel = compute_relative_poses(interpolated.to(self.context.device), framewise=True) + intrinsics_t = torch.as_tensor(intrinsics, dtype=torch.float32, device=self.context.device) + intrinsics_t = self._transform_intrinsics(self._resample_intrinsics(intrinsics_t, len(poses_rel))) + control = self._build_tensor(poses_rel, intrinsics_t, action) + chunks = list(control.to(dtype=torch.float32).split(self.context.chunk_size, dim=2)) + if any(chunk.shape[2] != self.context.chunk_size for chunk in chunks): + raise ValueError("Control sequence must contain complete chunks") + return chunks + + def _transform_intrinsics(self, intrinsics: torch.Tensor) -> torch.Tensor: + return get_ks_transformed( + intrinsics, + height_org=self.context.orig_height, + width_org=self.context.orig_width, + height_resize=self.context.height, + width_resize=self.context.width, + height_final=self.context.height, + width_final=self.context.width, + ) + + def _build_tensor(self, poses: torch.Tensor, intrinsics: torch.Tensor, action: object | None) -> torch.Tensor: + if self.context.control_type == "act": + if action is None: + raise ValueError("Action control mode requires an action sequence") + action_t = torch.as_tensor(action, dtype=torch.float32, device=self.context.device) + action_t = self._align_action_frames(action_t, len(poses)) + return build_action_control_chunk( + poses, + intrinsics, + action_t, + self.context.latent_h, + self.context.latent_w, + self.context.height, + self.context.width, + ) + if action is not None: + raise ValueError("Camera control mode does not accept an action sequence") + return build_camera_control_chunk( + poses, intrinsics, self.context.latent_h, self.context.latent_w, self.context.height, self.context.width + ) def load_camera_control_inputs(action_path: str | Path) -> tuple[np.ndarray, np.ndarray]: diff --git a/telefuser/pipelines/lingbot_world_fast/denoising.py b/telefuser/pipelines/lingbot_world_fast/denoising.py index ba5188c..620b851 100644 --- a/telefuser/pipelines/lingbot_world_fast/denoising.py +++ b/telefuser/pipelines/lingbot_world_fast/denoising.py @@ -4,6 +4,11 @@ import torch +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.distributed.device_mesh import create_device_mesh_from_config, get_ulysses_world_size +from telefuser.distributed.fsdp import shard_model +from telefuser.models.lingbot_world_fast_dit import LingBotWorldFastDiT from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler from telefuser.utils.logging import logger @@ -18,12 +23,113 @@ def select(self, scheduler: FlowUniPCMultistepScheduler, shift: float) -> torch. return scheduler.timesteps[list(self.indices)].clone() -class LingBotWorldFastDenoisingStage: - """Chunk-level denoising for LingBot-World-Fast.""" +@dataclass +class _DenoisingCacheState: + scheduler: FlowUniPCMultistepScheduler + timesteps: torch.Tensor + self_kv_cache: list[dict[str, torch.Tensor | int]] + crossattn_cache: list[dict[str, torch.Tensor | bool]] + generator: torch.Generator - def __init__(self, dit_model, torch_dtype: torch.dtype = torch.bfloat16) -> None: + +class LingBotWorldFastDenoisingStage(BaseStage): + """Chunk-level denoising stage with worker-local persistent KV caches.""" + + def __init__( + self, + name: str, + dit_model: LingBotWorldFastDiT, + model_runtime_config: ModelRuntimeConfig, + ) -> None: + super().__init__(name, model_runtime_config) self.dit = dit_model - self.torch_dtype = torch_dtype + self.dit.set_attention_config(model_runtime_config.attention_config) + self.model_names = ["dit"] + self._cache_registry: dict[int, _DenoisingCacheState] = {} + + def parallel_models(self) -> None: + """Configure Ulysses SP and optional FSDP inside a ParallelWorker.""" + parallel_config = self.model_runtime_config.parallel_config + self.dit.device_mesh = create_device_mesh_from_config(parallel_config) + self.dit.set_attention_config(self.model_runtime_config.attention_config) + if parallel_config.sp_ulysses_degree > 1: + self.dit.enable_usp(self.dit.device_mesh) + if parallel_config.enable_fsdp: + logger.info(f"Enabling FSDP for {self.name}") + self.dit = shard_model( + module=self.dit, + device_id=self.device, + wrap_module_names=self.dit.get_fsdp_module_names(), + param_dtype=self.torch_dtype, + reduce_dtype=self.torch_dtype, + buffer_dtype=self.torch_dtype, + ) + self.onload_models_flag = True + + def _init_self_kv_cache( + self, + batch_size: int, + kv_size: int, + ) -> list[dict[str, torch.Tensor | int]]: + head_dim = self.dit.dim // self.dit.num_heads + ulysses_world_size = get_ulysses_world_size(getattr(self.dit, "device_mesh", None)) + num_heads = self.dit.num_heads + if ulysses_world_size > 1: + num_heads = (num_heads + ulysses_world_size - 1) // ulysses_world_size + shape = (batch_size, kv_size, num_heads, head_dim) + return [ + { + "k": torch.zeros(shape, dtype=self.torch_dtype, device=self.device), + "v": torch.zeros(shape, dtype=self.torch_dtype, device=self.device), + "global_end_index": 0, + "local_end_index": 0, + } + for _ in range(self.dit.num_layers) + ] + + def _init_crossattn_cache( + self, + batch_size: int, + max_sequence_length: int, + ) -> list[dict[str, torch.Tensor | bool]]: + head_dim = self.dit.dim // self.dit.num_heads + shape = (batch_size, max_sequence_length, self.dit.num_heads, head_dim) + return [ + { + "k": torch.zeros(shape, dtype=self.torch_dtype, device=self.device), + "v": torch.zeros(shape, dtype=self.torch_dtype, device=self.device), + "is_init": False, + } + for _ in range(self.dit.num_layers) + ] + + @with_model_offload(["dit"]) + def initialize_cache( + self, + cache_handle: int, + batch_size: int, + kv_size: int, + max_sequence_length: int, + sample_shift: float, + generator_state: list[int], + ) -> 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) + generator = torch.Generator(device=self.device) + generator.set_state(torch.tensor(generator_state, dtype=torch.uint8)) + state = _DenoisingCacheState( + scheduler=scheduler, + timesteps=timesteps, + self_kv_cache=self._init_self_kv_cache(batch_size, kv_size), + crossattn_cache=self._init_crossattn_cache(batch_size, max_sequence_length), + generator=generator, + ) + self._cache_registry[cache_handle] = state + return True @staticmethod def _convert_flow_pred_to_x0( @@ -83,3 +189,57 @@ def denoise_chunk( logger.debug("LingBotWorldFast chunk denoised") return current_latent + + @with_model_offload(["dit"]) + def denoise_and_update_cache( + self, + cache_handle: int, + latent_chunk: torch.Tensor, + condition_chunk: torch.Tensor, + prompt_emb: torch.Tensor, + control_chunk: torch.Tensor | None, + current_start: int, + max_attention_size: int, + ) -> torch.Tensor: + """Denoise a chunk and commit its clean KV state inside each worker.""" + try: + state = self._cache_registry[cache_handle] + except KeyError as exc: + raise KeyError(f"Unknown cache handle {cache_handle}") from exc + denoised = self.denoise_chunk( + latent_chunk=latent_chunk, + condition_chunk=condition_chunk, + prompt_emb=prompt_emb, + timesteps=state.timesteps, + scheduler=state.scheduler, + control_chunk=control_chunk, + self_kv_cache=state.self_kv_cache, + crossattn_cache=state.crossattn_cache, + current_start=current_start, + 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, + ) + return denoised + + def has_cache(self, cache_handle: int) -> bool: + """Return whether this worker owns the requested cache handle.""" + return cache_handle in self._cache_registry + + def list_cache_handles(self) -> tuple[int, ...]: + """Return registered cache handles for diagnostics and tests.""" + return tuple(sorted(self._cache_registry)) + + def release_cache(self, cache_handle: int) -> bool: + """Idempotently release worker-local state for one generation session.""" + return self._cache_registry.pop(cache_handle, None) is not None diff --git a/telefuser/pipelines/lingbot_world_fast/pipeline.py b/telefuser/pipelines/lingbot_world_fast/pipeline.py index eebf5c9..d2615dc 100644 --- a/telefuser/pipelines/lingbot_world_fast/pipeline.py +++ b/telefuser/pipelines/lingbot_world_fast/pipeline.py @@ -1,44 +1,40 @@ from __future__ import annotations -import base64 import math from collections.abc import Callable from dataclasses import dataclass, field from pathlib import Path -import cv2 import numpy as np import torch from PIL import Image from telefuser.core.base_pipeline import BasePipeline -from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.config import AttentionConfig, ModelRuntimeConfig, ParallelConfig from telefuser.models.lingbot_world_fast_dit import LingBotWorldFastDiT from telefuser.models.t5_tokenizer import HuggingfaceTokenizer from telefuser.models.wan_video_text_encoder import WanTextEncoder from telefuser.models.wan_video_vae import WanVideoVAE -from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler from telefuser.utils.logging import logger from telefuser.utils.model_weight import load_state_dict from telefuser.utils.profiler import ProfilingContext4Debug - -from .control import ( - build_action_control_chunk, - build_camera_control_chunk, - compute_relative_poses, - get_ks_transformed, - interpolate_camera_poses, - load_action_control_inputs, - load_camera_control_inputs, +from telefuser.worker.parallel_worker import ParallelWorker + +from .control import LingBotWorldFastControlContext +from .denoising import LingBotWorldFastDenoisingStage +from .session import ( + LingBotWorldFastChunkRequest, + LingBotWorldFastChunkResult, + LingBotWorldFastGenerationSession, + LingBotWorldFastSessionConfig, + LingBotWorldFastSessionStatus, ) -from .denoising import LingBotWorldFastDenoisingStage, LingBotWorldFastTimesteps -from .session import LingBotWorldFastRuntimeState, LingBotWorldFastSessionConfig @dataclass class LingBotWorldFastPipelineConfig: checkpoint_dir: str = "" - fast_checkpoint_subdir: str = "lingbot_world_fast" + fast_checkpoint_path: str = "lingbot_world_fast" vae_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) dit_torch_dtype: torch.dtype = torch.bfloat16 @@ -48,18 +44,22 @@ class LingBotWorldFastPipelineConfig: max_area: int = 480 * 832 local_attn_size: int = -1 sink_size: int = 0 + parallel_config: ParallelConfig = field(default_factory=ParallelConfig) + attention_config: AttentionConfig = field(default_factory=AttentionConfig) class LingBotWorldFastPipeline(BasePipeline): """Pipeline wrapper for LingBot-World-Fast chunked causal generation.""" + clear_memory_after_call = False + def __init__(self, device: str, torch_dtype: torch.dtype = torch.bfloat16) -> None: super().__init__(device=device, torch_dtype=torch_dtype) self.height_division_factor = 16 self.width_division_factor = 16 def _get_stages(self) -> list: - return [] + return [self.denoise_stage] if hasattr(self, "denoise_stage") else [] @staticmethod def _notify_progress( @@ -81,7 +81,9 @@ def _runtime_device(self, runtime_config: ModelRuntimeConfig) -> torch.device: return torch.device(f"cuda:{runtime_config.device_id}") return torch.device(runtime_config.device_type) - def init(self, module_manager, config: LingBotWorldFastPipelineConfig) -> None: + def init(self, config: LingBotWorldFastPipelineConfig) -> None: + if config.control_type not in {"cam", "act"}: + raise ValueError(f"Unsupported LingBot control_type: {config.control_type!r}") self.config = config checkpoint_root = Path(config.checkpoint_dir).expanduser().resolve() self._model_info = [{"name": "lingbot_world_fast", "path": str(checkpoint_root)}] @@ -101,17 +103,27 @@ def init(self, module_manager, config: LingBotWorldFastPipelineConfig) -> None: self.vae.load_state_dict(vae_state_dict, strict=False) self.vae = self.vae.to(device=self.vae_device, dtype=self.torch_dtype).eval() - fast_path = checkpoint_root / config.fast_checkpoint_subdir + fast_path = checkpoint_root / config.fast_checkpoint_path + dit_device = "cpu" if config.parallel_config.world_size > 1 else self.device self.dit = LingBotWorldFastDiT.from_pretrained( str(fast_path), torch_dtype=config.dit_torch_dtype, control_type=config.control_type, config=self._build_dit_config(config), - ).to(self.device) + ).to(dit_device) self.dit.eval().requires_grad_(False) - self.denoise_stage = LingBotWorldFastDenoisingStage(self.dit, torch_dtype=self.torch_dtype) - self.timesteps = LingBotWorldFastTimesteps() + pipeline_device = torch.device(self.device) + dit_runtime_config = ModelRuntimeConfig( + device_type=pipeline_device.type, + device_id=pipeline_device.index or 0, + torch_dtype=config.dit_torch_dtype, + attention_config=config.attention_config, + parallel_config=config.parallel_config, + ) + denoise_stage = LingBotWorldFastDenoisingStage("lingbot_world_fast_denoise", self.dit, dit_runtime_config) + self.denoise_stage = ParallelWorker(denoise_stage) if config.parallel_config.world_size > 1 else denoise_stage + self._next_cache_handle = 0 @staticmethod def _build_dit_config(config: LingBotWorldFastPipelineConfig) -> dict[str, object]: @@ -146,12 +158,19 @@ def encode_prompt(self, prompt: str) -> torch.Tensor: return prompt_emb.to(self.device) @torch.inference_mode() - def decode_video_cached(self, latents: torch.Tensor, is_first_clip: bool, is_last_clip: bool) -> torch.Tensor: + def decode_video_cached( + self, + session: LingBotWorldFastGenerationSession, + latents: torch.Tensor, + is_first_clip: bool, + is_last_clip: bool, + ) -> torch.Tensor: return self.vae.cached_decode_withflag( latents, device=self.vae_device, is_first_clip=is_first_clip, is_last_clip=is_last_clip, + decode_state=session.decoder_state, ) @staticmethod @@ -176,6 +195,89 @@ def _best_output_size(w: int, h: int, expected_area: int, dw: int = 16, dh: int return ow1, oh1 return ow2, oh2 + def control_context(self, session_config: LingBotWorldFastSessionConfig) -> LingBotWorldFastControlContext: + """Return control geometry without allocating a generation runtime.""" + self._validate_session_config(session_config) + width, height = self._best_output_size( + session_config.image.width, + session_config.image.height, + 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 + if session_config.intrinsics is None: + focal = float(max(self.config.orig_width, self.config.orig_height)) + intrinsics = torch.tensor( + [focal, focal, self.config.orig_width * 0.5, self.config.orig_height * 0.5], + dtype=torch.float32, + device=self.device, + ) + else: + intrinsics = torch.as_tensor(session_config.intrinsics, dtype=torch.float32, device=self.device) + if intrinsics.ndim == 2: + if intrinsics.shape[0] < 1 or intrinsics.shape[1] != 4: + raise ValueError( + f"Session intrinsics must have shape (4,) or (frames, 4), got {tuple(intrinsics.shape)}" + ) + intrinsics = intrinsics[0] + if intrinsics.shape != (4,): + raise ValueError( + f"Session intrinsics must have shape (4,) or (frames, 4), got {tuple(intrinsics.shape)}" + ) + return LingBotWorldFastControlContext( + control_type=self.config.control_type, + device=self.device, + control_dtype=torch.float32, + orig_height=self.config.orig_height, + orig_width=self.config.orig_width, + height=height, + width=width, + latent_h=height // 8, + latent_w=width // 8, + latent_frames=latent_frames, + chunk_size=session_config.chunk_size, + intrinsics=intrinsics, + ) + + def _validate_session_config(self, session_config: LingBotWorldFastSessionConfig) -> None: + if session_config.control_mode != self.config.control_type: + raise ValueError( + f"Session control_mode {session_config.control_mode!r} does not match " + f"pipeline control_type {self.config.control_type!r}" + ) + if not isinstance(session_config.chunk_size, int) or isinstance(session_config.chunk_size, bool): + 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}" + ) + + def _validate_control(self, session: LingBotWorldFastGenerationSession, control: torch.Tensor) -> None: + expected_device = torch.device(self.device) + if control.device.type != expected_device.type or ( + expected_device.index is not None and control.device.index != expected_device.index + ): + raise ValueError(f"Control device must be compatible with {self.device}, got {control.device}") + if control.dtype != torch.float32: + raise ValueError(f"Control dtype must be torch.float32, got {control.dtype}") + channels_per_pixel = 7 if self.config.control_type == "act" else 6 + height_factor = max(1, session.height // session.latent_h) + width_factor = max(1, session.width // session.latent_w) + expected_channels = channels_per_pixel * height_factor * width_factor + expected_shape = (1, expected_channels, session.chunk_size, session.latent_h, session.latent_w) + if tuple(control.shape) != expected_shape: + 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 @@ -203,227 +305,196 @@ def _encode_condition_video(self, image_tensor: torch.Tensor, frame_num: int) -> mask = mask.view(1, mask.shape[1] // 4, 4, latent_h, latent_w).transpose(1, 2)[0] return torch.cat([mask, latent], dim=0) - def _init_self_kv_cache( - self, - batch_size: int, - kv_size: int, - dtype: torch.dtype, - device: str | torch.device, - ) -> list[dict[str, torch.Tensor | int]]: - head_dim = self.dit.dim // self.dit.num_heads - shape = (batch_size, kv_size, self.dit.num_heads, head_dim) - return [ - { - "k": torch.zeros(shape, dtype=dtype, device=device), - "v": torch.zeros(shape, dtype=dtype, device=device), - "global_end_index": 0, - "local_end_index": 0, - } - for _ in range(self.dit.num_layers) - ] - - def _init_crossattn_cache( - self, - batch_size: int, - dtype: torch.dtype, - device: str | torch.device, - max_sequence_length: int, - ) -> list[dict[str, torch.Tensor | bool]]: - head_dim = self.dit.dim // self.dit.num_heads - shape = (batch_size, self.dit.num_heads, max_sequence_length, head_dim) - return [ - { - "k": torch.zeros(shape, dtype=dtype, device=device), - "v": torch.zeros(shape, dtype=dtype, device=device), - "is_init": False, - } - for _ in range(self.dit.num_layers) - ] - - def _populate_session_controls(self, session_config: LingBotWorldFastSessionConfig) -> None: - if session_config.action_path is None: - return - if session_config.poses is not None and session_config.intrinsics is not None: - if session_config.control_mode != "act" or session_config.action is not None: - return - - if session_config.control_mode == "act": - poses, intrinsics, action = load_action_control_inputs(session_config.action_path) - if session_config.action is None: - session_config.action = action - else: - poses, intrinsics = load_camera_control_inputs(session_config.action_path) + def _release_session_cache(self, session: LingBotWorldFastGenerationSession) -> bool: + cache_handle = session.cache_handle + if cache_handle is None: + return True + try: + if isinstance(self.denoise_stage, ParallelWorker): + self.denoise_stage.release_cache(cache_handle, sync=True) + else: + self.denoise_stage.release_cache(cache_handle) + except Exception as exc: + logger.error(f"Failed to release LingBot cache handle {cache_handle}: {exc}") + return False + session.cache_handle = None + return True + + def release_session(self, session: LingBotWorldFastGenerationSession) -> None: + """Idempotently release cache and decoder state owned by a session.""" + with session.transaction_lock: + cache_released = self._release_session_cache(session) + session.decoder_state.feat_cache = [] + session.decoder_state.feat_idx = [0] + session.active = False + if not cache_released: + session.status = LingBotWorldFastSessionStatus.POISONED + session.poisoned_reason = f"Failed to release cache handle {session.cache_handle}" + elif session.status != LingBotWorldFastSessionStatus.POISONED: + session.status = LingBotWorldFastSessionStatus.RELEASED + + def close(self) -> None: + """Deterministically close the multi-process denoising worker group.""" + denoise_stage = getattr(self, "denoise_stage", None) + if isinstance(denoise_stage, ParallelWorker): + denoise_stage.close() + + def __del__(self) -> None: + """Best-effort fallback for callers that do not explicitly close the pipeline.""" + try: + self.close() + except Exception: + pass - if session_config.poses is None: - session_config.poses = poses - if session_config.intrinsics is None: - session_config.intrinsics = intrinsics + @torch.inference_mode() + def __call__( + self, + session: LingBotWorldFastGenerationSession, + request: LingBotWorldFastChunkRequest, + progress_callback: Callable[..., None] | None = None, + ) -> LingBotWorldFastChunkResult: + """Initialize a session on first use and generate one controlled chunk.""" + if not session.transaction_lock.acquire(blocking=False): + raise RuntimeError("LingBot session already has a chunk in progress") + try: + self._validate_chunk_request(session, request) + resolved_control: torch.Tensor | None = None + if session.status == LingBotWorldFastSessionStatus.NEW: + try: + + def materialize_first_control() -> None: + nonlocal resolved_control + resolved_control = self._resolve_control(request.control) + + initialized = self._create_initialized_session( + session.config, + progress_callback, + before_cache=materialize_first_control, + ) + session.__dict__.update( + (key, value) for key, value in initialized.__dict__.items() if key != "transaction_lock" + ) + except Exception as exc: + session.status = LingBotWorldFastSessionStatus.POISONED + session.poisoned_reason = f"{type(exc).__name__}: {exc}" + self.release_session(session) + raise + if resolved_control is None: + resolved_control = self._resolve_control(request.control) + self._validate_control(session, resolved_control) + return self._generate_session_chunk(session, request, resolved_control, progress_callback) + finally: + session.transaction_lock.release() @staticmethod - def _truncate_control_inputs_to_frame_num( - poses: object, - intrinsics: object, - action: object | None, - frame_num: int, - ) -> tuple[object, object, object | None, int]: - requested_frame_num = ((int(frame_num) - 1) // 4) * 4 + 1 - pose_frame_num = ((len(poses) - 1) // 4) * 4 + 1 - effective_frame_num = min(requested_frame_num, pose_frame_num) - - trimmed_poses = poses[:effective_frame_num] - intrinsics_arr = np.asarray(intrinsics) - if intrinsics_arr.ndim > 1: - trimmed_intrinsics = intrinsics[:effective_frame_num] - else: - trimmed_intrinsics = intrinsics - trimmed_action = action[:effective_frame_num] if action is not None else None - return trimmed_poses, trimmed_intrinsics, trimmed_action, effective_frame_num + def _validate_chunk_request( + session: LingBotWorldFastGenerationSession, + request: LingBotWorldFastChunkRequest, + ) -> None: + if session.status == LingBotWorldFastSessionStatus.POISONED: + 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: + raise RuntimeError("Cannot generate a chunk from an inactive LingBot session") + if request.chunk_index != session.current_chunk_index: + raise ValueError( + f"Chunk request index {request.chunk_index} does not match session index {session.current_chunk_index}" + ) @staticmethod - def _align_action_frames(action: torch.Tensor, target_frames: int) -> torch.Tensor: - if action.ndim == 1: - action = action.unsqueeze(0) - if action.shape[0] == target_frames: - return action - - sampled = action[::4] - if sampled.shape[0] >= target_frames: - return sampled[:target_frames] - if action.shape[0] == 1: - return action.repeat(target_frames, 1) - if action.shape[0] < target_frames: - raise ValueError(f"Action length {action.shape[0]} is shorter than target latent frames {target_frames}") - - indices = torch.linspace(0, action.shape[0] - 1, target_frames, device=action.device) - return action.index_select(0, indices.round().long()) - - def _prepare_control_chunks( + def _resolve_control(control: torch.Tensor | Callable[[], torch.Tensor]) -> torch.Tensor: + resolved = control() if callable(control) else control + if not isinstance(resolved, torch.Tensor): + raise TypeError("Deferred control factory must return a torch.Tensor") + return resolved + + def _generate_session_chunk( self, - session_config: LingBotWorldFastSessionConfig, - lat_f: int, - lat_h: int, - lat_w: int, - height: int, - width: int, - chunk_size: int, - ) -> list[torch.Tensor] | None: - if session_config.poses is None or session_config.intrinsics is None: - return None - - poses = torch.as_tensor(session_config.poses, dtype=torch.float32) - intrinsics = torch.as_tensor(session_config.intrinsics, dtype=torch.float32) - intrinsics = get_ks_transformed( - intrinsics, - height_org=self.config.orig_height, - width_org=self.config.orig_width, - height_resize=height, - width_resize=width, - height_final=height, - width_final=width, + session: LingBotWorldFastGenerationSession, + request: LingBotWorldFastChunkRequest, + control: torch.Tensor, + progress_callback: Callable[..., None] | None, + ) -> LingBotWorldFastChunkResult: + """Execute a chunk while the caller holds the session transaction lock.""" + session.status = LingBotWorldFastSessionStatus.RUNNING + try: + chunk_frames = self.generate_next_chunk( + session, + control=control, + progress_callback=progress_callback, + ) + 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: + 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" ) - - len_c2ws = len(poses) - lat_f_target = int(lat_f - (lat_f % chunk_size)) - poses_interp = interpolate_camera_poses( - src_indices=np.linspace(0, len_c2ws - 1, len_c2ws), - src_rot_mat=np.asarray(poses[:, :3, :3]), - src_trans_vec=np.asarray(poses[:, :3, 3]), - tgt_indices=np.linspace(0, len_c2ws - 1, lat_f_target), + return LingBotWorldFastChunkResult( + chunk_index=request.chunk_index, + frames=chunk_frames, + emitted_frames=session.emitted_frames, + done=not session.active, + session_id=request.session_id, ) - poses_rel = compute_relative_poses(poses_interp.to(self.device), framewise=True) - intrinsics = intrinsics[0].to(self.device).repeat(len(poses_rel), 1) - - if session_config.control_mode == "act" and session_config.action is not None: - action = torch.as_tensor(session_config.action, dtype=torch.float32, device=self.device) - action = self._align_action_frames(action, len(poses_rel)) - chunk = build_action_control_chunk(poses_rel, intrinsics, action, lat_h, lat_w, height, width) - else: - chunk = build_camera_control_chunk(poses_rel, intrinsics, lat_h, lat_w, height, width) - - return list(chunk.control_tensor.split(chunk_size, dim=2)) - def build_control_override( + @torch.inference_mode() + def generate_video( self, - runtime: LingBotWorldFastRuntimeState, - chunk: dict, - ) -> torch.Tensor | None: - if "control_tensor" in chunk: - return torch.as_tensor(chunk["control_tensor"], device=self.device, dtype=self.torch_dtype) - - poses = chunk.get("poses") - intrinsics = chunk.get("intrinsics") - if poses is None or intrinsics is None: - return None - - poses_t = torch.as_tensor(poses, dtype=torch.float32, device=self.device) - intrinsics_t = torch.as_tensor(intrinsics, dtype=torch.float32, device=self.device) - if intrinsics_t.ndim == 1: - intrinsics_t = intrinsics_t.unsqueeze(0).repeat(poses_t.shape[0], 1) - - intrinsics_t = get_ks_transformed( - intrinsics_t, - height_org=self.config.orig_height, - width_org=self.config.orig_width, - height_resize=runtime.height, - width_resize=runtime.width, - height_final=runtime.height, - width_final=runtime.width, - ) - poses_rel = compute_relative_poses(poses_t, framewise=True) - - if chunk.get("action") is not None: - action_t = torch.as_tensor(chunk["action"], dtype=torch.float32, device=self.device) - action_t = self._align_action_frames(action_t, len(poses_rel)) - control = build_action_control_chunk( - poses_rel, - intrinsics_t, - action_t, - runtime.latent_h, - runtime.latent_w, - runtime.height, - runtime.width, - ) - else: - control = build_camera_control_chunk( - poses_rel, - intrinsics_t, - runtime.latent_h, - runtime.latent_w, - runtime.height, - runtime.width, - ) - return control.control_tensor - - @ProfilingContext4Debug("create_runtime") + session_config: LingBotWorldFastSessionConfig, + controls: list[torch.Tensor | Callable[[], torch.Tensor]], + progress_callback: Callable[..., None] | None = None, + ) -> list[Image.Image]: + """Drain externally prepared controls through the single-chunk API.""" + session = LingBotWorldFastGenerationSession(config=session_config) + frames: list[Image.Image] = [] + try: + for chunk_index, control in enumerate(controls): + result = self( + session, + LingBotWorldFastChunkRequest( + chunk_index=chunk_index, + control=control, + ), + progress_callback=progress_callback, + ) + frames.extend(result.frames) + if session.active: + raise ValueError("Control sequence ended before the generation session completed") + finally: + self.release_session(session) + return frames + + @ProfilingContext4Debug("initialize_session") @torch.inference_mode() - def create_runtime( + def _create_initialized_session( self, session_config: LingBotWorldFastSessionConfig, progress_callback: Callable[..., None] | None = None, - ) -> LingBotWorldFastRuntimeState: - self._notify_progress(progress_callback, "loading_controls") - self._populate_session_controls(session_config) + before_cache: Callable[[], None] | None = None, + ) -> LingBotWorldFastGenerationSession: + """Allocate model state for the first chunk; callers never invoke this directly.""" + control_context = self.control_context(session_config) self._notify_progress(progress_callback, "encoding_prompt", device=str(self.text_device)) prompt_emb = self.encode_prompt(session_config.prompt) self._notify_progress(progress_callback, "prompt_encoded") - w0, h0 = session_config.image.size - width, height = self._best_output_size(w0, h0, self.config.max_area) - height, width = self.check_resize_height_width(height, width) + width = control_context.width + height = control_context.height 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 - 1) // 4) * 4 + 1 - if session_config.poses is not None and session_config.intrinsics is not None: - session_config.poses, session_config.intrinsics, session_config.action, frame_num = ( - self._truncate_control_inputs_to_frame_num( - poses=session_config.poses, - intrinsics=session_config.intrinsics, - action=session_config.action, - frame_num=frame_num, - ) - ) + frame_num = session_config.frame_num self._notify_progress( progress_callback, "encoding_condition_video", @@ -433,11 +504,9 @@ def create_runtime( latent_condition = self._encode_condition_video(image_tensor, frame_num) self._notify_progress(progress_callback, "condition_video_encoded") - lat_h = height // 8 - lat_w = width // 8 - lat_f = (frame_num - 1) // 4 + 1 - lat_f = int(lat_f - (lat_f % session_config.chunk_size)) - frame_num = (lat_f - 1) * 4 + 1 + lat_h = control_context.latent_h + lat_w = control_context.latent_w + lat_f = control_context.latent_frames patch_area = self.dit.patch_size[1] * self.dit.patch_size[2] frame_tokens = (lat_h * lat_w) // patch_area kv_size = self._resolve_self_kv_size( @@ -445,7 +514,6 @@ def create_runtime( latent_frames=lat_f, config=self.config, ) - max_seq_len = session_config.chunk_size * frame_tokens max_attention_size = ( kv_size if session_config.max_attention_size is None else int(session_config.max_attention_size) ) @@ -468,69 +536,60 @@ def create_runtime( ) 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)) - self._notify_progress(progress_callback, "preparing_controls") - control_chunks = self._prepare_control_chunks( - session_config=session_config, - lat_f=lat_f, - lat_h=lat_h, - lat_w=lat_w, - height=height, - width=width, - chunk_size=session_config.chunk_size, - ) - - runtime = LingBotWorldFastRuntimeState( + if before_cache is not None: + before_cache() + cache_handle = self._next_cache_handle + self._next_cache_handle += 1 + session = LingBotWorldFastGenerationSession( prompt_emb=prompt_emb, - encoded_image_latent=latent_condition, + config=session_config, noise_chunks=noise_chunks, condition_chunks=condition_chunks, - control_chunks=control_chunks, - timesteps=torch.empty(0, dtype=torch.int64), - self_kv_cache=self._init_self_kv_cache( - batch_size=1, - kv_size=kv_size, - dtype=self.torch_dtype, - device=self.device, - ), - crossattn_cache=self._init_crossattn_cache( - batch_size=1, - dtype=self.torch_dtype, - device=self.device, - max_sequence_length=session_config.max_sequence_length, - ), latent_h=lat_h, latent_w=lat_w, latent_f=lat_f, height=height, width=width, - max_seq_len=max_seq_len, frame_tokens=frame_tokens, chunk_size=session_config.chunk_size, max_attention_size=max_attention_size, - scheduler=FlowUniPCMultistepScheduler(num_train_timesteps=1000, shift=1, use_dynamic_shifting=False), - generator=generator, - kv_local_attn_size=self.config.local_attn_size, - kv_sink_size=self.config.sink_size if self.config.local_attn_size != -1 else 0, + cache_handle=cache_handle, ) - runtime.timesteps = self.timesteps.select(runtime.scheduler, session_config.sample_shift) + try: + initialize_cache_kwargs = dict( + cache_handle=cache_handle, + batch_size=1, + kv_size=kv_size, + max_sequence_length=session_config.max_sequence_length, + sample_shift=session_config.sample_shift, + generator_state=generator.get_state().tolist(), + ) + if isinstance(self.denoise_stage, ParallelWorker): + self.denoise_stage.initialize_cache(**initialize_cache_kwargs, sync=True) + else: + self.denoise_stage.initialize_cache(**initialize_cache_kwargs) + except Exception: + self._release_session_cache(session) + raise if session_config.world_kv_binding is not None: - runtime.world_kv_binding = session_config.world_kv_binding + session.world_kv_binding = session_config.world_kv_binding try: - runtime.world_kv_binding.on_runtime_created(runtime, session_config) - if runtime.world_kv_cached_latents: - logger.info(f"world_kv: fast-forward {len(runtime.world_kv_cached_latents)} chunks (decode-only)") + session.world_kv_binding.on_runtime_created(session, session_config) + if session.world_kv_cached_latents: + logger.info(f"world_kv: fast-forward {len(session.world_kv_cached_latents)} chunks (decode-only)") except Exception as exc: logger.warning(f"world_kv on_runtime_created failed; falling back to cold run: {exc}") - runtime.world_kv_cached_latents = {} + session.world_kv_cached_latents = {} self._notify_progress(progress_callback, "runtime_created", width=width, height=height, latent_frames=lat_f) logger.info(f"LingBot runtime created: {width}x{height}, latent={lat_f}x{lat_h}x{lat_w}") - return runtime + session.status = LingBotWorldFastSessionStatus.READY + return session @torch.inference_mode() def generate_next_chunk( self, - runtime: LingBotWorldFastRuntimeState, - control_override: torch.Tensor | None = None, + runtime: LingBotWorldFastGenerationSession, + control: torch.Tensor, progress_callback: Callable[..., None] | None = None, ) -> list[Image.Image]: if runtime.current_chunk_index >= len(runtime.noise_chunks): @@ -541,9 +600,7 @@ def generate_next_chunk( idx = runtime.current_chunk_index latent_chunk = runtime.noise_chunks[idx] condition_chunk = runtime.condition_chunks[idx] - control_chunk = control_override - if control_chunk is None and runtime.control_chunks is not None and idx < len(runtime.control_chunks): - control_chunk = runtime.control_chunks[idx] + control_chunk = control self._notify_progress(progress_callback, "denoising_chunk", index=idx) current_start = idx * runtime.chunk_size * runtime.frame_tokens @@ -555,33 +612,19 @@ def generate_next_chunk( denoised = cached_latent.to(device=self.device, dtype=self.torch_dtype) else: with ProfilingContext4Debug("denoise_chunk"): - denoised = self.denoise_stage.denoise_chunk( + denoise_kwargs = dict( + cache_handle=runtime.cache_handle, latent_chunk=latent_chunk, condition_chunk=condition_chunk, prompt_emb=runtime.prompt_emb, - timesteps=runtime.timesteps, - scheduler=runtime.scheduler, control_chunk=control_chunk, - self_kv_cache=runtime.self_kv_cache, - crossattn_cache=runtime.crossattn_cache, - current_start=current_start, - max_attention_size=runtime.max_attention_size, - generator=runtime.generator, - ) - - self._notify_progress(progress_callback, "updating_cache", index=idx) - with ProfilingContext4Debug("kv_cache_update_forward"): - self.dit( - x=denoised.to(dtype=self.torch_dtype), - timestep=torch.zeros((1,), dtype=torch.float32, device=self.device), - context=runtime.prompt_emb, - y=condition_chunk, - control_tensor=control_chunk, - kv_cache=runtime.self_kv_cache, - crossattn_cache=runtime.crossattn_cache, current_start=current_start, max_attention_size=runtime.max_attention_size, ) + if isinstance(self.denoise_stage, ParallelWorker): + denoised = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs, sync=True) + else: + denoised = self.denoise_stage.denoise_and_update_cache(**denoise_kwargs) if runtime.world_kv_binding is not None: try: @@ -592,6 +635,7 @@ def generate_next_chunk( self._notify_progress(progress_callback, "decoding_chunk", index=idx, device=str(self.vae_device)) with ProfilingContext4Debug("vae_decode"): frames = self.decode_video_cached( + runtime, denoised, is_first_clip=(idx == 0), is_last_clip=(idx == len(runtime.noise_chunks) - 1), @@ -603,14 +647,3 @@ def generate_next_chunk( if runtime.current_chunk_index >= len(runtime.noise_chunks): runtime.active = False return images - - @staticmethod - def encode_frames_to_b64(frames: list[Image.Image], quality: int = 85) -> list[str]: - encoded: list[str] = [] - for frame in frames: - rgb = np.asarray(frame.convert("RGB")) - bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) - ok, buf = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, int(quality)]) - if ok: - encoded.append(base64.b64encode(buf.tobytes()).decode("ascii")) - return encoded diff --git a/telefuser/pipelines/lingbot_world_fast/service.py b/telefuser/pipelines/lingbot_world_fast/service.py index 1ea3a72..51b9e7f 100644 --- a/telefuser/pipelines/lingbot_world_fast/service.py +++ b/telefuser/pipelines/lingbot_world_fast/service.py @@ -1,6 +1,7 @@ from __future__ import annotations import asyncio +import base64 import gc import math import queue @@ -8,41 +9,55 @@ import time import uuid from collections.abc import AsyncGenerator, Callable +from pathlib import Path +import cv2 +import numpy as np import torch from PIL import Image, ImageDraw from telefuser.utils.logging import logger from telefuser.utils.profiler import ProfilingContext4Debug +from .control import LingBotWorldFastControlBuilder, LingBotWorldFastControlContext from .pipeline import LingBotWorldFastPipeline from .session import ( - LingBotWorldFastRuntimeState, + LingBotWorldFastChunkRequest, + LingBotWorldFastGenerationSession, LingBotWorldFastSessionConfig, LingBotWorldFastSessionState, ) _DIRECTION_ALIASES = { - "ArrowUp": "up", - "ArrowDown": "down", - "ArrowLeft": "left", - "ArrowRight": "right", - "KeyW": "up", - "KeyS": "down", - "KeyA": "left", - "KeyD": "right", - "w": "up", - "s": "down", - "a": "left", - "d": "right", - "up": "up", - "down": "down", - "left": "left", - "right": "right", - "forward": "up", - "backward": "down", + "ArrowUp": "w", + "ArrowDown": "s", + "ArrowLeft": "j", + "ArrowRight": "l", + "KeyW": "w", + "KeyA": "a", + "KeyS": "s", + "KeyD": "d", + "KeyI": "i", + "KeyJ": "j", + "KeyK": "k", + "KeyL": "l", + "w": "w", + "a": "a", + "s": "s", + "d": "d", + "i": "i", + "j": "j", + "k": "k", + "l": "l", + "up": "w", + "down": "s", + "left": "j", + "right": "l", + "forward": "w", + "backward": "s", } -_ACTION_DIRECTIONS = ("up", "down", "left", "right") +_ACTION_DIRECTIONS = ("w", "a", "s", "d") +MAX_GENERATION_SECONDS = 20.0 class LingBotWorldFastService: @@ -59,6 +74,7 @@ def start(self) -> None: def stop(self) -> None: for session_id in list(self._sessions.keys()): self.close_session(session_id) + self.pipeline.close() def has_session(self, session_id: str) -> bool: return session_id in self._sessions @@ -86,30 +102,44 @@ def create_session(self, config: dict) -> str: session_id = config.get("session_id") or str(uuid.uuid4()) image = self._load_image(config) + intrinsics = config.get("intrinsics") + if intrinsics is None and config.get("action_path"): + intrinsics = np.load(Path(config["action_path"]) / "intrinsics.npy") + + fps = int(config.get("fps") or self.default_fps) + frame_num = int(config.get("frame_num", 81)) + if fps <= 0: + raise ValueError(f"fps must be positive, got {fps}") + duration_seconds = (frame_num - 1) / fps + if duration_seconds > MAX_GENERATION_SECONDS: + raise ValueError( + f"LingBot streaming duration must not exceed {MAX_GENERATION_SECONDS:g} seconds, " + f"got {duration_seconds:g} seconds" + ) session_config = LingBotWorldFastSessionConfig( prompt=config.get("prompt", ""), image=image, control_mode=config.get("control_mode", "cam"), - fps=int(config.get("fps") or self.default_fps), + fps=fps, chunk_size=int(config.get("chunk_size", 3)), - frame_num=int(config.get("frame_num", 81)), - sample_shift=float(config.get("sample_shift", 5.0)), + 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"), - offload_model=bool(config.get("offload_model", False)), max_sequence_length=int(config.get("max_sequence_length", 512)), - action_path=config.get("action_path"), - poses=config.get("poses"), - intrinsics=config.get("intrinsics"), - action=config.get("action"), - control_move_step=float(config.get("control_move_step", 0.18)), - control_yaw_step_degrees=float(config.get("control_yaw_step_degrees", 10.0)), - control_lateral_step=float(config.get("control_lateral_step", 0.12)), + 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_context = self.pipeline.control_context(session_config) state = LingBotWorldFastSessionState( config=session_config, + control_context=control_context, output_queue=asyncio.Queue(), loop=None, ) @@ -127,10 +157,21 @@ def _put_output(state: LingBotWorldFastSessionState, payload: dict) -> None: logger.warning(f"Failed to enqueue LingBotWorld output: {exc}") @staticmethod - def _release_runtime(state: LingBotWorldFastSessionState) -> None: - if state.runtime is not None: - state.runtime.active = False - state.runtime = None + def _encode_frames_to_b64(frames: list[Image.Image], quality: int = 85) -> list[str]: + """Serialize generated frames for the streaming transport.""" + encoded: list[str] = [] + for frame in frames: + rgb = np.asarray(frame.convert("RGB")) + bgr = cv2.cvtColor(rgb, cv2.COLOR_RGB2BGR) + ok, buf = cv2.imencode(".jpg", bgr, [cv2.IMWRITE_JPEG_QUALITY, int(quality)]) + if ok: + encoded.append(base64.b64encode(buf.tobytes()).decode("ascii")) + return encoded + + def _release_generation_session(self, state: LingBotWorldFastSessionState) -> None: + if state.generation_session is not None: + self.pipeline.release_session(state.generation_session) + state.generation_session = None with state.control_lock: state.pressed_controls.clear() state.queued_controls.clear() @@ -140,8 +181,8 @@ def _release_runtime(state: LingBotWorldFastSessionState) -> None: def _emit_preview_frame(self, state: LingBotWorldFastSessionState) -> None: image = state.config.image.convert("RGB") - width, height = self.pipeline._best_output_size(image.width, image.height, self.pipeline.config.max_area) - height, width = self.pipeline.check_resize_height_width(height, width) + control_context = state.control_context or self.pipeline.control_context(state.config) + width, height = control_context.width, control_context.height preview = image.resize((width, height), Image.BICUBIC) self._put_output( state, @@ -150,7 +191,7 @@ def _emit_preview_frame(self, state: LingBotWorldFastSessionState) -> None: "index": -1, "fps": state.config.fps, "timestamp": time.time(), - "frames_b64": self.pipeline.encode_frames_to_b64([preview]), + "frames_b64": self._encode_frames_to_b64([preview]), }, ) @@ -163,24 +204,71 @@ def _direction_from_chunk(chunk: dict) -> str | None: @staticmethod def _is_explicit_control_chunk(chunk: dict) -> bool: - return "control_tensor" in chunk or (chunk.get("poses") is not None and chunk.get("intrinsics") is not None) + return "control_tensor" in chunk or chunk.get("poses") is not None @staticmethod - def _pose_matrix(yaw: float, position: list[float]) -> list[list[float]]: - cos_y = math.cos(yaw) - sin_y = math.sin(yaw) - return [ - [cos_y, 0.0, sin_y, position[0]], - [0.0, 1.0, 0.0, position[1]], - [-sin_y, 0.0, cos_y, position[2]], - [0.0, 0.0, 0.0, 1.0], - ] - - def _default_intrinsics(self) -> list[list[float]]: - width = float(self.pipeline.config.orig_width) - height = float(self.pipeline.config.orig_height) - focal = max(width, height) - return [[focal, focal, width * 0.5, height * 0.5]] + def _rotation_matrix(axis: str, angle: float) -> np.ndarray: + cos_a = math.cos(angle) + sin_a = math.sin(angle) + if axis == "x": + return np.asarray([[1.0, 0.0, 0.0], [0.0, cos_a, -sin_a], [0.0, sin_a, cos_a]]) + if axis == "y": + return np.asarray([[cos_a, 0.0, sin_a], [0.0, 1.0, 0.0], [-sin_a, 0.0, cos_a]]) + raise ValueError(f"Unsupported rotation axis: {axis}") + + @classmethod + def _integrate_camera_step( + cls, + c2w: np.ndarray, + pitch: float, + controls: set[str], + config: LingBotWorldFastSessionConfig, + ) -> tuple[np.ndarray, float]: + pitch_delta = 0.0 + pitch_step = math.radians(float(config.control_pitch_step_degrees)) + if "i" in controls: + pitch_delta += pitch_step + if "k" in controls: + pitch_delta -= pitch_step + pitch_limit = math.radians(float(config.control_pitch_limit_degrees)) + new_pitch = pitch + pitch_delta + if -pitch_limit <= new_pitch <= pitch_limit: + pitch = new_pitch + else: + pitch_delta = 0.0 + + yaw_delta = 0.0 + yaw_step = math.radians(float(config.control_yaw_step_degrees)) + if "j" in controls: + yaw_delta -= yaw_step + if "l" in controls: + yaw_delta += yaw_step + + rotation = c2w[:3, :3] + rotation_new = cls._rotation_matrix("y", yaw_delta) @ rotation @ cls._rotation_matrix("x", pitch_delta) + forward = np.asarray([rotation_new[0, 2], 0.0, rotation_new[2, 2]]) + right = np.asarray([rotation_new[0, 0], 0.0, rotation_new[2, 0]]) + forward_norm = np.linalg.norm(forward) + right_norm = np.linalg.norm(right) + if forward_norm > 0: + forward /= forward_norm + 1e-6 + if right_norm > 0: + right /= right_norm + 1e-6 + + movement = np.zeros(3) + if "w" in controls: + movement += forward * float(config.control_move_step) + if "s" in controls: + movement -= forward * float(config.control_move_step) + if "d" in controls: + movement += right * float(config.control_lateral_step) + if "a" in controls: + movement -= right * float(config.control_lateral_step) + + result = np.eye(4) + result[:3, :3] = rotation_new + result[:3, 3] = c2w[:3, 3] + movement + return result, pitch @staticmethod def _draw_triangle( @@ -208,7 +296,17 @@ def _overlay_control_hud(cls, frames: list[Image.Image], controls: list[str] | N if not controls: return frames - active = set(controls) + controls_active = set(controls) + active = { + direction + for direction, source_controls in { + "up": {"w"}, + "down": {"s"}, + "left": {"a", "j"}, + "right": {"d", "l"}, + }.items() + if controls_active & source_controls + } out: list[Image.Image] = [] for frame in frames: image = frame.convert("RGB") @@ -245,66 +343,46 @@ def _overlay_control_hud(cls, frames: list[Image.Image], controls: list[str] | N def _build_directional_control_chunk( self, state: LingBotWorldFastSessionState, - runtime: LingBotWorldFastRuntimeState, + control_context: LingBotWorldFastControlContext, ) -> dict | None: with state.control_lock: controls = set(state.pressed_controls) | set(state.queued_controls) state.queued_controls.clear() - yaw = float(state.control_yaw) - position = [float(v) for v in state.control_position] + c2w = np.asarray(state.control_c2w, dtype=np.float64) + pitch = float(state.control_pitch) + initialized = state.control_initialized - if not controls or runtime.current_chunk_index >= len(runtime.noise_chunks): + if not controls: return None - latent_frames = int(runtime.noise_chunks[runtime.current_chunk_index].shape[2]) - move_step = float(state.config.control_move_step) - yaw_step = math.radians(float(state.config.control_yaw_step_degrees)) - lateral_step = float(state.config.control_lateral_step) + latent_frames = control_context.chunk_size poses: list[list[list[float]]] = [] action_rows: list[list[float]] = [] + previous_pose = c2w.copy() if initialized else None - for _ in range(latent_frames): - strafe = 0 - if "left" in controls and "right" not in controls: - yaw += yaw_step - strafe = -1 - elif "right" in controls and "left" not in controls: - yaw -= yaw_step - strafe = 1 - - forward_x = math.sin(yaw) - forward_z = math.cos(yaw) - right_x = math.cos(yaw) - right_z = -math.sin(yaw) - if "up" in controls and "down" not in controls: - position[0] += forward_x * move_step - position[2] += forward_z * move_step - elif "down" in controls and "up" not in controls: - position[0] -= forward_x * move_step - position[2] -= forward_z * move_step - if strafe: - position[0] += right_x * lateral_step * strafe - position[2] += right_z * lateral_step * strafe - - poses.append(self._pose_matrix(yaw, position)) + if not initialized: + poses.append(c2w.tolist()) + action_rows.append([1.0 if name in controls else 0.0 for name in _ACTION_DIRECTIONS]) + intervals = latent_frames if initialized else latent_frames - 1 + for _ in range(intervals): + for _ in range(4): + c2w, pitch = self._integrate_camera_step(c2w, pitch, controls, state.config) + poses.append(c2w.tolist()) action_rows.append([1.0 if name in controls else 0.0 for name in _ACTION_DIRECTIONS]) with state.control_lock: - state.control_yaw = yaw - state.control_position = position + state.control_c2w = c2w.tolist() + state.control_pitch = pitch + state.control_initialized = True chunk: dict = { "type": "control", "poses": poses, - "intrinsics": self._default_intrinsics(), "controls": sorted(controls), } - model_control_type = getattr( - getattr(self.pipeline, "dit", None), - "control_type", - self.pipeline.config.control_type, - ) - if model_control_type == "act": + if previous_pose is not None: + chunk["previous_pose"] = previous_pose.tolist() + if control_context.control_type == "act": chunk["action"] = action_rows return chunk @@ -320,8 +398,9 @@ def _update_direction_controls(self, state: LingBotWorldFastSessionState, chunk: elif event == "reset": state.pressed_controls.clear() state.queued_controls.clear() - state.control_position = [0.0, 0.0, 0.0] - state.control_yaw = 0.0 + state.control_c2w = np.eye(4).tolist() + state.control_pitch = 0.0 + state.control_initialized = False else: state.pressed_controls.add(direction) state.queued_controls.add(direction) @@ -363,20 +442,32 @@ def _run_worker_loop( ) -> None: try: self._emit_preview_frame(state) - emit_status("initializing_runtime") - state.runtime = self.pipeline.create_runtime(state.config, progress_callback=emit_status) - runtime = state.runtime + control_context = state.control_context or self.pipeline.control_context(state.config) + control_builder = LingBotWorldFastControlBuilder(control_context) + state.generation_session = LingBotWorldFastGenerationSession(config=state.config) + runtime = state.generation_session emit_status( "runtime_ready", - width=runtime.width, - height=runtime.height, - latent_frames=runtime.latent_f, - total_chunks=len(runtime.noise_chunks), + width=control_context.width, + height=control_context.height, + latent_frames=control_context.latent_frames, + total_chunks=control_context.latent_frames // control_context.chunk_size, ) chunk_index = 0 while state.active and runtime.active: - control_override = None - incoming = None + with state.control_lock: + controls_held = bool(state.pressed_controls) + if controls_held: + try: + incoming = state.pending_inputs.get_nowait() + except queue.Empty: + incoming = {"type": "direction_control"} + else: + incoming = state.pending_inputs.get() + if incoming.get("type") == "stop": + break + + explicit_control = incoming if self._is_explicit_control_chunk(incoming) else None applied_controls = None while True: try: @@ -387,32 +478,39 @@ def _run_worker_loop( incoming = next_item break if self._is_explicit_control_chunk(next_item): - incoming = next_item + explicit_control = next_item if incoming and incoming.get("type") == "stop": break - if incoming: - control_override = self.pipeline.build_control_override(runtime, incoming) - if control_override is None: - directional_chunk = self._build_directional_control_chunk(state, runtime) - if directional_chunk is not None: - applied_controls = directional_chunk["controls"] - emit_status( - "applying_direction_control", - index=chunk_index, - controls=applied_controls, - move_step=state.config.control_move_step, - yaw_step_degrees=state.config.control_yaw_step_degrees, - lateral_step=state.config.control_lateral_step, - ) - control_override = self.pipeline.build_control_override(runtime, directional_chunk) + if explicit_control is not None: + control = control_builder.defer(explicit_control) + else: + directional_chunk = self._build_directional_control_chunk(state, control_context) + if directional_chunk is None: + continue + applied_controls = directional_chunk["controls"] + emit_status( + "applying_direction_control", + index=chunk_index, + controls=applied_controls, + move_step=state.config.control_move_step, + yaw_step_degrees=state.config.control_yaw_step_degrees, + lateral_step=state.config.control_lateral_step, + pitch_step_degrees=state.config.control_pitch_step_degrees, + ) + control = control_builder.defer(directional_chunk) emit_status("generating_chunk", index=chunk_index) - frames = self.pipeline.generate_next_chunk( + result = self.pipeline( runtime, - control_override=control_override, + LingBotWorldFastChunkRequest( + chunk_index=runtime.current_chunk_index, + session_id=session_id, + control=control, + ), progress_callback=emit_status, ) + frames = result.frames if not frames: break if state.config.show_control_hud and applied_controls: @@ -422,7 +520,7 @@ def _run_worker_loop( "index": chunk_index, "fps": state.config.fps, "timestamp": time.time(), - "frames_b64": self.pipeline.encode_frames_to_b64(frames), + "frames_b64": self._encode_frames_to_b64(frames), } self._put_output(state, payload) emit_status("chunk_sent", index=chunk_index, frames=len(frames)) @@ -441,13 +539,17 @@ def _run_worker_loop( finally: state.active = False self._put_output(state, {"type": "done"}) - self._release_runtime(state) + self._release_generation_session(state) def push_chunk(self, session_id: str, chunk: dict) -> None: state = self._sessions.get(session_id) if state is None: return - if chunk.get("type") == "control" and self._update_direction_controls(state, chunk): + is_direction_action = chunk.get("type") == "control" and self._update_direction_controls(state, chunk) + if is_direction_action: + event = str(chunk.get("event") or chunk.get("action") or "press").lower() + if event not in {"release", "keyup", "end", "reset"}: + state.pending_inputs.put({"type": "direction_control"}) return state.pending_inputs.put(chunk) @@ -467,7 +569,7 @@ async def pull_chunks(self, session_id: str) -> AsyncGenerator[dict, None]: ) state.worker_thread.start() - while state.active: + while True: chunk = await state.output_queue.get() if chunk.get("type") == "done": break @@ -479,5 +581,5 @@ def close_session(self, session_id: str) -> None: return state.active = False state.pending_inputs.put({"type": "stop"}) - self._release_runtime(state) + self._release_generation_session(state) logger.info(f"LingBotWorld session closed: {session_id}") diff --git a/telefuser/pipelines/lingbot_world_fast/session.py b/telefuser/pipelines/lingbot_world_fast/session.py index ad6e429..3a25c5a 100644 --- a/telefuser/pipelines/lingbot_world_fast/session.py +++ b/telefuser/pipelines/lingbot_world_fast/session.py @@ -3,12 +3,16 @@ import asyncio import queue import threading +from collections.abc import Callable from dataclasses import dataclass, field +from enum import Enum import torch from PIL import Image -from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler +from telefuser.models.wan_video_vae import WanVideoVAEStreamingDecodeState + +from .control import LingBotWorldFastControlContext @dataclass @@ -19,50 +23,82 @@ class LingBotWorldFastSessionConfig: fps: int = 16 chunk_size: int = 3 frame_num: int = 81 - sample_shift: float = 5.0 + sample_shift: float = 10.0 seed: int = 42 max_attention_size: int | None = None - offload_model: bool = False max_sequence_length: int = 512 - action_path: str | None = None - poses: object | None = None - intrinsics: object | None = None - action: object | None = None # Optional CacheSeek world_kv reuse. None preserves baseline behavior. world_kv_binding: object | None = None - control_move_step: float = 0.18 - control_yaw_step_degrees: float = 10.0 - control_lateral_step: float = 0.12 + intrinsics: object | None = None + control_move_step: float = 0.05 + control_yaw_step_degrees: float = 2.0 + control_lateral_step: float = 0.05 + control_pitch_step_degrees: float = 2.0 + control_pitch_limit_degrees: float = 85.0 show_control_hud: bool = True @dataclass -class LingBotWorldFastRuntimeState: - prompt_emb: torch.Tensor - encoded_image_latent: torch.Tensor - noise_chunks: list[torch.Tensor] - condition_chunks: list[torch.Tensor] - control_chunks: list[torch.Tensor] | None - timesteps: torch.Tensor - self_kv_cache: list[dict[str, torch.Tensor | int]] - crossattn_cache: list[dict[str, torch.Tensor | bool]] - latent_h: int - latent_w: int - latent_f: int - height: int - width: int - max_seq_len: int - frame_tokens: int - chunk_size: int - max_attention_size: int - scheduler: FlowUniPCMultistepScheduler +class LingBotWorldFastChunkRequest: + """Inputs for one explicitly indexed LingBot video chunk.""" + + chunk_index: int + control: torch.Tensor | Callable[[], torch.Tensor] = field(repr=False) + session_id: str | None = None + + def __post_init__(self) -> None: + if self.chunk_index < 0: + raise ValueError(f"chunk_index must be non-negative, got {self.chunk_index}") + if not isinstance(self.control, torch.Tensor) and not callable(self.control): + raise TypeError("Each chunk request requires a model control tensor or deferred control factory") + + +@dataclass +class LingBotWorldFastChunkResult: + """Output and progress metadata for one generated LingBot chunk.""" + + chunk_index: int + frames: list[Image.Image] = field(repr=False) + emitted_frames: int = 0 + done: bool = False + session_id: str | None = None + + +class LingBotWorldFastSessionStatus(str, Enum): + """Transaction and lifecycle state for one generation session.""" + + READY = "ready" + RUNNING = "running" + NEW = "new" + COMMITTED = "committed" + POISONED = "poisoned" + RELEASED = "released" + + +@dataclass +class LingBotWorldFastGenerationSession: + """Externally owned state for one chunked LingBot generation.""" + + 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) + latent_h: int = 0 + latent_w: int = 0 + latent_f: int = 0 + height: int = 0 + width: int = 0 + frame_tokens: int = 0 + chunk_size: int = 0 + max_attention_size: int = 0 + cache_handle: int | None = None + decoder_state: WanVideoVAEStreamingDecodeState = field(default_factory=WanVideoVAEStreamingDecodeState) current_chunk_index: int = 0 emitted_frames: int = 0 active: bool = True - generator: torch.Generator | None = None - # KV geometry in latent frames; -1 means full-length KV. - kv_local_attn_size: int = -1 - kv_sink_size: int = 0 + status: LingBotWorldFastSessionStatus = LingBotWorldFastSessionStatus.NEW + poisoned_reason: str | None = None + transaction_lock: object = field(default_factory=threading.RLock, repr=False) # CacheSeek world_kv binding and decode-only latents for fast-forward hits. world_kv_binding: object | None = None world_kv_cached_latents: dict[int, torch.Tensor] = field(default_factory=dict) @@ -71,7 +107,8 @@ class LingBotWorldFastRuntimeState: @dataclass class LingBotWorldFastSessionState: config: LingBotWorldFastSessionConfig - runtime: LingBotWorldFastRuntimeState | None = None + control_context: LingBotWorldFastControlContext | None = None + generation_session: LingBotWorldFastGenerationSession | None = None pending_inputs: "queue.Queue[dict]" = field(default_factory=queue.Queue) output_queue: asyncio.Queue | None = None worker_thread: threading.Thread | None = None @@ -79,6 +116,14 @@ class LingBotWorldFastSessionState: active: bool = True pressed_controls: set[str] = field(default_factory=set) queued_controls: set[str] = field(default_factory=set) - control_position: list[float] = field(default_factory=lambda: [0.0, 0.0, 0.0]) - control_yaw: float = 0.0 + control_c2w: list[list[float]] = field( + default_factory=lambda: [ + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 0.0, 1.0], + ] + ) + control_pitch: float = 0.0 + control_initialized: bool = False control_lock: object = field(default_factory=threading.Lock) diff --git a/telefuser/service/core/container.py b/telefuser/service/core/container.py index 4f7a598..c04f607 100644 --- a/telefuser/service/core/container.py +++ b/telefuser/service/core/container.py @@ -250,6 +250,7 @@ def initialize_stream_service( self, pipe_path: str, skip_validation: bool = False, + gpu_num: int = 1, ) -> bool: """Initialize stream pipeline service (alternative to initialize_all).""" self.stream_pipeline_service = StreamPipelineService( @@ -258,6 +259,7 @@ def initialize_stream_service( ) return self.stream_pipeline_service.start_service( ppl_file=pipe_path, + gpu_num=gpu_num, skip_validation=skip_validation, ) diff --git a/telefuser/service/core/stream_pipeline_service.py b/telefuser/service/core/stream_pipeline_service.py index e7095d1..7b6f0b5 100644 --- a/telefuser/service/core/stream_pipeline_service.py +++ b/telefuser/service/core/stream_pipeline_service.py @@ -17,6 +17,7 @@ from __future__ import annotations import asyncio +import inspect import threading from collections.abc import AsyncGenerator from types import ModuleType @@ -79,7 +80,9 @@ class StreamPipelineService: """Loads a stream pipeline module and drives streaming execution. Pipeline file convention: - def get_service() -> ServerPushService | BidirectionalService + def get_service(gpu_num: int = 1) -> ServerPushService | BidirectionalService + + Factories without a ``gpu_num`` parameter remain supported. """ def __init__( @@ -109,7 +112,7 @@ def __init__( # -- lifecycle ----------------------------------------------------------- - def start_service(self, ppl_file: str, skip_validation: bool = False) -> bool: + def start_service(self, ppl_file: str, skip_validation: bool = False, gpu_num: int = 1) -> bool: """Load module, call get_service(), detect mode, and start.""" if self.is_running: logger.warning("Stream service is already running") @@ -135,7 +138,12 @@ def start_service(self, ppl_file: str, skip_validation: bool = False) -> bool: "returning a ServerPushService or BidirectionalService" ) - self.service = self._module.get_service() + get_service = self._module.get_service + signature = inspect.signature(get_service) + accepts_gpu_num = "gpu_num" in signature.parameters or any( + parameter.kind is inspect.Parameter.VAR_KEYWORD for parameter in signature.parameters.values() + ) + self.service = get_service(gpu_num=gpu_num) if accepts_gpu_num else get_service() self.stream_mode = self._detect_mode(self.service) self.service.start() self.is_running = True diff --git a/telefuser/service/main.py b/telefuser/service/main.py index 4f5aeff..2bd88ed 100644 --- a/telefuser/service/main.py +++ b/telefuser/service/main.py @@ -94,6 +94,7 @@ def run_stream_server( enable_rate_limit: bool = True, skip_validation: bool = False, security_level: str | None = None, + gpu_num: int = 1, ) -> None: """Run the TeleFuser stream server. @@ -109,7 +110,11 @@ def run_stream_server( container = ServiceContainer.create(config=server_config) - if not container.initialize_stream_service(pipe_path=pipe_path, skip_validation=skip_validation): + if not container.initialize_stream_service( + pipe_path=pipe_path, + gpu_num=gpu_num, + skip_validation=skip_validation, + ): raise RuntimeError("Failed to initialize stream service") logger.info("Stream service initialized successfully") diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index fa27332..de9dff5 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -8,6 +8,7 @@ import gc import os +import threading import time from collections.abc import Callable from datetime import timedelta @@ -63,6 +64,8 @@ def _worker_loop( from telefuser.utils.profiler import mark_as_worker_process mark_as_worker_process() + args = None + kwargs = None try: device = stage.device if world_size > 1: @@ -124,9 +127,15 @@ def _worker_loop( traceback.print_exc() logger.error(f"Error in worker loop (rank {rank}): {e}") - queue_out.put(e) # any exception caught in the worker will be raised to the main process + message = f"Parallel worker rank {rank} failed with {type(e).__name__}: {e}" + try: + worker_error = type(e)(message) + except Exception: + worker_error = RuntimeError(message) + queue_out.put(worker_error) finally: - del stage, args, kwargs + args = None + kwargs = None current_platform.synchronize() gc.collect() current_platform.empty_cache() @@ -155,6 +164,10 @@ def __init__( self.name: str = f"Parallel Worker {stage.name}" self.queue_with_cpu: bool = parallel_config.queue_with_cpu self.timeout: int = parallel_config.timeout + self._lifecycle_lock = threading.Lock() + self._failed = False + self._closed = False + self._failure_reason: str | None = None # Use spawn to start processes regardless of world_size current_method = mp.get_start_method(allow_none=True) @@ -185,6 +198,64 @@ def __init__( join=False, ) + @property + def failed(self) -> bool: + """Return whether this worker group encountered an unrecoverable failure.""" + return self._failed + + @property + def closed(self) -> bool: + """Return whether this worker group has been deterministically closed.""" + return self._closed + + @property + def failure_reason(self) -> str | None: + """Return the first failure that made this worker group unusable.""" + return self._failure_reason + + def _ensure_usable(self) -> None: + if self._closed: + raise RuntimeError(f"ParallelWorker:{self.name} is closed") + if self._failed: + raise RuntimeError(f"ParallelWorker:{self.name} has failed: {self._failure_reason}") + + def _terminate_processes(self) -> None: + if not hasattr(self, "ctx"): + return + for process in self.ctx.processes: + try: + if process.is_alive(): + process.kill() + process.join(timeout=2) + except Exception as exc: + logger.warning(f"Failed to terminate {self.name} process: {exc}") + + def _mark_failed(self, reason: str) -> None: + with self._lifecycle_lock: + if self._failed: + return + self._failed = True + self._failure_reason = reason + logger.error(f"ParallelWorker:{self.name} marked failed: {reason}") + self._terminate_processes() + + def _wait_result(self, method_name: str) -> Any: + try: + result = self.queue_out.get(timeout=self.timeout) + except Empty as exc: + reason = f"{method_name} timeout after {self.timeout} seconds" + self._mark_failed(reason) + raise RuntimeError(f"ParallelWorker:{self.name} {reason}") from exc + except Exception as exc: + reason = f"{method_name} result queue failed: {exc}" + self._mark_failed(reason) + raise RuntimeError(f"ParallelWorker:{self.name} {reason}") from exc + if isinstance(result, Exception): + reason = f"{method_name} failed: {result}" + self._mark_failed(reason) + raise RuntimeError(f"ParallelWorker:{self.name} {reason}") from result + return result + def enable_metrics(self, registry: Any | None = None) -> None: """Enable metrics collection on the wrapped stage. @@ -210,6 +281,7 @@ def _metrics_hook(self, value: StageMetricContext | None) -> None: def put_data(self, data: Any) -> None: """Send data to all worker processes.""" + self._ensure_usable() if self.queue_with_cpu: data = to_device(data, "cpu") for i, q in enumerate(self.queue_in): @@ -218,22 +290,13 @@ def put_data(self, data: Any) -> None: def __call__(self, *args: Any, **kwargs: Any) -> Any | Callable[[], Any]: """Submit __call__ task to all workers.""" + self._ensure_usable() sync = kwargs.pop("sync", False) data = ["__call__", args, kwargs] self.put_data(data) def wait() -> Any: - try: - res = self.queue_out.get(timeout=self.timeout) - except Empty: - logger.error(f"ParallelWorker:{self.name} __call__ timeout") - raise RuntimeError(f"ParallelWorker:{self.name} __call__ timeout") - except Exception as e: - logger.error(f"ParallelWorker:{self.name} __call__ error: {e}") - raise RuntimeError(f"ParallelWorker:{self.name} __call__ error: {e}") - if isinstance(res, Exception): - raise res - return res + return self._wait_result("__call__") if sync: return wait() @@ -244,6 +307,7 @@ def __getattr__(self, name: str) -> Callable[..., Any]: """Submit arbitrary method call to all workers.""" def wrapped_func(*args: Any, **kwargs: Any) -> Any | Callable[[], Any]: + self._ensure_usable() sync = kwargs.pop("sync", False) data = [name, args, kwargs] self.put_data(data) @@ -254,23 +318,17 @@ def wrapped_func(*args: Any, **kwargs: Any) -> Any | Callable[[], Any]: def wait() -> Any: start_time = time.perf_counter() + success = False try: - res = self.queue_out.get(timeout=self.timeout) - except Empty: - logger.error(f"ParallelWorker:{self.name} {name} timeout") - raise RuntimeError(f"ParallelWorker:{self.name} {name} timeout") - except Exception as e: - logger.error(f"ParallelWorker:{self.name} {name} error: {e}") - raise RuntimeError(f"ParallelWorker:{self.name} {name} error: {e}") + result = self._wait_result(name) + success = True + logger.info(f"ParallelWorker:{self.name} {name} done") + return result finally: if hook is not None: duration = time.perf_counter() - start_time - hook.record_execution(duration, success=True) + hook.record_execution(duration, success=success) hook.exit() - if isinstance(res, Exception): - raise res - logger.info(f"ParallelWorker:{self.name} {name} done") - return res if sync: return wait() @@ -279,14 +337,41 @@ def wait() -> Any: return wrapped_func - def __del__(self) -> None: - """Cleanup worker processes on deletion.""" + def close(self) -> None: + """Idempotently stop all ranks and close multiprocessing queues.""" + with self._lifecycle_lock: + if self._closed: + return + self._closed = True + if hasattr(self, "ctx"): - self.put_data(["exit", None, None]) - for p in self.ctx.processes: - p.join(timeout=10) - if p.is_alive(): - p.kill() - for q in self.queue_in: - q.close() - self.queue_out.close() + if not self._failed: + for queue in self.queue_in: + try: + queue.put(["exit", None, None]) + except Exception as exc: + logger.warning(f"Failed to send exit to {self.name}: {exc}") + for process in self.ctx.processes: + try: + process.join(timeout=10) + if process.is_alive(): + process.kill() + process.join(timeout=2) + except Exception as exc: + logger.warning(f"Failed to join {self.name} process: {exc}") + for queue in self.queue_in: + try: + queue.close() + except Exception: + pass + try: + self.queue_out.close() + except Exception: + pass + + def __del__(self) -> None: + """Best-effort fallback; callers should use close explicitly.""" + try: + self.close() + except Exception: + pass diff --git a/tests/unit/models/test_lingbot_world_fast_dit.py b/tests/unit/models/test_lingbot_world_fast_dit.py new file mode 100644 index 0000000..6198884 --- /dev/null +++ b/tests/unit/models/test_lingbot_world_fast_dit.py @@ -0,0 +1,97 @@ +from unittest.mock import patch + +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.models.lingbot_world_fast_dit import ( + CachedCrossAttention, + CausalSelfAttention, + LingBotWorldFastDiT, +) +from telefuser.models.wan_video_dit import precompute_freqs_cis_3d + + +def test_causal_self_attention_uses_unified_attention() -> None: + attention = CausalSelfAttention(dim=32, num_heads=4) + attention_config = AttentionConfig.dense_attention(AttnImplType.SAGE_ATTN_2_8_8_SM90) + attention.attention_config = attention_config + + freqs = precompute_freqs_cis_3d(8) + freqs_cos = torch.cat([freq.real for freq in freqs], dim=-1) + freqs_sin = torch.cat([freq.imag for freq in freqs], dim=-1) + cache = { + "k": torch.zeros(1, 12, 4, 8), + "v": torch.zeros(1, 12, 4, 8), + "global_end_index": 0, + "local_end_index": 0, + } + captured: dict[str, object] = {} + + def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: object) -> torch.Tensor: + captured.update(q_shape=q.shape, k_shape=k.shape, v_shape=v.shape, **kwargs) + return q + + with patch("telefuser.models.lingbot_world_fast_dit.attn_func", side_effect=fake_attention): + output = attention( + torch.randn(1, 6, 32), + freqs_cos, + freqs_sin, + (1, 2, 3), + cache, + current_start=0, + max_attention_size=12, + ) + + assert output.shape == (1, 6, 32) + assert captured["q_shape"] == torch.Size([1, 6, 4, 8]) + assert captured["k_shape"] == torch.Size([1, 6, 4, 8]) + assert captured["v_shape"] == torch.Size([1, 6, 4, 8]) + assert captured["attention_config"] is attention_config + assert captured["input_layout"] == "BSND" + assert captured["output_layout"] == "BSND" + + +def test_cached_cross_attention_uses_unified_attention_and_bsnd_cache() -> None: + attention = CachedCrossAttention(dim=32, num_heads=4) + attention_config = AttentionConfig.dense_attention(AttnImplType.SAGE_ATTN_2_8_8_SM90) + attention.attention_config = attention_config + cache: dict[str, torch.Tensor | bool] = {"is_init": False} + calls: list[tuple[torch.Size, torch.Size, AttentionConfig]] = [] + + def fake_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, **kwargs: object) -> torch.Tensor: + assert k.shape == v.shape + calls.append((q.shape, k.shape, kwargs["attention_config"])) + return q + + with patch("telefuser.models.lingbot_world_fast_dit.attn_func", side_effect=fake_attention): + output = attention(torch.randn(1, 6, 32), torch.randn(1, 5, 32), cache) + cached_output = attention(torch.randn(1, 6, 32), torch.randn(1, 5, 32), cache) + + assert output.shape == cached_output.shape == (1, 6, 32) + assert calls == [ + (torch.Size([1, 6, 4, 8]), torch.Size([1, 5, 4, 8]), attention_config), + (torch.Size([1, 6, 4, 8]), torch.Size([1, 5, 4, 8]), attention_config), + ] + assert cache["k"].shape == torch.Size([1, 5, 4, 8]) + assert cache["v"].shape == torch.Size([1, 5, 4, 8]) + assert cache["is_init"] is True + + +def test_set_attention_config_updates_all_blocks() -> None: + model = LingBotWorldFastDiT( + in_dim=4, + dim=32, + ffn_dim=64, + freq_dim=8, + text_dim=16, + out_dim=4, + num_heads=4, + num_layers=2, + ) + attention_config = AttentionConfig.dense_attention(AttnImplType.SAGE_ATTN_2_8_8_SM90) + + model.set_attention_config(attention_config) + + for block in model.blocks: + assert block.self_attn.attention_config is attention_config + assert block.cross_attn.attention_config is attention_config diff --git a/tests/unit/ops/test_attention_backends.py b/tests/unit/ops/test_attention_backends.py new file mode 100644 index 0000000..56155b5 --- /dev/null +++ b/tests/unit/ops/test_attention_backends.py @@ -0,0 +1,31 @@ +from types import ModuleType +from unittest.mock import patch + +from telefuser.ops.attention import backends + + +def test_sage_attention_prefers_tf_kernel() -> None: + imported_modules: list[str] = [] + tf_kernel_module = ModuleType("tf_kernel.sageattn2") + previous_available = backends.SAGE_ATTN_AVAILABLE + previous_backend = backends.sageattention + + def import_module(name: str) -> ModuleType: + imported_modules.append(name) + return tf_kernel_module + + try: + backends.SAGE_ATTN_AVAILABLE = False + backends.sageattention = None + with ( + patch("telefuser.ops.attention.backends.importlib.util.find_spec", return_value=object()), + patch("telefuser.ops.attention.backends.importlib.import_module", side_effect=import_module), + ): + backends._try_import_sage_attn() + + assert imported_modules == ["tf_kernel.sageattn2"] + assert backends.SAGE_ATTN_AVAILABLE is True + assert backends.sageattention is tf_kernel_module + finally: + backends.SAGE_ATTN_AVAILABLE = previous_available + backends.sageattention = previous_backend diff --git a/tests/unit/pipelines/lingbot_world_fast/test_control_alignment.py b/tests/unit/pipelines/lingbot_world_fast/test_control_alignment.py index 14809a5..54f5121 100644 --- a/tests/unit/pipelines/lingbot_world_fast/test_control_alignment.py +++ b/tests/unit/pipelines/lingbot_world_fast/test_control_alignment.py @@ -1,62 +1,210 @@ import numpy as np +import pytest +import torch -from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline +from telefuser.pipelines.lingbot_world_fast.control import ( + LingBotWorldFastControlBuilder, + LingBotWorldFastControlContext, + LingBotWorldFastOfflineControlSource, + build_camera_control_chunk, + compute_relative_poses, + get_ks_transformed, + interpolate_camera_poses, + truncate_control_sequence, +) -def test_control_inputs_are_truncated_to_available_4n_plus_1_frames() -> None: - poses = np.zeros((10, 4, 4), dtype=np.float32) - intrinsics = np.zeros((10, 4), dtype=np.float32) - action = np.zeros((10, 6), dtype=np.float32) - - trimmed_poses, trimmed_intrinsics, trimmed_action, frame_num = ( - LingBotWorldFastPipeline._truncate_control_inputs_to_frame_num( - poses=poses, - intrinsics=intrinsics, - action=action, - frame_num=81, +def _builder() -> LingBotWorldFastControlBuilder: + return LingBotWorldFastControlBuilder( + LingBotWorldFastControlContext( + control_type="act", + device="cpu", + control_dtype=torch.float32, + orig_height=8, + orig_width=8, + height=8, + width=8, + latent_h=1, + latent_w=1, + latent_frames=3, + chunk_size=3, + intrinsics=torch.tensor([8.0, 8.0, 4.0, 4.0]), ) ) - assert frame_num == 9 - assert trimmed_poses.shape == (9, 4, 4) - assert trimmed_intrinsics.shape == (9, 4) - assert trimmed_action.shape == (9, 6) +def test_action_alignment_samples_external_video_rate_actions() -> None: + action = torch.arange(9, dtype=torch.float32).unsqueeze(1) -def test_scalar_intrinsics_are_preserved_when_truncating_controls() -> None: - poses = np.zeros((10, 4, 4), dtype=np.float32) - intrinsics = np.zeros((4,), dtype=np.float32) + aligned = _builder()._align_action_frames(action, target_frames=3) - trimmed_poses, trimmed_intrinsics, trimmed_action, frame_num = ( - LingBotWorldFastPipeline._truncate_control_inputs_to_frame_num( - poses=poses, - intrinsics=intrinsics, - action=None, - frame_num=5, - ) + torch.testing.assert_close(aligned[:, 0], torch.tensor([0.0, 4.0, 8.0])) + + +def test_action_alignment_rejects_an_incomplete_chunk_action() -> None: + with pytest.raises(ValueError, match="must be"): + _builder()._align_action_frames(torch.zeros(2, 4), target_frames=3) + + +def test_prebuilt_control_is_the_only_pipeline_input() -> None: + control = torch.ones(1, 12, 3, 1, 1) + + built = _builder().build({"control_tensor": control}) + + assert built is control + + +def test_offline_control_source_materializes_once_when_first_requested() -> None: + builder = _builder() + calls = 0 + + def build_sequence(*args): + nonlocal calls + calls += 1 + return [torch.tensor([1]), torch.tensor([2])] + + builder.build_sequence = build_sequence + source = LingBotWorldFastOfflineControlSource(builder, poses=object(), intrinsics=object()) + + second = source.control_at(1) + first = source.control_at(0) + + assert calls == 0 + assert torch.equal(first(), torch.tensor([1])) + assert calls == 1 + assert torch.equal(second(), torch.tensor([2])) + assert calls == 1 + + +def test_external_builder_matches_legacy_offline_camera_control_math() -> None: + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 9, axis=0) + poses[:, 2, 3] = np.linspace(0.0, 1.0, len(poses)) + intrinsics = np.repeat(np.array([[8.0, 8.0, 4.0, 4.0]], dtype=np.float32), len(poses), axis=0) + context = LingBotWorldFastControlContext( + control_type="cam", + device="cpu", + control_dtype=torch.float32, + orig_height=8, + orig_width=8, + height=8, + width=8, + latent_h=1, + latent_w=1, + latent_frames=3, + chunk_size=3, + intrinsics=torch.tensor([8.0, 8.0, 4.0, 4.0]), ) + controls = torch.cat(LingBotWorldFastControlBuilder(context).build_sequence(poses, intrinsics), dim=2) - assert frame_num == 5 - assert trimmed_poses.shape == (5, 4, 4) - assert trimmed_intrinsics.shape == (4,) - assert trimmed_action is None + poses_t = torch.as_tensor(poses, dtype=torch.float32) + intrinsics_t = get_ks_transformed( + torch.as_tensor(intrinsics, dtype=torch.float32), + height_org=context.orig_height, + width_org=context.orig_width, + height_resize=context.height, + width_resize=context.width, + height_final=context.height, + width_final=context.width, + ) + interpolated = interpolate_camera_poses( + src_indices=np.linspace(0, len(poses_t) - 1, len(poses_t)), + src_rot_mat=np.asarray(poses_t[:, :3, :3]), + src_trans_vec=np.asarray(poses_t[:, :3, 3]), + tgt_indices=np.linspace(0, len(poses_t) - 1, context.latent_frames), + ) + relative = compute_relative_poses(interpolated, framewise=True) + legacy = build_camera_control_chunk( + relative, + intrinsics_t[0].repeat(len(relative), 1), + context.latent_h, + context.latent_w, + context.height, + context.width, + ) + assert torch.equal(controls, legacy) -def test_requested_frame_count_caps_longer_controls() -> None: - poses = np.zeros((269, 4, 4), dtype=np.float32) - intrinsics = np.zeros((269, 4), dtype=np.float32) - action = np.zeros((269, 6), dtype=np.float32) - trimmed_poses, trimmed_intrinsics, trimmed_action, frame_num = ( - LingBotWorldFastPipeline._truncate_control_inputs_to_frame_num( - poses=poses, - intrinsics=intrinsics, - action=action, - frame_num=81, - ) +def test_offline_control_window_keeps_first_intrinsics_and_rejects_short_action() -> None: + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 9, axis=0) + intrinsics = np.arange(32, dtype=np.float32).reshape(8, 4) + + with pytest.raises(ValueError, match="Action sequence"): + truncate_control_sequence(poses, np.ones(4, dtype=np.float32), np.ones((8, 4), dtype=np.float32), frame_num=9) + + _, fixed_intrinsics, _ = truncate_control_sequence(poses, intrinsics, None, frame_num=9) + np.testing.assert_array_equal(fixed_intrinsics, intrinsics[0]) + + +def test_offline_control_window_rejects_empty_intrinsics() -> None: + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 9, axis=0) + + with pytest.raises(ValueError, match="at least one row"): + truncate_control_sequence(poses, np.empty((0, 4), dtype=np.float32), None, frame_num=9) + + +def test_action_mode_requires_actions_and_camera_mode_fixes_per_frame_intrinsics() -> None: + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 3, axis=0) + intrinsics = np.repeat(np.array([[8.0, 8.0, 4.0, 4.0]], dtype=np.float32), 3, axis=0) + + with pytest.raises(ValueError, match="requires an action"): + _builder().build({"poses": poses, "intrinsics": intrinsics}) + + camera_context = LingBotWorldFastControlContext( + control_type="cam", + device="cpu", + control_dtype=torch.float32, + orig_height=8, + orig_width=8, + height=8, + width=8, + latent_h=1, + latent_w=1, + latent_frames=3, + chunk_size=3, + intrinsics=torch.tensor([8.0, 8.0, 4.0, 4.0]), + ) + control = LingBotWorldFastControlBuilder(camera_context).build({"poses": poses, "intrinsics": intrinsics}) + assert control.shape == (1, 384, 3, 1, 1) + + resampled = LingBotWorldFastControlBuilder._resample_intrinsics( + torch.tensor([[8.0, 8.0, 4.0, 4.0], [9.0, 9.0, 5.0, 5.0]]), + target_frames=3, ) + torch.testing.assert_close(resampled, torch.tensor([[8.0, 8.0, 4.0, 4.0]]).repeat(3, 1)) + + +def test_online_builder_uses_session_intrinsics_and_preserves_chunk_boundary_delta() -> None: + context = LingBotWorldFastControlContext( + control_type="cam", + device="cpu", + control_dtype=torch.float32, + orig_height=8, + orig_width=8, + height=8, + width=8, + latent_h=1, + latent_w=1, + latent_frames=6, + chunk_size=3, + intrinsics=torch.tensor([8.0, 8.0, 4.0, 4.0]), + ) + builder = LingBotWorldFastControlBuilder(context) + captured_intrinsics: list[torch.Tensor] = [] + + def capture(poses: torch.Tensor, intrinsics: torch.Tensor, action: object | None) -> torch.Tensor: + captured_intrinsics.append(intrinsics) + return poses + + builder._build_tensor = capture + poses = np.repeat(np.eye(4, dtype=np.float32)[None], 3, axis=0) + poses[:, 2, 3] = [0.3, 0.4, 0.5] + previous_pose = np.eye(4, dtype=np.float32) + previous_pose[2, 3] = 0.2 + + relative = builder.build({"poses": poses, "previous_pose": previous_pose}) - assert frame_num == 81 - assert trimmed_poses.shape == (81, 4, 4) - assert trimmed_intrinsics.shape == (81, 4) - assert trimmed_action.shape == (81, 6) + assert not torch.equal(relative[0], torch.eye(4)) + torch.testing.assert_close(relative[:, 2, 3], torch.ones(3)) + assert len(captured_intrinsics) == 1 + torch.testing.assert_close(captured_intrinsics[0], context.intrinsics.repeat(3, 1)) diff --git a/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py b/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py new file mode 100644 index 0000000..2cb47dd --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_parallelism.py @@ -0,0 +1,76 @@ +from unittest.mock import MagicMock, patch + +import torch + +from telefuser.core.config import AttentionConfig, ModelRuntimeConfig, ParallelConfig +from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage + + +def test_denoising_stage_parallel_models_enables_ulysses_and_fsdp() -> None: + dit = MagicMock() + dit.get_fsdp_module_names.return_value = ["blocks"] + parallel_config = ParallelConfig( + device_ids=[0, 1, 2, 3], + sp_ulysses_degree=4, + enable_fsdp=True, + ) + runtime_config = ModelRuntimeConfig( + device_type="cuda", + device_id=0, + torch_dtype=torch.bfloat16, + attention_config=AttentionConfig(), + parallel_config=parallel_config, + ) + stage = LingBotWorldFastDenoisingStage("denoise", dit, runtime_config) + device_mesh = MagicMock() + fsdp_model = MagicMock() + + with ( + patch( + "telefuser.pipelines.lingbot_world_fast.denoising.create_device_mesh_from_config", + return_value=device_mesh, + ) as create_mesh, + patch( + "telefuser.pipelines.lingbot_world_fast.denoising.shard_model", + return_value=fsdp_model, + ) as shard, + ): + stage.parallel_models() + + create_mesh.assert_called_once_with(parallel_config) + dit.enable_usp.assert_called_once_with(device_mesh) + shard.assert_called_once_with( + module=dit, + device_id=stage.device, + wrap_module_names=["blocks"], + param_dtype=torch.bfloat16, + reduce_dtype=torch.bfloat16, + buffer_dtype=torch.bfloat16, + ) + assert stage.dit is fsdp_model + assert stage.onload_models_flag is True + + +def test_denoising_stage_parallel_models_without_fsdp_keeps_model() -> None: + dit = MagicMock() + parallel_config = ParallelConfig(device_ids=[0, 1], sp_ulysses_degree=2) + runtime_config = ModelRuntimeConfig( + device_type="cuda", + device_id=0, + parallel_config=parallel_config, + ) + stage = LingBotWorldFastDenoisingStage("denoise", dit, runtime_config) + device_mesh = MagicMock() + + with ( + patch( + "telefuser.pipelines.lingbot_world_fast.denoising.create_device_mesh_from_config", + return_value=device_mesh, + ), + patch("telefuser.pipelines.lingbot_world_fast.denoising.shard_model") as shard, + ): + stage.parallel_models() + + dit.enable_usp.assert_called_once_with(device_mesh) + shard.assert_not_called() + assert stage.dit is dit diff --git a/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py new file mode 100644 index 0000000..a97e73b --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_pipeline_call.py @@ -0,0 +1,299 @@ +import threading +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_world_fast.pipeline import LingBotWorldFastPipeline +from telefuser.pipelines.lingbot_world_fast.session import ( + LingBotWorldFastChunkRequest, + LingBotWorldFastGenerationSession, + LingBotWorldFastSessionConfig, + LingBotWorldFastSessionStatus, +) +from telefuser.worker.parallel_worker import ParallelWorker + + +def _session( + *, + active: bool = True, + 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, + 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, + height=8, + width=8, + frame_tokens=1, + chunk_size=1, + max_attention_size=1, + cache_handle=7, + active=active, + current_chunk_index=current_chunk_index, + ) + + +def _control() -> torch.Tensor: + return torch.zeros(1, 384, 1, 1, 1, dtype=torch.float32) + + +def _pipeline() -> LingBotWorldFastPipeline: + pipeline = LingBotWorldFastPipeline(device="cpu") + pipeline.config = SimpleNamespace(control_type="cam") + return pipeline + + +def test_pipeline_call_generates_one_explicitly_indexed_chunk() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + runtime = _session() + expected = [Image.new("RGB", (8, 8), "red") for _ in range(9)] + + def generate_next_chunk(runtime_state, control=None, progress_callback=None): + runtime_state.current_chunk_index += 1 + runtime_state.emitted_frames += len(expected) + return expected + + pipeline.generate_next_chunk = MagicMock(side_effect=generate_next_chunk) + progress_callback = MagicMock() + control = _control() + + result = pipeline( + runtime, + LingBotWorldFastChunkRequest( + chunk_index=0, + session_id="session-a", + control=control, + ), + progress_callback=progress_callback, + ) + + assert result.frames == expected + assert result.chunk_index == 0 + assert result.session_id == "session-a" + assert result.emitted_frames == 9 + assert result.done is False + assert runtime.status == LingBotWorldFastSessionStatus.COMMITTED + pipeline.generate_next_chunk.assert_called_once_with( + runtime, + control=control, + progress_callback=progress_callback, + ) + + +def test_chunk_request_rejects_invalid_inputs() -> None: + with pytest.raises(ValueError, match="non-negative"): + LingBotWorldFastChunkRequest(chunk_index=-1, control=torch.zeros(1)) + + with pytest.raises(TypeError, match="model control tensor"): + LingBotWorldFastChunkRequest(chunk_index=0, control=object()) + + +def test_pipeline_call_rejects_inactive_runtime() -> None: + pipeline = _pipeline() + runtime = _session(active=False) + pipeline.generate_next_chunk = MagicMock() + + with pytest.raises(RuntimeError, match="inactive"): + pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=torch.zeros(1))) + + pipeline.generate_next_chunk.assert_not_called() + + +def test_pipeline_call_rejects_out_of_order_chunk() -> None: + pipeline = _pipeline() + runtime = _session(current_chunk_index=1) + pipeline.generate_next_chunk = MagicMock() + + with pytest.raises(ValueError, match="does not match session index"): + pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=torch.zeros(1))) + + pipeline.generate_next_chunk.assert_not_called() + + +def test_new_session_rejects_out_of_order_chunk_without_initializing() -> None: + pipeline = _pipeline() + session = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))) + ) + deferred_control = MagicMock(return_value=_control()) + pipeline._create_initialized_session = MagicMock() + + with pytest.raises(ValueError, match="does not match session index"): + pipeline(session, LingBotWorldFastChunkRequest(chunk_index=1, control=deferred_control)) + + deferred_control.assert_not_called() + pipeline._create_initialized_session.assert_not_called() + + +def test_first_pipeline_call_initializes_the_external_session() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + initialized = _session() + session = LingBotWorldFastGenerationSession( + config=LingBotWorldFastSessionConfig(prompt="test", image=Image.new("RGB", (8, 8))) + ) + control = _control() + events: list[str] = [] + + def initialize(config, progress_callback, before_cache): + assert config is session.config + assert progress_callback is None + events.append("initialization_started") + before_cache() + events.append("cache_initialization_ready") + return initialized + + def deferred_control() -> torch.Tensor: + events.append("control_materialized") + return control + + pipeline.generate_next_chunk = MagicMock(return_value=[]) + pipeline._create_initialized_session = MagicMock(side_effect=initialize) + + pipeline(session, LingBotWorldFastChunkRequest(chunk_index=0, control=deferred_control)) + + pipeline._create_initialized_session.assert_called_once() + assert events == ["initialization_started", "control_materialized", "cache_initialization_ready"] + pipeline.generate_next_chunk.assert_called_once_with(session, control=control, progress_callback=None) + + +def test_pipeline_call_releases_runtime_when_generation_fails() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + runtime = _session(chunk_count=1) + pipeline.generate_next_chunk = MagicMock(side_effect=RuntimeError("generation failed")) + + 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) + + with pytest.raises(RuntimeError, match="poisoned"): + pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=_control())) + + +def test_pipeline_call_rejects_control_with_wrong_shape() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + runtime = _session() + + with pytest.raises(ValueError, match="Control shape"): + pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=torch.zeros(1, dtype=torch.float32))) + + pipeline.denoise_stage.release_cache.assert_not_called() + + +def test_final_chunk_releases_decoder_state_and_cache() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + runtime = _session(chunk_count=1) + runtime.decoder_state.feat_cache = [torch.ones(1)] + runtime.decoder_state.feat_idx = [1] + + 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) + + result = pipeline(runtime, LingBotWorldFastChunkRequest(chunk_index=0, control=_control())) + + assert result.done is True + assert runtime.cache_handle is None + assert runtime.decoder_state.feat_cache == [] + assert runtime.decoder_state.feat_idx == [0] + assert runtime.status == LingBotWorldFastSessionStatus.RELEASED + pipeline.denoise_stage.release_cache.assert_called_once_with(7) + + +def test_release_session_is_idempotent() -> None: + pipeline = _pipeline() + pipeline.denoise_stage = MagicMock() + session = _session() + + pipeline.release_session(session) + pipeline.release_session(session) + + pipeline.denoise_stage.release_cache.assert_called_once_with(7) + assert session.cache_handle is None + assert session.status == LingBotWorldFastSessionStatus.RELEASED + + +def test_concurrent_chunk_on_same_session_is_rejected() -> None: + pipeline = _pipeline() + session = _session() + lock_acquired = threading.Event() + release_lock = threading.Event() + + def hold_transaction() -> None: + with session.transaction_lock: + lock_acquired.set() + release_lock.wait(timeout=2.0) + + holder = threading.Thread(target=hold_transaction, daemon=True) + holder.start() + assert lock_acquired.wait(timeout=1.0) + + try: + with pytest.raises(RuntimeError, match="already has a chunk in progress"): + pipeline( + session, + LingBotWorldFastChunkRequest( + chunk_index=0, + control=torch.zeros(1), + ), + ) + finally: + release_lock.set() + holder.join(timeout=1.0) + + assert not holder.is_alive() + assert session.status == LingBotWorldFastSessionStatus.READY + + +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)) + 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]) + + config = LingBotWorldFastSessionConfig(prompt="test", image=frame) + + with patch.object(LingBotWorldFastPipeline, "__call__", side_effect=generate) as generate_chunk: + frames = pipeline.generate_video(config, controls=[torch.tensor([1]), torch.tensor([2])]) + + assert frames == [frame, frame] + assert generate_chunk.call_count == 2 + pipeline.release_session.assert_called_once() + + +def test_pipeline_close_delegates_to_parallel_worker() -> None: + pipeline = LingBotWorldFastPipeline(device="cpu") + worker = object.__new__(ParallelWorker) + worker.close = MagicMock() + pipeline.denoise_stage = worker + + pipeline.close() + + worker.close.assert_called_once_with() diff --git a/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py new file mode 100644 index 0000000..aafe92b --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_runtime_baseline.py @@ -0,0 +1,218 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +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.session import LingBotWorldFastSessionConfig + + +def _build_runtime_pipeline() -> LingBotWorldFastPipeline: + pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32) + pipeline.text_device = torch.device("cpu") + pipeline.vae_device = torch.device("cpu") + pipeline.config = SimpleNamespace( + control_type="cam", + max_area=16 * 16, + orig_height=16, + orig_width=16, + local_attn_size=-1, + sink_size=0, + ) + pipeline.dit = SimpleNamespace( + patch_size=(1, 2, 2), + dim=8, + num_heads=2, + num_layers=1, + ) + 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 + + +def _create_runtime(frame_num: int, seed: int = 42): + pipeline = _build_runtime_pipeline() + runtime = pipeline._create_initialized_session( + LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + frame_num=frame_num, + chunk_size=3, + seed=seed, + ) + ) + return pipeline, runtime + + +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.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: + pipeline = _build_runtime_pipeline() + config = LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + frame_num=9, + chunk_size=3, + ) + + first = pipeline._create_initialized_session(config) + second = pipeline._create_initialized_session(config) + + assert first.cache_handle == 0 + assert second.cache_handle == 1 + assert first.decoder_state is not second.decoder_state + assert pipeline.denoise_stage.initialize_cache.call_count == 2 + + +def test_cache_initialization_failure_triggers_global_cleanup() -> None: + pipeline = _build_runtime_pipeline() + pipeline.denoise_stage.initialize_cache.side_effect = RuntimeError("rank initialization failed") + + with pytest.raises(RuntimeError, match="rank initialization failed"): + pipeline._create_initialized_session( + LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + frame_num=9, + ) + ) + + pipeline.denoise_stage.release_cache.assert_called_once_with(0) + + +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_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) + + torch.testing.assert_close(first_noise, repeated_noise) + assert not torch.equal(first_noise, different_noise) + + +def test_denoising_generator_state_advances_between_chunks() -> None: + class Scheduler: + sigmas = torch.tensor([1.0, 0.0]) + timesteps = torch.tensor([10.0, 0.0]) + + @staticmethod + def add_noise(x0: torch.Tensor, noise: torch.Tensor, _timestep: torch.Tensor) -> torch.Tensor: + return x0 + noise + + stage = LingBotWorldFastDenoisingStage.__new__(LingBotWorldFastDenoisingStage) + stage.torch_dtype = torch.float32 + stage.dit = MagicMock(side_effect=lambda **kwargs: torch.zeros_like(kwargs["x"])) + timesteps = torch.tensor([10.0, 0.0]) + latent = torch.zeros(1, 1, 1, 1, 1) + generator = torch.Generator(device="cpu").manual_seed(123) + + def denoise(active_generator: torch.Generator) -> torch.Tensor: + return stage.denoise_chunk( + latent_chunk=latent, + condition_chunk=latent, + prompt_emb=torch.zeros(1, 1, 1), + timesteps=timesteps, + scheduler=Scheduler(), + control_chunk=None, + self_kv_cache=[], + crossattn_cache=[], + current_start=0, + max_attention_size=1, + generator=active_generator, + ) + + first = denoise(generator) + second = denoise(generator) + repeated = denoise(torch.Generator(device="cpu").manual_seed(123)) + + torch.testing.assert_close(first, repeated) + assert not torch.equal(first, second) + + +def test_final_chunk_marks_runtime_inactive_and_releases_denoise_state() -> None: + pipeline = LingBotWorldFastPipeline(device="cpu", torch_dtype=torch.float32) + pipeline.vae_device = torch.device("cpu") + pipeline.denoise_stage = object() + pipeline.decode_video_cached = MagicMock(return_value=torch.zeros(1)) + 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) + runtime = SimpleNamespace( + current_chunk_index=0, + noise_chunks=[torch.zeros_like(cached_latent)], + condition_chunks=[torch.zeros_like(cached_latent)], + 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, + ) + + frames = pipeline.generate_next_chunk(runtime, control=torch.zeros(1)) + + assert frames == expected_frames + assert runtime.current_chunk_index == 1 + assert runtime.emitted_frames == 9 + assert runtime.active is False + + +def test_runtime_rejects_non_aligned_frame_count() -> None: + with pytest.raises(ValueError, match="frame_num"): + _create_runtime(frame_num=13) + + +def test_runtime_rejects_frame_count_smaller_than_first_chunk() -> None: + with pytest.raises(ValueError, match="frame_num"): + _create_runtime(frame_num=5) + + +def test_runtime_rejects_non_positive_chunk_size() -> None: + pipeline = _build_runtime_pipeline() + with pytest.raises(ValueError, match="chunk_size"): + pipeline.control_context( + LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + chunk_size=0, + frame_num=9, + ) + ) + + +def test_control_mode_must_match_the_initialized_pipeline() -> None: + pipeline = _build_runtime_pipeline() + with pytest.raises(ValueError, match="does not match"): + pipeline.control_context( + LingBotWorldFastSessionConfig( + prompt="baseline", + image=Image.new("RGB", (16, 16)), + control_mode="act", + frame_num=9, + ) + ) 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 new file mode 100644 index 0000000..28ec427 --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_service_action_loop.py @@ -0,0 +1,224 @@ +import asyncio +import threading +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import numpy as np +import pytest +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_world_fast.service import LingBotWorldFastService +from telefuser.pipelines.lingbot_world_fast.session import ( + LingBotWorldFastSessionConfig, + LingBotWorldFastSessionState, +) + + +def _state() -> LingBotWorldFastSessionState: + return LingBotWorldFastSessionState( + config=LingBotWorldFastSessionConfig( + prompt="test", + image=Image.new("RGB", (8, 8)), + frame_num=9, + ) + ) + + +def test_online_worker_waits_for_external_action_before_generating() -> None: + pipeline = MagicMock() + pipeline._best_output_size.return_value = (8, 8) + pipeline.check_resize_height_width.return_value = (8, 8) + service = LingBotWorldFastService(pipeline) + state = _state() + runtime_ready = threading.Event() + + def control_context(*args, **kwargs): + runtime_ready.set() + return SimpleNamespace( + control_type="cam", + chunk_size=3, + width=8, + height=8, + latent_frames=3, + ) + + 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))]) + + pipeline.control_context.side_effect = control_context + pipeline.side_effect = generate_chunk + builder = MagicMock() + control = torch.ones(1) + deferred_control = MagicMock(return_value=control) + builder.defer.return_value = deferred_control + + with patch("telefuser.pipelines.lingbot_world_fast.service.LingBotWorldFastControlBuilder", return_value=builder): + worker = threading.Thread( + target=service._run_worker_loop, + args=("session-a", state, MagicMock()), + daemon=True, + ) + worker.start() + + assert runtime_ready.wait(timeout=1.0) + assert pipeline.call_count == 0 + + state.pending_inputs.put({"type": "control", "control_tensor": control}) + worker.join(timeout=2.0) + + assert not worker.is_alive() + assert pipeline.call_count == 1 + builder.defer.assert_called_once_with({"type": "control", "control_tensor": control}) + pipeline.release_session.assert_called_once() + + +def test_direction_action_updates_state_and_wakes_worker() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + service._sessions["session-a"] = state + + service.push_chunk( + "session-a", + {"type": "control", "direction": "up", "event": "press"}, + ) + + assert state.pressed_controls == {"w"} + assert state.pending_inputs.get_nowait() == {"type": "direction_control"} + + +def test_release_stops_control_without_scheduling_stationary_generation() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + service._sessions["session-a"] = state + state.pressed_controls.add("w") + + service.push_chunk("session-a", {"type": "control", "key": "ArrowUp", "event": "release"}) + + assert state.pressed_controls == set() + assert state.pending_inputs.empty() + + +def test_directional_chunks_match_source_video_rate_integration_and_boundary() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + state.pressed_controls.add("w") + context = SimpleNamespace(control_type="cam", chunk_size=3) + + first = service._build_directional_control_chunk(state, context) + second = service._build_directional_control_chunk(state, context) + + assert first is not None + assert second is not None + assert "previous_pose" not in first + first_poses = np.asarray(first["poses"]) + second_poses = np.asarray(second["poses"]) + np.testing.assert_allclose(first_poses[:, 2, 3], [0.0, 0.2, 0.4], rtol=0, atol=1e-6) + np.testing.assert_allclose(np.asarray(second["previous_pose"])[2, 3], 0.4, rtol=0, atol=1e-6) + np.testing.assert_allclose(second_poses[:, 2, 3], [0.6, 0.8, 1.0], rtol=0, atol=1e-6) + + +def test_wasd_ijkl_and_arrow_aliases_have_distinct_translation_and_rotation_controls() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + context = SimpleNamespace(control_type="cam", chunk_size=3) + + state.pressed_controls.add("j") + yaw_chunk = service._build_directional_control_chunk(state, context) + + assert yaw_chunk is not None + expected_yaw = np.deg2rad(-16.0) + np.testing.assert_allclose(np.asarray(yaw_chunk["poses"])[-1, 0, 0], np.cos(expected_yaw), atol=1e-6) + assert service._direction_from_chunk({"key": "ArrowLeft"}) == "j" + assert service._direction_from_chunk({"key": "KeyA"}) == "a" + assert service._direction_from_chunk({"key": "KeyI"}) == "i" + + +def test_service_stop_closes_sessions_before_pipeline() -> None: + pipeline = MagicMock() + service = LingBotWorldFastService(pipeline) + service._sessions = {"session-a": MagicMock(), "session-b": MagicMock()} + service.close_session = MagicMock() + + service.stop() + + assert service.close_session.call_args_list == [ + (("session-a",),), + (("session-b",),), + ] + pipeline.close.assert_called_once_with() + + +def test_create_session_rejects_invalid_pipeline_configuration() -> None: + pipeline = MagicMock() + pipeline.control_context.side_effect = ValueError("invalid session configuration") + service = LingBotWorldFastService(pipeline) + + with pytest.raises(ValueError, match="invalid session"): + service.create_session({"image": Image.new("RGB", (8, 8))}) + + assert service._sessions == {} + + +def test_create_session_limits_stream_generation_to_20_seconds() -> None: + pipeline = MagicMock() + service = LingBotWorldFastService(pipeline) + + session_id = service.create_session( + { + "image": Image.new("RGB", (8, 8)), + "fps": 16, + "frame_num": 321, + } + ) + assert service._sessions[session_id].config.frame_num == 321 + service.close_session(session_id) + + with pytest.raises(ValueError, match="must not exceed 20 seconds"): + service.create_session( + { + "image": Image.new("RGB", (8, 8)), + "fps": 16, + "frame_num": 333, + } + ) + + +def test_create_session_initializes_fixed_intrinsics_from_action_path() -> None: + pipeline = MagicMock() + service = LingBotWorldFastService(pipeline) + intrinsics = np.asarray([[8.0, 8.0, 4.0, 4.0], [9.0, 9.0, 4.0, 4.0]]) + + with patch("telefuser.pipelines.lingbot_world_fast.service.np.load", return_value=intrinsics) as load: + session_id = service.create_session( + { + "image": Image.new("RGB", (8, 8)), + "action_path": "/controls", + } + ) + + load.assert_called_once_with(Path("/controls") / "intrinsics.npy") + session_config = pipeline.control_context.call_args.args[0] + assert session_config.intrinsics is intrinsics + assert service._sessions[session_id].control_context is pipeline.control_context.return_value + + +def test_pull_chunks_drains_terminal_messages_after_session_becomes_inactive() -> None: + service = LingBotWorldFastService(MagicMock()) + state = _state() + state.active = False + state.worker_thread = MagicMock() + state.worker_thread.is_alive.return_value = True + state.output_queue = asyncio.Queue() + state.output_queue.put_nowait({"type": "preview"}) + state.output_queue.put_nowait({"type": "error"}) + state.output_queue.put_nowait({"type": "done"}) + service._sessions["session-a"] = state + + async def collect() -> list[dict]: + return [chunk async for chunk in service.pull_chunks("session-a")] + + assert asyncio.run(collect()) == [{"type": "preview"}, {"type": "error"}] diff --git a/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py new file mode 100644 index 0000000..263b1cd --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_session_cache.py @@ -0,0 +1,115 @@ +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch +import torch.nn as nn + +from telefuser.models.wan_video_vae import WanVideoVAE, WanVideoVAEStreamingDecodeState +from telefuser.pipelines.lingbot_world_fast.denoising import LingBotWorldFastDenoisingStage + + +def _cache_stage() -> LingBotWorldFastDenoisingStage: + stage = LingBotWorldFastDenoisingStage.__new__(LingBotWorldFastDenoisingStage) + stage.device = torch.device("cpu") + stage._cache_registry = {} + stage._init_self_kv_cache = MagicMock(side_effect=lambda *_args: [{"owner": object()}]) + stage._init_crossattn_cache = MagicMock(side_effect=lambda *_args: [{"owner": object()}]) + return stage + + +def _initialize_cache(stage: LingBotWorldFastDenoisingStage, cache_handle: int) -> None: + generator_state = torch.Generator(device="cpu").manual_seed(cache_handle).get_state().tolist() + LingBotWorldFastDenoisingStage.initialize_cache.__wrapped__( + stage, + cache_handle=cache_handle, + batch_size=1, + kv_size=4, + max_sequence_length=8, + sample_shift=10.0, + generator_state=generator_state, + ) + + +def test_worker_cache_registry_isolates_handles_and_releases_idempotently() -> None: + stage = _cache_stage() + + _initialize_cache(stage, 11) + _initialize_cache(stage, 12) + + assert stage.list_cache_handles() == (11, 12) + assert stage.has_cache(11) + assert stage._cache_registry[11] is not stage._cache_registry[12] + assert stage._cache_registry[11].self_kv_cache is not stage._cache_registry[12].self_kv_cache + + with pytest.raises(ValueError, match="already registered"): + _initialize_cache(stage, 11) + + assert stage.release_cache(11) is True + assert stage.release_cache(11) is False + assert stage.list_cache_handles() == (12,) + + +def test_worker_rejects_unknown_cache_handle() -> None: + stage = _cache_stage() + latent = torch.zeros(1, 1, 1, 1, 1) + + with pytest.raises(KeyError, match="Unknown cache handle 99"): + LingBotWorldFastDenoisingStage.denoise_and_update_cache.__wrapped__( + stage, + cache_handle=99, + latent_chunk=latent, + condition_chunk=latent, + prompt_emb=torch.zeros(1, 1, 1), + control_chunk=None, + current_start=0, + max_attention_size=1, + ) + + +class _RecordingDecoder(nn.Module): + def __init__(self) -> None: + super().__init__() + self.cache_ids: list[int] = [] + + def forward(self, x, feat_cache, feat_idx): + self.cache_ids.append(id(feat_cache)) + feat_cache.append(float(x.flatten()[0])) + feat_idx[0] += 1 + return x + + +def test_vae_streaming_decode_state_is_session_scoped() -> None: + decoder = _RecordingDecoder() + vae = SimpleNamespace( + model=SimpleNamespace(conv2=lambda value: value, decoder=decoder), + scale=[0.0, 1.0], + z_dim=1, + _feat_cache=[], + _feat_idx=[0], + ) + first = WanVideoVAEStreamingDecodeState() + second = WanVideoVAEStreamingDecodeState() + + WanVideoVAE.cached_decode_withflag( + vae, + torch.ones(1, 1, 1, 1), + device=torch.device("cpu"), + is_first_clip=True, + is_last_clip=False, + decode_state=first, + ) + WanVideoVAE.cached_decode_withflag( + vae, + torch.full((1, 1, 1, 1), 2.0), + device=torch.device("cpu"), + is_first_clip=True, + is_last_clip=False, + decode_state=second, + ) + + assert first.feat_cache == [1.0] + assert second.feat_cache == [2.0] + assert first.feat_cache is not second.feat_cache + assert vae._feat_cache == [] + assert decoder.cache_ids == [id(first.feat_cache), id(second.feat_cache)] diff --git a/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py b/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py new file mode 100644 index 0000000..698b629 --- /dev/null +++ b/tests/unit/pipelines/lingbot_world_fast/test_stream_example.py @@ -0,0 +1,75 @@ +from unittest.mock import MagicMock, patch + +import torch + +from examples.lingbot import lingbot_world_fast_image_to_video_h100 as offline_example +from examples.lingbot import stream_lingbot_world_fast as stream_example +from examples.stream_server import webrtc_bidirectional_demo as webrtc_demo +from telefuser.core.config import AttnImplType +from telefuser.pipelines.lingbot_world_fast.service import MAX_GENERATION_SECONDS + + +def test_webrtc_demo_defaults_match_offline_h100_example() -> None: + assert webrtc_demo.DEFAULT_SAMPLE_SHIFT == offline_example.PPL_CONFIG["sample_shift"] + assert webrtc_demo.DEFAULT_PROMPT == offline_example.DEFAULT_PROMPT + assert webrtc_demo.MAX_GENERATION_SECONDS == MAX_GENERATION_SECONDS + + +def test_stream_get_pipeline_maps_ppl_config_to_internal_workers() -> None: + pipeline = MagicMock() + + with patch.object(stream_example, "LingBotWorldFastPipeline", return_value=pipeline) as pipeline_cls: + result = stream_example.get_pipeline( + parallelism=4, + model_root="/models/Wan2.2-I2V-A14B", + fast_model_root="/models/lingbot-world-fast", + ) + + assert result is pipeline + pipeline_cls.assert_called_once_with(device="cuda", torch_dtype=torch.bfloat16) + + assert len(pipeline.init.call_args.args) == 1 + config = pipeline.init.call_args.args[0] + assert config.checkpoint_dir == "/models/Wan2.2-I2V-A14B" + assert config.fast_checkpoint_path == "/models/lingbot-world-fast" + assert config.control_type == "cam" + assert config.max_area == 480 * 832 + assert config.local_attn_size == -1 + assert config.attention_config.attn_impl == AttnImplType.SAGE_ATTN_2_8_8_SM90 + assert config.parallel_config.device_ids == [0, 1, 2, 3] + assert config.parallel_config.sp_ulysses_degree == 4 + assert config.parallel_config.enable_fsdp is True + + +def test_stream_get_service_uses_passed_gpu_num_and_ppl_fps() -> 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) + assert service.pipeline is pipeline + assert service.default_fps == stream_example.PPL_CONFIG["target_fps"] + + +def test_stream_get_service_retains_example_default_gpu_num() -> None: + pipeline = MagicMock() + + with patch.object(stream_example, "get_pipeline", return_value=pipeline) as get_pipeline: + service = stream_example.get_service() + + get_pipeline.assert_called_once_with(parallelism=stream_example.PPL_CONFIG["parallelism"]) + assert service.pipeline is pipeline + + +def test_stream_get_pipeline_reads_tf_model_zoo_path(monkeypatch) -> None: + monkeypatch.setenv("TF_MODEL_ZOO_PATH", "/models") + + pipeline = MagicMock() + with patch.object(stream_example, "LingBotWorldFastPipeline", return_value=pipeline): + stream_example.get_pipeline() + + assert len(pipeline.init.call_args.args) == 1 + config = pipeline.init.call_args.args[0] + assert config.checkpoint_dir == "/models/Wan2.2-I2V-A14B" + assert config.fast_checkpoint_path == "/models/lingbot/lingbot-world-fast" diff --git a/tests/unit/service/test_service_routes.py b/tests/unit/service/test_service_routes.py index f982f8d..974ce76 100644 --- a/tests/unit/service/test_service_routes.py +++ b/tests/unit/service/test_service_routes.py @@ -308,6 +308,8 @@ def fake_run_stream_server(pipe_path, port, host, **kwargs): [ "stream-serve", "stream_pipeline.py", + "--gpu-num", + "3", "--security-level", "none", ], @@ -315,6 +317,7 @@ def fake_run_stream_server(pipe_path, port, host, **kwargs): assert result.exit_code == 0 assert captured["validated"] == "stream_pipeline.py" + assert captured["gpu_num"] == 3 assert captured["security_level"] == "none" assert captured["skip_validation"] is False @@ -419,6 +422,44 @@ def test_stream_pipeline_service_uses_injected_config() -> None: assert service.validation_config.strict_validation is False +def test_stream_pipeline_service_passes_gpu_num_to_factory(monkeypatch: pytest.MonkeyPatch) -> None: + captured = {} + + class FakeBidirectionalService: + def start(self) -> None: + captured["started"] = True + + def stop(self) -> None: + pass + + def create_session(self, config: dict) -> str: + return "session" + + def push_chunk(self, session_id: str, chunk: dict) -> None: + pass + + async def pull_chunks(self, session_id: str): + if False: + yield {} + + def close_session(self, session_id: str) -> None: + pass + + def get_service(gpu_num: int) -> FakeBidirectionalService: + captured["gpu_num"] = gpu_num + return FakeBidirectionalService() + + module = types.SimpleNamespace(get_service=get_service) + monkeypatch.setattr( + "telefuser.service.core.stream_pipeline_service.load_pipeline_module", + lambda *args, **kwargs: (module, "test_stream_module"), + ) + service = StreamPipelineService(config=ServerConfig(security_level=SecurityLevel.NONE)) + + assert service.start_service("stream_pipeline.py", gpu_num=3, skip_validation=True) + assert captured == {"gpu_num": 3, "started": True} + + def test_webrtc_routes_use_api_server_config(monkeypatch: pytest.MonkeyPatch) -> None: captured = {} diff --git a/tests/unit/worker/test_parallel_worker_lifecycle.py b/tests/unit/worker/test_parallel_worker_lifecycle.py new file mode 100644 index 0000000..7254455 --- /dev/null +++ b/tests/unit/worker/test_parallel_worker_lifecycle.py @@ -0,0 +1,87 @@ +import threading +from queue import Empty +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest + +from telefuser.worker.parallel_worker import ParallelWorker + + +def _worker() -> ParallelWorker: + worker = ParallelWorker.__new__(ParallelWorker) + worker.name = "Parallel Worker TestStage" + worker.timeout = 1 + worker._lifecycle_lock = threading.Lock() + worker._failed = False + worker._closed = False + worker._failure_reason = None + worker.queue_out = MagicMock() + worker.queue_in = [MagicMock(), MagicMock()] + worker.ctx = SimpleNamespace(processes=[MagicMock(), MagicMock()]) + for process in worker.ctx.processes: + process.is_alive.return_value = True + return worker + + +def test_timeout_marks_group_failed_terminates_all_ranks_and_rejects_reuse() -> None: + worker = _worker() + worker.queue_out.get.side_effect = Empty() + + with pytest.raises(RuntimeError, match="timeout after 1 seconds"): + worker._wait_result("denoise") + + assert worker.failed is True + assert worker.closed is False + assert worker.failure_reason == "denoise timeout after 1 seconds" + for process in worker.ctx.processes: + process.kill.assert_called_once_with() + process.join.assert_called_once_with(timeout=2) + + with pytest.raises(RuntimeError, match="has failed"): + worker._ensure_usable() + + +def test_remote_rank_error_marks_entire_group_failed() -> None: + worker = _worker() + worker.queue_out.get.return_value = RuntimeError("Parallel worker rank 2 failed") + + with pytest.raises(RuntimeError, match="rank 2 failed"): + worker._wait_result("initialize_cache") + + assert worker.failed is True + assert "initialize_cache failed" in worker.failure_reason + for process in worker.ctx.processes: + process.kill.assert_called_once_with() + + +def test_close_is_idempotent_and_rejects_new_work() -> None: + worker = _worker() + for process in worker.ctx.processes: + process.is_alive.return_value = False + + worker.close() + worker.close() + + assert worker.closed is True + assert worker.failed is False + for queue in worker.queue_in: + queue.put.assert_called_once_with(["exit", None, None]) + queue.close.assert_called_once_with() + worker.queue_out.close.assert_called_once_with() + + with pytest.raises(RuntimeError, match="is closed"): + worker._ensure_usable() + + +def test_close_after_failure_does_not_submit_more_worker_commands() -> None: + worker = _worker() + worker._mark_failed("rank 1 failed") + + worker.close() + + assert worker.failed is True + assert worker.closed is True + for queue in worker.queue_in: + queue.put.assert_not_called() + queue.close.assert_called_once_with()