diff --git a/CLAUDE.md b/CLAUDE.md index 1b98247..8d55923 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -6,7 +6,7 @@ TeleFuser is a high-performance framework for efficient multimodal generation mo **Tech Stack:** Python 3.10-3.13, PyTorch 2.6+, CUDA 12.8+, FastAPI, Ray -**Supported Models:** WanVideo (Wan2.1/2.2), Qwen-Image, Z-Image, FlashVSR, HunyuanVideo, Flux2 Klein, LTX Video, LiveAct, LongCat-Video, LingBot-World +**Supported Models:** WanVideo (Wan2.1/2.2), Qwen-Image, Z-Image, FlashVSR, HunyuanVideo, Flux2 Klein, LTX Video, LiveAct, LongCat-Video, LingBot-World, LingBot-Video ## Commands @@ -43,6 +43,7 @@ telefuser/ │ ├── longcat_video/ # LongCat-Video: T2V, I2V │ ├── lingbot_world_fast/ # LingBot shared causal-fast engine │ ├── lingbot_world_v2/ # LingBot-World v2 causal-fast facade +│ ├── lingbot_video/ # LingBot-Video Dense/MoE/refiner runtime │ └── common/ # Shared pipeline utilities ├── models/ # Model architectures: DiT, VAE, text encoders ├── ops/ # Custom operations: attention, FFN, normalization @@ -82,6 +83,14 @@ telefuser/ - Interpret chunk period as output cadence: real-time operation requires p95 to stay below the media duration represented by one chunk, with margin for transport and encoding. +### LingBot-Video Single-Process Runtime + +- Dense and MoE LingBot-Video requests use structured JSON captions. Spatial height and width must be divisible by 16: the Wan VAE downsamples by 8 and the DiT uses a spatial patch size of 2. +- TI2V conditions are independent Qwen3-VL visual tokens and a VAE clean frame-zero latent. Preserve both paths and reapply the latent condition after every denoising step. +- The MoE refiner is a separately loaded stage. For a shared GPU, release/offload base stages before loading it, and retain the native RGB handoff rather than introducing an MP4 round trip. +- The sorted eager MoE path is the validated single-GPU correctness implementation. Four-GPU MoE uses native grouped GEMM when `torch._grouped_mm` is available and retains sorted eager as an explicit fallback. FP8 and expert parallelism require separate parity and benchmark evidence before being enabled. +- Distributed LingBot base sampling keeps the complete scheduler loop inside each stage worker. Do not move per-step latents through the parent process. Spawned distributed workers use one PyTorch intra-op CPU thread per rank, matching `torchrun` and avoiding host launch-thread oversubscription. + ### Layer Architecture Principles For Models TeleFuser's model follows a strict layered architecture for operations: @@ -145,6 +154,7 @@ When adding or porting a pipeline, preserve upstream behavior and reuse TeleFuse - Before editing, select the closest maintained pipeline, public example, and tests as structural baselines. Read the relevant adding-new-example, adding-new-model, adding-new-stage, model-loading, configuration, and service documentation. - Inventory required model-specific classes and configuration fields and map them to upstream behavior and the selected baseline. - List every proposed framework-level or cross-pipeline interface, general-purpose configuration field, environment variable, loader, registry, CLI option, or service schema deviation. The expected list is empty. +- Reusing a framework API does not authorize changing that API. During model or pipeline integration, do not add or broaden shared loading, registry, configuration, orchestration, or service behavior merely for convenience. First use the existing API exactly as-is, including its supported file-list and wildcard inputs. If it cannot express the requirement, report the precise gap, alternatives, affected callers, and compatibility impact, then obtain explicit user approval before editing shared framework code. - Reuse `BasePipeline`, `BaseStage`, `ModuleManager`, existing configuration dataclasses, example contracts, and service schemas. Do not create parallel interfaces or attach ad-hoc configuration attributes for convenience. - If existing extension points cannot express a requirement, report the exact gap and obtain user approval before introducing a new public interface or configuration mechanism. - Do not add an environment variable during pipeline integration unless the user explicitly requests it or an existing documented variable already has the exact required semantics. Prefer function parameters for request-scoped inputs, dataclass fields for runtime configuration, CLI options for command-line inputs, and service schemas for API inputs. diff --git a/README.md b/README.md index cb9027a..c5b43c4 100644 --- a/README.md +++ b/README.md @@ -17,6 +17,7 @@ TeleFuser is a high-performance runtime for world model inference and multimodal ## News 📰 +- ✨ **2026-07-22**: **NEW** Added [**LingBot-Video**](examples/lingbot_video/README.md) support for Dense and MoE T2I/T2V/TI2V generation, native four-GPU CFG/SP execution, and in-memory MoE refinement. - ✨ **2026-07-15**: Added [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) support for offline generation, interactive WebRTC streaming, and multi-GPU inference. - ✨ **2026-07-06**: Added external **CacheSeek** latent cache integration for service-mode cross-request reuse. Cache hits can skip the first N denoising steps; the Wan2.2 cache-enabled service example snapshots `[5, 10, 15, 20, 25]` by default. See [docs/en/latent_cache.md](docs/en/latent_cache.md). @@ -195,6 +196,7 @@ telefuser/ | `HunyuanVideo` | T2V, I2V | Supported via [examples/hunyuan_video/README.md](examples/hunyuan_video/README.md) | | `LTX Video` | I2V + Audio | Unified audio-video generation via [examples/ltx_video/README.md](examples/ltx_video/README.md) | | `LongCat-Video` | T2V, I2V, VC | Long-form generation and continuation via [examples/longcat_video/README.md](examples/longcat_video/README.md) | +| **NEW** `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | Dense/MoE generation with native CFG/SP and an in-memory base-to-refiner path; see [examples/lingbot_video/README.md](examples/lingbot_video/README.md) | ### Image Generation and Other Multimodal Pipelines diff --git a/README_zh.md b/README_zh.md index 82d1ed3..c18f74f 100644 --- a/README_zh.md +++ b/README_zh.md @@ -17,6 +17,7 @@ TeleFuser 是一个面向世界模型推理与多模态生成的高性能运行 ## News 📰 +- ✨ **2026-07-22**：**NEW** 新增 [**LingBot-Video**](examples/lingbot_video/README.md) 支持，覆盖 Dense/MoE T2I、T2V、TI2V、原生四卡 CFG/SP 推理与内存直传 MoE refiner。 - ✨ **2026-07-15**：新增 [**LingBot-World v2**](https://github.com/Robbyant/lingbot-world-v2) 支持，支持离线生成、交互式 WebRTC 流和多卡推理。 - ✨ **2026-07-06**：新增外部 **CacheSeek** latent cache 集成，支持服务模式下跨请求复用；命中后可跳过前 N 步去噪。Wan2.2 服务示例默认快照 `[5, 10, 15, 20, 25]`。配置和安装方式见 [docs/zh/latent_cache.md](docs/zh/latent_cache.md)。 @@ -194,6 +195,7 @@ telefuser/ | `HunyuanVideo` | T2V, I2V | 见 [examples/hunyuan_video/README.md](examples/hunyuan_video/README.md) | | `LTX Video` | I2V + Audio | 统一音视频生成，见 [examples/ltx_video/README.md](examples/ltx_video/README.md) | | `LongCat-Video` | T2V, I2V, VC | 长视频生成与续写，见 [examples/longcat_video/README.md](examples/longcat_video/README.md) | +| **NEW** `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | 支持原生 CFG/SP 的 Dense/MoE 生成与内存直传 base-to-refiner，见 [examples/lingbot_video/README.md](examples/lingbot_video/README.md) | ### 图像与其他多模态生成 diff --git a/docs/en/configuration.md b/docs/en/configuration.md index 38b7fbc..e5d56ef 100644 --- a/docs/en/configuration.md +++ b/docs/en/configuration.md @@ -251,6 +251,7 @@ class ParallelConfig: pp_degree: int = 1 # Pipeline parallelism tp_degree: int = 1 # Tensor parallelism enable_fsdp: bool = False + worker_intra_op_threads: int = 1 # CPU intra-op threads per worker def validate(self) -> None: """Validate that device count matches parallelism degrees.""" diff --git a/docs/en/index.md b/docs/en/index.md index c96457e..6a249e1 100644 --- a/docs/en/index.md +++ b/docs/en/index.md @@ -56,7 +56,8 @@ Reusable stages, model configs, schedulers, and pipeline orchestration. | Model | Tasks | Description | |-------|-------|-------------| -| LingBot-World-Fast | Bidirectional streaming | Interactive world model via WebRTC DataChannel | +| LingBot-World v2 | Bidirectional streaming | Camera-controlled interactive world model via WebRTC | +| LingBot-World-Fast | Bidirectional streaming | Legacy/causal-fast interactive world model via WebRTC DataChannel | ### Video Generation @@ -68,6 +69,7 @@ Reusable stages, model configs, schedulers, and pipeline orchestration. | FlashVSR | VSR | Video super-resolution | | LiveAct | S2V | Speech-to-video | | LongCat-Video | T2V, I2V | Long video generation | +| LingBot-Video | T2I, T2V, TI2V, MoE refiner | Precision-first Dense and MoE video generation | ### Image Generation diff --git a/docs/en/parallel.md b/docs/en/parallel.md index ce171d6..b2ee03a 100644 --- a/docs/en/parallel.md +++ b/docs/en/parallel.md @@ -385,8 +385,14 @@ class ParallelConfig: enable_fsdp: bool = False # Enable FSDP timeout: int = 600 # Timeout in seconds queue_with_cpu: bool = False # Use CPU queue + worker_intra_op_threads: int = 1 # CPU intra-op threads per worker ``` +The default of one thread matches `torchrun` and prevents CPU oversubscription +across GPU worker processes. Increase it only for distributed stages with +substantial CPU-side computation; it does not change the parent process thread +pool. + ### Validation Rules ```python @@ -611,4 +617,4 @@ RuntimeError: ParallelWorker timeout - [Ulysses: Sequence Parallelism](https://arxiv.org/abs/2309.14509) - [Ring Attention](https://arxiv.org/abs/2310.01889) - [GPipe: Pipeline Parallelism](https://arxiv.org/abs/1811.06965) -- [PyTorch Distributed](https://pytorch.org/docs/stable/distributed.html) \ No newline at end of file +- [PyTorch Distributed](https://pytorch.org/docs/stable/distributed.html) diff --git a/docs/en/service.md b/docs/en/service.md index 0af6aa3..bf8ea50 100644 --- a/docs/en/service.md +++ b/docs/en/service.md @@ -100,6 +100,7 @@ See the [Stream Server Guide](stream_server.md) for full streaming documentation | `LiveAct` | S2V (speech-to-video) | Speech-driven talking head generation | | `FlashVSR` | VSR | Streaming video super-resolution | | `LongCat-Video` | T2V, I2V, VC | Long-form generation and continuation | +| `LingBot-World v2` | Bidirectional world-model streaming | Camera-controlled WebRTC loop — see [Stream Server Guide](stream_server.md) | ### Video Generation @@ -108,6 +109,7 @@ See the [Stream Server Guide](stream_server.md) for full streaming documentation | `WanVideo` (Wan2.1 / Wan2.2) | T2V, I2V, FL2V | Main video generation family | | `HunyuanVideo` | T2V, I2V | Supported via service examples | | `LTX Video` | I2V + Audio | Unified audio-video generation | +| `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | Dense/MoE generation with native CFG/SP and in-memory refinement; see `examples/lingbot_video/README.md` | ### Image Generation diff --git a/docs/zh/configuration.md b/docs/zh/configuration.md index 4e1b79d..2864a08 100644 --- a/docs/zh/configuration.md +++ b/docs/zh/configuration.md @@ -251,6 +251,7 @@ class ParallelConfig: pp_degree: int = 1 # 流水线并行 tp_degree: int = 1 # 张量并行 enable_fsdp: bool = False + worker_intra_op_threads: int = 1 # 每个 worker 的 CPU intra-op 线程数 def validate(self) -> None: """验证设备数量与并行度匹配。""" diff --git a/docs/zh/index.md b/docs/zh/index.md index 2690f9d..82c1b03 100644 --- a/docs/zh/index.md +++ b/docs/zh/index.md @@ -55,7 +55,8 @@ AdaTaylorCache 和运行时缓存控制，面向重复生成工作负载。 | 模型 | 任务 | 描述 | |------|------|------| -| LingBot-World-Fast | 双向流式推理 | 通过 WebRTC DataChannel 的交互式世界模型 | +| LingBot-World v2 | 双向流式推理 | 通过 WebRTC 的相机控制交互式世界模型 | +| LingBot-World-Fast | 双向流式推理 | Legacy/causal-fast WebRTC 交互式世界模型 | ### 视频生成 @@ -67,6 +68,7 @@ AdaTaylorCache 和运行时缓存控制，面向重复生成工作负载。 | FlashVSR | VSR | 视频超分辨率 | | LiveAct | S2V | 语音转视频 | | LongCat-Video | T2V, I2V | 长视频生成 | +| LingBot-Video | T2I, T2V, TI2V, MoE refiner | 精度优先的 Dense/MoE 视频生成 | ### 图像生成 diff --git a/docs/zh/parallel.md b/docs/zh/parallel.md index 088f05a..45c63ca 100644 --- a/docs/zh/parallel.md +++ b/docs/zh/parallel.md @@ -385,8 +385,12 @@ class ParallelConfig: enable_fsdp: bool = False # 启用 FSDP timeout: int = 600 # 超时时间（秒） queue_with_cpu: bool = False # 使用 CPU 队列 + worker_intra_op_threads: int = 1 # 每个 worker 的 CPU intra-op 线程数 ``` +默认值 1 与 `torchrun` 一致，可避免多个 GPU worker 的 CPU 线程过度竞争。只有分布式 +stage 包含大量 CPU 计算时才需要提高该值；它不会修改父进程的线程池。 + ### 验证规则 ```python @@ -611,4 +615,4 @@ RuntimeError: ParallelWorker timeout - [Ulysses: Sequence Parallelism](https://arxiv.org/abs/2309.14509) - [Ring Attention](https://arxiv.org/abs/2310.01889) - [GPipe: Pipeline Parallelism](https://arxiv.org/abs/1811.06965) -- [PyTorch Distributed](https://pytorch.org/docs/stable/distributed.html) \ No newline at end of file +- [PyTorch Distributed](https://pytorch.org/docs/stable/distributed.html) diff --git a/docs/zh/service.md b/docs/zh/service.md index bd1e90c..4d0e3a8 100644 --- a/docs/zh/service.md +++ b/docs/zh/service.md @@ -100,6 +100,7 @@ TeleFuser 提供两种服务命令，针对不同工作负载类型优化： | `LiveAct` | S2V（语音转视频） | 语音驱动说话人头部生成 | | `FlashVSR` | VSR | 流式视频超分辨率 | | `LongCat-Video` | T2V, I2V, VC | 长视频生成和续写 | +| `LingBot-World v2` | 双向世界模型流式推理 | 相机控制的 WebRTC 闭环 — 参见[流式服务指南](stream_server.md) | ### 视频生成 @@ -108,6 +109,7 @@ TeleFuser 提供两种服务命令，针对不同工作负载类型优化： | `WanVideo`（Wan2.1 / Wan2.2） | T2V, I2V, FL2V | 主要视频生成系列 | | `HunyuanVideo` | T2V, I2V | 通过服务示例支持 | | `LTX Video` | I2V + Audio | 音视频统一生成 | +| `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | 支持原生 CFG/SP 的 Dense/MoE 生成与内存直传 refinement；参见 `examples/lingbot_video/README.md` | ### 图像生成 diff --git a/examples/example_config.yaml b/examples/example_config.yaml index a0a05d7..e0138ac 100644 --- a/examples/example_config.yaml +++ b/examples/example_config.yaml @@ -116,6 +116,24 @@ pipelines: attn_impl: TORCH_SDPA frame_num: 13 + # ============================================================ + # lingbot_video — Fixed Dense and MoE examples + # ============================================================ + + lingbot_video_dense_1_3b_t2v: + script: lingbot_video/lingbot_video_dense_1_3b.py + gpu_count: 1 + output_type: video + timeout_seconds: 1800 + resolution: "480p" + + lingbot_video_moe_30b_refiner_t2v: + script: lingbot_video/lingbot_video_moe_30b.py + gpu_count: 4 + output_type: video + timeout_seconds: 3600 + resolution: "480p" + # ============================================================ # hunyuan_video # ============================================================ diff --git a/examples/lingbot_video/README.md b/examples/lingbot_video/README.md new file mode 100644 index 0000000..aedb331 --- /dev/null +++ b/examples/lingbot_video/README.md @@ -0,0 +1,84 @@ +# LingBot-Video + +Dense 1.3B and MoE 30B are separate examples. Each module exposes +`PPL_CONFIG`, `CONTRACT`, `get_pipeline`, `run`, and `run_with_file` for the +shared CLI runner and TeleFuser service. +Both examples default to the official five-second structured caption in `assets/t2v_5s.json.example` and the validated 832x480 LingBot landscape geometry. +Their default checkpoints are resolved from `TF_MODEL_ZOO_PATH`, which defaults +to `/hhb-data/aigc/model_zoo` in the current environment. + +The model-specific files also contain their checkpoint loading, stage assembly, request +handling, refiner lifecycle, and output encoding. Shared behavior uses TeleFuser contract +templates and video utilities directly. + +Use a structured JSON caption produced by the LingBot rewriter. Dense 1.3B T2V: + +```bash +python examples/lingbot_video/lingbot_video_dense_1_3b.py \ + --model_root /path/to/lingbot-video-dense-1.3b \ + --prompt "$(cat /path/to/caption.json)" --output_path result.mp4 +``` + +Pass `--task i2v --first_image_path first_frame.png` for TI2V. + +For the MoE checkpoint refiner, use the in-memory base-to-refiner path: + +```bash +python examples/lingbot_video/lingbot_video_moe_30b.py \ + --model_root /path/to/lingbot-video-moe-30b-a3b --refine \ + --prompt "$(cat /path/to/caption.json)" --output_path result.mp4 +``` + +On four GPUs the base and refiner DiTs remain resident together during refiner +denoising by default. Both worker groups are released before the high-resolution +VAE decode. Use `--no-refiner_co_resident` for the lower-memory sequential +lifecycle, or `--no-refine` for MoE base-only generation. `expert_backend=auto` keeps +the validated sorted eager path on one GPU and selects native grouped GEMM for +four-GPU inference. The grouped path requires a CUDA PyTorch build that exposes +`torch._grouped_mm`; use `--expert_backend sorted` as the explicit fallback. +`--expert_backend fp8` quantizes routed expert weights per output channel and +uses native dynamic W8A8 scaled GEMMs. It reduces expert residency but is an +explicit memory-oriented backend until a grouped FP8 kernel is available. + +Four-GPU base and refiner stages can split the devices as CFG2 x SP2: + +```bash +python examples/lingbot_video/lingbot_video_moe_30b.py \ + --gpu_num 4 --cfg_parallel_degree 2 \ + --refiner_gpu_num 4 --refiner_cfg_parallel_degree 2 \ + --refiner_co_resident --expert_backend fp8 \ + --model_root /path/to/lingbot-video-moe-30b-a3b \ + --output_path result.mp4 +``` + +CFG parallel and batch CFG are mutually exclusive. Use +`--cfg_parallel_degree 2` for Dense or MoE base generation and +`--refiner_cfg_parallel_degree 2` for the refiner. A degree of one retains SP4. + +Set `PPL_CONFIG["model_root"]` in the selected example, then serve structured-caption T2I/T2V/TI2V requests with: + +```bash +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --port 8000 +``` + +Serve MoE independently with `lingbot_video_moe_30b.py`. Refiner requests are +enabled by default and can override the `refine` contract parameter. Set +`PPL_CONFIG["refiner_parallelism"] = 4` to select its distributed FSDP stage +explicitly; otherwise it inherits service parallelism. The CFG degrees determine +whether four workers use SP4 or CFG2 x SP2. + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 +``` + +Dense and MoE base checkpoints support TeleFuser-native four-GPU FSDP plus Ulysses sequence parallelism: + +```bash +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 +``` + +For MoE, set `PPL_CONFIG["model_root"]` in the MoE example and serve +`lingbot_video_moe_30b.py`. Do not use `torchrun` for either service: TeleFuser +creates and manages the workers. Configure `refiner_co_resident=False` when the +selected dtype, CFG layout, or GPU memory cannot hold both DiTs during refiner +denoising. diff --git a/examples/lingbot_video/assets/t2v_5s.json.example b/examples/lingbot_video/assets/t2v_5s.json.example new file mode 100644 index 0000000..92afdf6 --- /dev/null +++ b/examples/lingbot_video/assets/t2v_5s.json.example @@ -0,0 +1,77 @@ +{ + "caption": { + "comprehensive_description": { + "scene_content_description": "A young woman with long, wavy brown hair is standing in a bright, modern apartment living room. She is wearing a stylish, oversized cream-colored knit cardigan over a white tank top, paired with high-waisted, wide-leg beige trousers. She holds a small, structured tan leather handbag in her left hand. The background features a neutral-toned interior with a beige sofa, a potted plant, and large windows that let in soft, natural light, creating a warm and inviting atmosphere. The woman is smiling and looking directly at the camera, showcasing her outfit with a confident and friendly demeanor.", + "camera_movement_description": "The camera is positioned at eye level and remains essentially stationary throughout the video, maintaining a medium shot that captures the subject from the waist up. There is a very shallow depth of field, keeping the woman in sharp focus while the background remains softly blurred." + }, + "camera_info": { + "color": "Warm", + "frame_size": "Medium", + "shot_type_angle": "Eye level", + "lens_size": "Medium", + "composition": "Center", + "lighting": "Soft light", + "lighting_type": "Daylight" + }, + "world_knowledge": [], + "prominent_elements": [ + { + "name": "young woman", + "description": "A woman with long, wavy brown hair and a friendly expression, modeling a fashion outfit.", + "actions": [ + { + "timestamp": "[0.0s - 0.5s]", + "action": "stands still, smiling at the camera" + }, + { + "timestamp": "[0.5s - 2.0s]", + "action": "shifts her weight and turns her body slightly to the right" + }, + { + "timestamp": "[2.0s - 3.5s]", + "action": "adjusts the collar of her cardigan with her right hand" + }, + { + "timestamp": "[3.5s - 5.0s]", + "action": "returns to a neutral pose, smiling at the camera" + } + ], + "location": "center of the frame", + "relative_size": "dominant", + "shape_and_color": "slender build; wearing cream, white, and beige", + "texture": "soft knit cardigan, smooth fabric trousers", + "appearance_details": "long wavy brown hair, gold hoop earrings, tan leather handbag", + "relationship": "the main subject of the video, standing in front of a blurred apartment background", + "orientation": "upright, facing the camera", + "pose": "standing, shifting weight, and adjusting clothing", + "expression": "smiling and confident", + "clothing": "oversized cream knit cardigan, white tank top, high-waisted wide-leg beige trousers", + "gender": "female", + "skin_tone_and_texture": "fair skin with a smooth texture" + }, + { + "name": "tan handbag", + "description": "A small, structured leather handbag with a top handle.", + "actions": [ + { + "timestamp": "[0.0s - 5.0s]", + "action": "held steady in the woman's left hand" + } + ], + "location": "held in the woman's left hand, lower center of the frame", + "relative_size": "small", + "shape_and_color": "rectangular, tan or light brown", + "texture": "smooth leather", + "appearance_details": "structured shape with a top handle", + "relationship": "held by the woman as an accessory", + "orientation": "upright", + "pose": "", + "expression": "", + "clothing": "", + "gender": "", + "skin_tone_and_texture": "" + } + ] + }, + "duration": 5 +} diff --git a/examples/lingbot_video/lingbot_video_dense_1_3b.py b/examples/lingbot_video/lingbot_video_dense_1_3b.py new file mode 100644 index 0000000..641f829 --- /dev/null +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -0,0 +1,356 @@ +"""LingBot-Video Dense 1.3B CLI and service example.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import click +import numpy as np +import torch +from PIL import Image +from diffusers.utils import export_to_video + +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + ModelRuntimeConfig, + OffloadConfig, + ParallelConfig, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video import ( + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_NEGATIVE_PROMPT_IMAGE, + LingBotVideoPipeline, + LingBotVideoPipelineConfig, + LingBotVideoRequest, + default_negative_caption, + load_lingbot_video_model_config, + num_frames_from_duration, + parse_lingbot_video_prompt, +) +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler +from telefuser.service.core.contract_templates import build_pipeline_manifest, build_task_contract_template +from telefuser.utils.video import get_target_video_size_from_ratio + +_DEFAULT_PROMPT = (Path(__file__).parent / "assets" / "t2v_5s.json.example").read_text(encoding="utf-8") +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "/hhb-data/aigc/model_zoo") + +PPL_CONFIG: dict[str, Any] = { + "name": "lingbot_video_dense_1_3b", + "model_root": TF_MODEL_ZOO_PATH + "/lingbot/lingbot-video-dense-1.3b", + "variant": "dense", + "supports_refiner": False, + "prompt": _DEFAULT_PROMPT, + "num_inference_steps": 40, + "guidance_scale": 3.0, + "fps": 24, + "target_video_length": 5, + "resolution": "480p", + "aspect_ratio": "16:9", + "seed": 42, + "cfg_parallel_degree": 1, +} + + +def _build_contract() -> dict[str, Any]: + """Build the standard service contract with LingBot-specific defaults.""" + task_contracts = {} + for task in ("t2i", "t2v", "i2v"): + overrides = { + "negative_prompt": {"default": DEFAULT_NEGATIVE_PROMPT_IMAGE if task == "t2i" else DEFAULT_NEGATIVE_PROMPT}, + "seed": {"default": PPL_CONFIG["seed"]}, + "resolution": {"default": PPL_CONFIG["resolution"]}, + "aspect_ratio": {"default": PPL_CONFIG["aspect_ratio"]}, + } + if task != "t2i": + overrides["target_video_length"] = {"default": PPL_CONFIG["target_video_length"]} + task_contracts[task] = build_task_contract_template(task, parameter_overrides=overrides) + return build_pipeline_manifest( + pipeline_name=PPL_CONFIG["name"], + supported_tasks=task_contracts, + task_contracts=task_contracts, + ) + + +PIPELINE_CONTRACT = _build_contract() + + +def build_pipeline( + model_root: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + cpu_offload: bool = False, + guidance_scale: float = PPL_CONFIG["guidance_scale"], + num_inference_steps: int = PPL_CONFIG["num_inference_steps"], + shift: float = 3.0, + attention_config: AttentionConfig | None = None, + parallel_config: ParallelConfig | None = None, + batch_cfg: bool | None = None, +) -> LingBotVideoPipeline: + """Load Dense modules and initialize the complete pipeline.""" + try: + from diffusers import AutoencoderKLWan + from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + except ImportError as exc: + raise RuntimeError("LingBot-Video requires diffusers and transformers") from exc + + parallel_config = parallel_config or ParallelConfig() + if batch_cfg is None: + batch_cfg = parallel_config.world_size > 1 and parallel_config.cfg_degree == 1 + if batch_cfg and parallel_config.cfg_degree > 1: + raise ValueError("LingBot CFG parallel and batch CFG are mutually exclusive") + if parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot FSDP inference requires cpu_offload=False") + runtime_device = torch.device(device) + root = Path(model_root) + transformer_dir = root / "transformer" + offload_config = OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD + ) + runtime_config = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch_dtype, + attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + offload_config=offload_config, + parallel_config=parallel_config, + ) + + module_manager = ModuleManager(torch_dtype=torch_dtype, device="cpu") + module_manager.load_model(str(transformer_dir), name="transformer", torch_dtype=torch_dtype) + transformer = module_manager.fetch_module("transformer") + if transformer is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {transformer_dir}") + transformer.promote_stability_layers_to_fp32() + module_manager.load_from_huggingface( + str(root / "processor"), + module_source="transformers", + module_name="processor", + module_class=AutoProcessor, + ) + module_manager.load_from_huggingface( + str(root / "text_encoder"), + module_source="transformers", + module_name="text_encoder", + module_class=Qwen3VLForConditionalGeneration, + torch_dtype=torch_dtype, + attn_implementation="sdpa", + ) + module_manager.load_from_huggingface( + str(root / "vae"), + module_source="diffusers", + module_name="vae", + module_class=AutoencoderKLWan, + torch_dtype=torch.float32, + ) + module_manager.load_from_huggingface( + str(root / "scheduler"), + module_source="diffusers", + module_name="scheduler", + module_class=FlowUniPCMultistepScheduler, + ) + + pipeline = LingBotVideoPipeline(device=device, torch_dtype=torch_dtype) + pipeline.init( + module_manager, + LingBotVideoPipelineConfig( + model=load_lingbot_video_model_config(transformer_dir, variant="dense"), + text_encoding_config=runtime_config, + dit_config=runtime_config, + vae_config=runtime_config, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + shift=shift, + batch_cfg=batch_cfg, + enable_denoising_parallel=parallel_config.world_size > 1, + ), + ) + return pipeline + + +def get_pipeline( + parallelism: int = 1, + model_root: str = PPL_CONFIG["model_root"], + cfg_parallel_degree: int | None = None, +) -> object: + """Load the fixed Dense 1.3B checkpoint on one or four GPUs.""" + if parallelism not in {1, 4}: + raise ValueError("LingBot-Video Dense supports parallelism=1 or 4") + cfg_parallel_degree = int(PPL_CONFIG["cfg_parallel_degree"]) if cfg_parallel_degree is None else cfg_parallel_degree + if cfg_parallel_degree not in {1, 2} or parallelism % cfg_parallel_degree: + raise ValueError("LingBot CFG parallel degree must be 1 or 2 and divide parallelism") + return build_pipeline( + model_root, + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)), + cfg_degree=cfg_parallel_degree, + sp_ulysses_degree=parallelism // cfg_parallel_degree, + enable_fsdp=parallelism > 1, + ), + ) + + +def _load_condition_image(path: str) -> torch.Tensor: + """Load an RGB image path as a raw [0,255] TI2V tensor.""" + pixels = np.asarray(Image.open(path).convert("RGB")).copy() + return torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).float() + + +def _resolve_video_size(aspect_ratio: str, resolution: str) -> tuple[int, int]: + """Resolve output geometry using the validated LingBot 480p landscape preset.""" + if resolution == "480p" and aspect_ratio == "16:9": + return 832, 480 + width, height = get_target_video_size_from_ratio( + aspect_ratio, + resolution, + height_division_factor=16, + width_division_factor=16, + ) + if width is None or height is None: + raise ValueError(f"unsupported LingBot-Video resolution: {resolution}") + return width, height + + +def _generator(pipeline: object, seed: int) -> torch.Generator: + """Create a deterministic generator on the pipeline device type.""" + device_type = torch.device(getattr(pipeline, "device", "cpu")).type + return torch.Generator(device_type).manual_seed(seed) + + +def run( + pipeline: object, + prompt: str = PPL_CONFIG["prompt"], + negative_prompt: str | None = None, + seed: int = PPL_CONFIG["seed"], + resolution: str = PPL_CONFIG["resolution"], + aspect_ratio: str = PPL_CONFIG["aspect_ratio"], + target_video_length: float | None = None, + task: str = "t2v", + first_image_path: str = "", +) -> torch.Tensor: + """Run Dense T2I, T2V, or I2V and return normalized RGB frames.""" + if task not in {"t2i", "t2v", "i2v"}: + raise ValueError(f"unsupported LingBot-Video task: {task}") + if task == "i2v" and not first_image_path: + raise ValueError("LingBot-Video i2v requires first_image_path") + + caption, prompt_duration = parse_lingbot_video_prompt(json.loads(prompt)) + width, height = _resolve_video_size(aspect_ratio, resolution) + duration = target_video_length if target_video_length is not None else prompt_duration + if duration is None: + duration = float(PPL_CONFIG["target_video_length"]) + num_frames = 1 if task == "t2i" else num_frames_from_duration(duration, fps=int(PPL_CONFIG["fps"])) + resolved_negative_prompt = negative_prompt if negative_prompt is not None else default_negative_caption(num_frames) + return pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=height, + width=width, + num_frames=num_frames, + image=_load_condition_image(first_image_path) if first_image_path else None, + ), + negative_caption=resolved_negative_prompt, + generator=_generator(pipeline, seed), + ).output + + +def _save_output(frames: torch.Tensor, output_path: str) -> dict[str, str]: + """Encode normalized RGB without an intermediate uint8 video conversion.""" + path = Path(output_path) + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + if video.shape[0] == 1: + if path.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}: + path = path.with_suffix(".png") + path.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray((video[0] * 255).round().astype(np.uint8)).save(path) + else: + path.parent.mkdir(parents=True, exist_ok=True) + export_to_video(list(video), str(path), fps=int(PPL_CONFIG["fps"])) + return {"output_path": str(path)} + + +def run_with_file( + pipeline: object, + prompt: str = PPL_CONFIG["prompt"], + negative_prompt: str | None = None, + seed: int = PPL_CONFIG["seed"], + output_path: str = "output.mp4", + target_video_length: float | None = None, + resolution: str = PPL_CONFIG["resolution"], + aspect_ratio: str = PPL_CONFIG["aspect_ratio"], + task: str = "t2v", + first_image_path: str = "", + **_: object, +) -> dict[str, str]: + """Run the Dense example and save its image or video output.""" + frames = run( + pipeline, + prompt, + negative_prompt, + seed, + resolution, + aspect_ratio, + target_video_length, + task, + first_image_path, + ) + return _save_output(frames, output_path) + + +@click.command() +@click.option("--gpu_num", default=1, type=int, help="Number of GPUs: 1 or 4") +@click.option("--cfg_parallel_degree", default=PPL_CONFIG["cfg_parallel_degree"], type=click.Choice([1, 2])) +@click.option("--model_root", default=PPL_CONFIG["model_root"], help="Dense 1.3B checkpoint root") +@click.option("--prompt", default=PPL_CONFIG["prompt"], help="Structured JSON caption") +@click.option("--negative_prompt", default=None, help="Optional structured negative caption") +@click.option("--seed", default=PPL_CONFIG["seed"], type=int) +@click.option("--resolution", default=PPL_CONFIG["resolution"]) +@click.option("--aspect_ratio", default=PPL_CONFIG["aspect_ratio"]) +@click.option("--target_video_length", default=None, type=float) +@click.option("--task", default="t2v", type=click.Choice(["t2i", "t2v", "i2v"])) +@click.option("--first_image_path", default="", help="Required for i2v") +@click.option("--output_path", default="output.mp4") +def main( + gpu_num: int, + cfg_parallel_degree: int, + model_root: str, + prompt: str, + negative_prompt: str | None, + seed: int, + resolution: str, + aspect_ratio: str, + target_video_length: float | None, + task: str, + first_image_path: str, + output_path: str, +) -> None: + """Generate with the LingBot-Video Dense 1.3B checkpoint.""" + pipeline = get_pipeline(gpu_num, model_root, cfg_parallel_degree) + try: + result = run_with_file( + pipeline, + prompt, + negative_prompt, + seed, + output_path, + target_video_length, + resolution, + aspect_ratio, + task, + first_image_path, + ) + click.echo(f"Output saved to {result['output_path']}") + finally: + stop = getattr(pipeline, "stop", None) + if callable(stop): + stop() + + +if __name__ == "__main__": + main() diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py new file mode 100644 index 0000000..09f134a --- /dev/null +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -0,0 +1,619 @@ +"""LingBot-Video MoE 30B plus refiner CLI and service example.""" + +from __future__ import annotations + +import json +import os +from pathlib import Path +from typing import Any + +import click +import numpy as np +import torch +from PIL import Image +from diffusers.utils import export_to_video + +from telefuser.core.config import ( + AttentionConfig, + AttnImplType, + ModelRuntimeConfig, + OffloadConfig, + ParallelConfig, + WeightOffloadType, +) +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video import ( + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_NEGATIVE_PROMPT_IMAGE, + LingBotVideoDenoisingStage, + LingBotVideoPipeline, + LingBotVideoPipelineConfig, + LingBotVideoRefinerStage, + LingBotVideoRequest, + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + default_negative_caption, + load_lingbot_video_model_config, + load_refiner_first_frame, + num_frames_from_duration, + parse_lingbot_video_prompt, + prepare_refiner_video, +) +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler +from telefuser.service.core.contract_templates import build_pipeline_manifest, build_task_contract_template +from telefuser.utils.video import get_target_video_size_from_ratio +from telefuser.worker.parallel_worker import ParallelWorker + +_DEFAULT_PROMPT = (Path(__file__).parent / "assets" / "t2v_5s.json.example").read_text(encoding="utf-8") +TF_MODEL_ZOO_PATH = os.environ.get("TF_MODEL_ZOO_PATH", "/hhb-data/aigc/model_zoo") + +PPL_CONFIG: dict[str, Any] = { + "name": "lingbot_video_moe_30b", + "model_root": TF_MODEL_ZOO_PATH + "/lingbot/lingbot-video-moe-30b-a3b", + "variant": "moe", + "supports_refiner": True, + "prompt": _DEFAULT_PROMPT, + "num_inference_steps": 40, + "guidance_scale": 3.0, + "fps": 24, + "target_video_length": 5, + "resolution": "480p", + "aspect_ratio": "16:9", + "seed": 42, + "expert_backend": "auto", + "enable_refiner": True, + "refiner_height": 1088, + "refiner_width": 1920, + "refiner_steps": 8, + "refiner_guidance_scale": 3.0, + "refiner_shift": 3.0, + "refiner_t_thresh": 0.85, + "refiner_tail_steps": 2, + "refiner_parallelism": 0, + "refiner_batch_cfg": False, + "refiner_null_cond_clone_zero": True, + "cfg_parallel_degree": 1, + "refiner_cfg_parallel_degree": 1, + "refiner_co_resident": True, +} + + +def _build_contract() -> dict[str, Any]: + """Build the standard service contract with LingBot-specific refiner defaults.""" + task_contracts = {} + for task in ("t2i", "t2v", "i2v"): + overrides = { + "negative_prompt": {"default": DEFAULT_NEGATIVE_PROMPT_IMAGE if task == "t2i" else DEFAULT_NEGATIVE_PROMPT}, + "seed": {"default": PPL_CONFIG["seed"]}, + "resolution": {"default": PPL_CONFIG["resolution"]}, + "aspect_ratio": {"default": PPL_CONFIG["aspect_ratio"]}, + } + if task != "t2i": + overrides["target_video_length"] = {"default": PPL_CONFIG["target_video_length"]} + overrides["refine"] = {"type": "boolean", "default": PPL_CONFIG["enable_refiner"]} + task_contracts[task] = build_task_contract_template(task, parameter_overrides=overrides) + return build_pipeline_manifest( + pipeline_name=PPL_CONFIG["name"], + supported_tasks=task_contracts, + task_contracts=task_contracts, + ) + + +PIPELINE_CONTRACT = _build_contract() + + +def build_pipeline( + model_root: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + cpu_offload: bool = True, + guidance_scale: float = PPL_CONFIG["guidance_scale"], + num_inference_steps: int = PPL_CONFIG["num_inference_steps"], + shift: float = 3.0, + attention_config: AttentionConfig | None = None, + parallel_config: ParallelConfig | None = None, + batch_cfg: bool | None = None, + expert_backend: str = PPL_CONFIG["expert_backend"], +) -> LingBotVideoPipeline: + """Load MoE base modules and initialize the complete pipeline.""" + try: + from diffusers import AutoencoderKLWan + from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + except ImportError as exc: + raise RuntimeError("LingBot-Video requires diffusers and transformers") from exc + + parallel_config = parallel_config or ParallelConfig() + if batch_cfg is None: + batch_cfg = parallel_config.world_size > 1 and parallel_config.cfg_degree == 1 + if batch_cfg and parallel_config.cfg_degree > 1: + raise ValueError("LingBot CFG parallel and batch CFG are mutually exclusive") + if parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot FSDP inference requires cpu_offload=False") + runtime_device = torch.device(device) + root = Path(model_root) + transformer_dir = root / "transformer" + offload_config = OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD + ) + runtime_config = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch_dtype, + attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + offload_config=offload_config, + parallel_config=parallel_config, + ) + + if expert_backend == "auto": + expert_backend = "grouped_mm" if parallel_config.world_size > 1 else "sorted" + if expert_backend not in {"fp8", "grouped_mm", "sorted"}: + raise ValueError("LingBot MoE expert_backend must be 'auto', 'fp8', 'grouped_mm', or 'sorted'") + if expert_backend == "grouped_mm" and not hasattr(torch, "_grouped_mm"): + raise RuntimeError("LingBot grouped_mm requires a recent CUDA-enabled PyTorch build") + module_manager = ModuleManager(torch_dtype=torch_dtype, device="cpu") + module_manager.load_model(str(transformer_dir), name="transformer", torch_dtype=torch_dtype) + transformer = module_manager.fetch_module("transformer") + if transformer is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {transformer_dir}") + transformer.promote_stability_layers_to_fp32() + if expert_backend == "fp8": + if not hasattr(torch, "_scaled_mm"): + raise RuntimeError("LingBot FP8 experts require torch._scaled_mm") + transformer.quantize_experts_fp8_() + transformer.set_expert_execution_backend(expert_backend) + + module_manager.load_from_huggingface( + str(root / "processor"), + module_source="transformers", + module_name="processor", + module_class=AutoProcessor, + ) + module_manager.load_from_huggingface( + str(root / "text_encoder"), + module_source="transformers", + module_name="text_encoder", + module_class=Qwen3VLForConditionalGeneration, + torch_dtype=torch_dtype, + attn_implementation="sdpa", + ) + module_manager.load_from_huggingface( + str(root / "vae"), + module_source="diffusers", + module_name="vae", + module_class=AutoencoderKLWan, + torch_dtype=torch.float32, + ) + module_manager.load_from_huggingface( + str(root / "scheduler"), + module_source="diffusers", + module_name="scheduler", + module_class=FlowUniPCMultistepScheduler, + ) + + pipeline = LingBotVideoPipeline(device=device, torch_dtype=torch_dtype) + pipeline.init( + module_manager, + LingBotVideoPipelineConfig( + model=load_lingbot_video_model_config(transformer_dir, variant="moe"), + text_encoding_config=runtime_config, + dit_config=runtime_config, + vae_config=runtime_config, + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + shift=shift, + batch_cfg=batch_cfg, + enable_denoising_parallel=parallel_config.world_size > 1, + ), + ) + return pipeline + + +def build_refiner( + model_root: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + cpu_offload: bool = True, + attention_config: AttentionConfig | None = None, + parallel_config: ParallelConfig | None = None, + batch_cfg: bool | None = None, + expert_backend: str = PPL_CONFIG["expert_backend"], +) -> LingBotVideoRefinerStage: + """Load refiner modules through ModuleManager and assemble its stages.""" + try: + from diffusers import AutoencoderKLWan + except ImportError as exc: + raise RuntimeError("LingBot-Video refiner requires diffusers") from exc + + parallel_config = parallel_config or ParallelConfig() + if batch_cfg is None: + batch_cfg = False + if batch_cfg and parallel_config.cfg_degree > 1: + raise ValueError("LingBot CFG parallel and batch CFG are mutually exclusive") + if parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot refiner FSDP inference requires cpu_offload=False") + runtime_device = torch.device(device) + root = Path(model_root) + offload_config = OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD + ) + runtime_config = ModelRuntimeConfig( + device_type=runtime_device.type, + device_id=runtime_device.index or 0, + torch_dtype=torch_dtype, + attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + offload_config=offload_config, + parallel_config=parallel_config, + ) + + if expert_backend == "auto": + expert_backend = "grouped_mm" if parallel_config.world_size > 1 else "sorted" + if expert_backend not in {"fp8", "grouped_mm", "sorted"}: + raise ValueError("LingBot refiner expert_backend must be 'auto', 'fp8', 'grouped_mm', or 'sorted'") + module_manager = ModuleManager(torch_dtype=torch_dtype, device="cpu") + module_manager.load_model(str(root / "refiner"), name="transformer", torch_dtype=torch_dtype) + transformer = module_manager.fetch_module("transformer") + if transformer is None: + raise RuntimeError(f"Unable to load LingBot-Video refiner from {root / 'refiner'}") + transformer.promote_stability_layers_to_fp32() + if expert_backend == "fp8": + if not hasattr(torch, "_scaled_mm"): + raise RuntimeError("LingBot FP8 experts require torch._scaled_mm") + transformer.quantize_experts_fp8_() + transformer.set_expert_execution_backend(expert_backend) + + module_manager.load_from_huggingface( + str(root / "vae"), + module_source="diffusers", + module_name="vae", + module_class=AutoencoderKLWan, + torch_dtype=torch.float32, + ) + module_manager.load_from_huggingface( + str(root / "scheduler"), + module_source="diffusers", + module_name="scheduler", + module_class=FlowUniPCMultistepScheduler, + ) + scheduler = module_manager.fetch_module("scheduler") + denoising_stage = LingBotVideoDenoisingStage("refiner", module_manager, runtime_config, batch_cfg=batch_cfg) + refiner = LingBotVideoRefinerStage( + denoising_stage=denoising_stage, + vae_encode_stage=LingBotVideoVAEEncodeStage("refiner_vae_encode", module_manager, runtime_config), + vae_decode_stage=LingBotVideoVAEDecodeStage("refiner_vae_decode", module_manager, runtime_config), + scheduler=scheduler, + ) + if parallel_config.world_size > 1 and not torch.distributed.is_initialized(): + refiner.denoising_stage = ParallelWorker(denoising_stage) + else: + denoising_stage.parallel_models() + return refiner + + +def get_pipeline( + parallelism: int = 1, + model_root: str = PPL_CONFIG["model_root"], + refiner_parallelism: int | None = None, + refiner_batch_cfg: bool | None = None, + cfg_parallel_degree: int | None = None, + refiner_cfg_parallel_degree: int | None = None, + refiner_co_resident: bool | None = None, + expert_backend: str = PPL_CONFIG["expert_backend"], +) -> object: + """Load the fixed MoE 30B checkpoint and configure its separate refiner.""" + if parallelism not in {1, 4}: + raise ValueError("LingBot-Video MoE supports parallelism=1 or 4") + cfg_parallel_degree = int(PPL_CONFIG["cfg_parallel_degree"]) if cfg_parallel_degree is None else cfg_parallel_degree + if cfg_parallel_degree not in {1, 2} or parallelism % cfg_parallel_degree: + raise ValueError("LingBot CFG parallel degree must be 1 or 2 and divide parallelism") + pipeline = build_pipeline( + model_root, + cpu_offload=parallelism == 1, + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)), + cfg_degree=cfg_parallel_degree, + sp_ulysses_degree=parallelism // cfg_parallel_degree, + enable_fsdp=parallelism > 1, + ), + expert_backend=expert_backend, + ) + refiner_parallelism = ( + int(PPL_CONFIG["refiner_parallelism"]) if refiner_parallelism is None else refiner_parallelism + ) or parallelism + if refiner_parallelism not in {1, 4}: + raise ValueError("LingBot-Video refiner supports parallelism=1 or 4") + refiner_cfg_parallel_degree = ( + int(PPL_CONFIG["refiner_cfg_parallel_degree"]) + if refiner_cfg_parallel_degree is None + else refiner_cfg_parallel_degree + ) + if refiner_cfg_parallel_degree not in {1, 2} or refiner_parallelism % refiner_cfg_parallel_degree: + raise ValueError("LingBot refiner CFG parallel degree must be 1 or 2 and divide refiner parallelism") + pipeline.refiner_parallel_config = ParallelConfig( + device_ids=list(range(refiner_parallelism)), + cfg_degree=refiner_cfg_parallel_degree, + sp_ulysses_degree=refiner_parallelism // refiner_cfg_parallel_degree, + enable_fsdp=refiner_parallelism > 1, + ) + pipeline.refiner_cpu_offload = refiner_parallelism == 1 + pipeline.refiner_batch_cfg = ( + bool(PPL_CONFIG["refiner_batch_cfg"]) if refiner_batch_cfg is None else refiner_batch_cfg + ) + if pipeline.refiner_batch_cfg and refiner_cfg_parallel_degree > 1: + raise ValueError("LingBot refiner CFG parallel and batch CFG are mutually exclusive") + pipeline.refiner_co_resident = ( + bool(PPL_CONFIG["refiner_co_resident"]) if refiner_co_resident is None else refiner_co_resident + ) and parallelism > 1 + pipeline.refiner_expert_backend = expert_backend + return pipeline + + +def _load_condition_image(path: str) -> torch.Tensor: + """Load an RGB image path as a raw [0,255] TI2V tensor.""" + pixels = np.asarray(Image.open(path).convert("RGB")).copy() + return torch.from_numpy(pixels).permute(2, 0, 1).unsqueeze(0).float() + + +def _resolve_video_size(aspect_ratio: str, resolution: str) -> tuple[int, int]: + """Resolve output geometry using the validated LingBot 480p landscape preset.""" + if resolution == "480p" and aspect_ratio == "16:9": + return 832, 480 + width, height = get_target_video_size_from_ratio( + aspect_ratio, + resolution, + height_division_factor=16, + width_division_factor=16, + ) + if width is None or height is None: + raise ValueError(f"unsupported LingBot-Video resolution: {resolution}") + return width, height + + +def _generator(pipeline: object, seed: int) -> torch.Generator: + """Create a deterministic generator on the pipeline device type.""" + device_type = torch.device(getattr(pipeline, "device", "cpu")).type + return torch.Generator(device_type).manual_seed(seed) + + +def run( + pipeline: object, + prompt: str = PPL_CONFIG["prompt"], + negative_prompt: str | None = None, + seed: int = PPL_CONFIG["seed"], + resolution: str = PPL_CONFIG["resolution"], + aspect_ratio: str = PPL_CONFIG["aspect_ratio"], + target_video_length: float | None = None, + task: str = "t2v", + first_image_path: str = "", + refine: bool | None = None, + model_root: str = PPL_CONFIG["model_root"], +) -> torch.Tensor: + """Run MoE T2I, T2V, or I2V with an optional refiner stage.""" + if task not in {"t2i", "t2v", "i2v"}: + raise ValueError(f"unsupported LingBot-Video task: {task}") + if task == "i2v" and not first_image_path: + raise ValueError("LingBot-Video i2v requires first_image_path") + if task == "t2i" and refine is True: + raise ValueError("LingBot-Video refiner does not support t2i") + if refine is None: + refine = task != "t2i" and bool(PPL_CONFIG["enable_refiner"]) + + caption, prompt_duration = parse_lingbot_video_prompt(json.loads(prompt)) + width, height = _resolve_video_size(aspect_ratio, resolution) + duration = target_video_length if target_video_length is not None else prompt_duration + if duration is None: + duration = float(PPL_CONFIG["target_video_length"]) + num_frames = 1 if task == "t2i" else num_frames_from_duration(duration, fps=int(PPL_CONFIG["fps"])) + resolved_negative_prompt = negative_prompt if negative_prompt is not None else default_negative_caption(num_frames) + generation = pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=height, + width=width, + num_frames=num_frames, + image=_load_condition_image(first_image_path) if first_image_path else None, + ), + negative_caption=resolved_negative_prompt, + generator=_generator(pipeline, seed), + ) + frames = generation.output + if not refine: + return frames + if pipeline.text_stage is None: + raise RuntimeError("LingBot-Video refiner requires the base text stage") + + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = pipeline.text_stage.encode(caption) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + if PPL_CONFIG["refiner_null_cond_clone_zero"]: + negative = torch.zeros_like(positive) + negative_mask = positive_mask.clone() + elif ( + generation.prompt_conditions.has_visual_condition or generation.prompt_conditions.negative_prompt_embeds is None + ): + negative, negative_mask = pipeline.text_stage.encode(resolved_negative_prompt) + else: + negative = generation.prompt_conditions.negative_prompt_embeds + negative_mask = generation.prompt_conditions.negative_attention_mask + + if not getattr(pipeline, "refiner_co_resident", False): + pipeline.release_gpu_resources() + refiner = build_refiner( + model_root, + cpu_offload=getattr(pipeline, "refiner_cpu_offload", True), + parallel_config=getattr(pipeline, "refiner_parallel_config", None), + batch_cfg=getattr(pipeline, "refiner_batch_cfg", False), + expert_backend=getattr(pipeline, "refiner_expert_backend", PPL_CONFIG["expert_backend"]), + ) + try: + lowres_video, _ = prepare_refiner_video( + frames, + source_fps=float(PPL_CONFIG["fps"]), + height=int(PPL_CONFIG["refiner_height"]), + width=int(PPL_CONFIG["refiner_width"]), + ) + clean_first_frame = ( + load_refiner_first_frame( + first_image_path, + target_height=int(PPL_CONFIG["refiner_height"]), + target_width=int(PPL_CONFIG["refiner_width"]), + geometry_height=height, + geometry_width=width, + ) + if first_image_path + else None + ) + return refiner.refine( + lowres_video, + positive, + negative, + positive_mask, + negative_mask, + num_inference_steps=int(PPL_CONFIG["refiner_steps"]), + guidance_scale=float(PPL_CONFIG["refiner_guidance_scale"]), + shift=float(PPL_CONFIG["refiner_shift"]), + t_thresh=float(PPL_CONFIG["refiner_t_thresh"]), + tail_steps=int(PPL_CONFIG["refiner_tail_steps"]), + clean_first_frame=clean_first_frame, + generator=_generator(pipeline, seed), + before_decode=(pipeline.release_gpu_resources if getattr(pipeline, "refiner_co_resident", False) else None), + ) + finally: + refiner.close() + + +def _save_output(frames: torch.Tensor, output_path: str) -> dict[str, str]: + """Encode normalized RGB without an intermediate uint8 video conversion.""" + path = Path(output_path) + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + if video.shape[0] == 1: + if path.suffix.lower() not in {".png", ".jpg", ".jpeg", ".webp", ".bmp"}: + path = path.with_suffix(".png") + path.parent.mkdir(parents=True, exist_ok=True) + Image.fromarray((video[0] * 255).round().astype(np.uint8)).save(path) + else: + path.parent.mkdir(parents=True, exist_ok=True) + export_to_video(list(video), str(path), fps=int(PPL_CONFIG["fps"])) + return {"output_path": str(path)} + + +def run_with_file( + pipeline: object, + prompt: str = PPL_CONFIG["prompt"], + negative_prompt: str | None = None, + seed: int = PPL_CONFIG["seed"], + output_path: str = "output.mp4", + target_video_length: float | None = None, + resolution: str = PPL_CONFIG["resolution"], + aspect_ratio: str = PPL_CONFIG["aspect_ratio"], + task: str = "t2v", + first_image_path: str = "", + refine: bool | None = None, + model_root: str = PPL_CONFIG["model_root"], + **_: object, +) -> dict[str, str]: + """Run the MoE example and save its base or refined output.""" + frames = run( + pipeline, + prompt, + negative_prompt, + seed, + resolution, + aspect_ratio, + target_video_length, + task, + first_image_path, + refine, + model_root=model_root, + ) + return _save_output(frames, output_path) + + +@click.command() +@click.option("--gpu_num", default=1, type=int, help="Number of GPUs: 1 or 4") +@click.option("--cfg_parallel_degree", default=PPL_CONFIG["cfg_parallel_degree"], type=click.Choice([1, 2])) +@click.option("--model_root", default=PPL_CONFIG["model_root"], help="MoE 30B checkpoint root") +@click.option("--prompt", default=PPL_CONFIG["prompt"], help="Structured JSON caption") +@click.option("--negative_prompt", default=None, help="Optional structured negative caption") +@click.option("--seed", default=PPL_CONFIG["seed"], type=int) +@click.option( + "--expert_backend", + default=PPL_CONFIG["expert_backend"], + type=click.Choice(["auto", "fp8", "grouped_mm", "sorted"]), + help="MoE expert backend; fp8 uses native dynamic W8A8 scaled GEMMs", +) +@click.option("--resolution", default=PPL_CONFIG["resolution"]) +@click.option("--aspect_ratio", default=PPL_CONFIG["aspect_ratio"]) +@click.option("--target_video_length", default=None, type=float) +@click.option("--task", default="t2v", type=click.Choice(["t2i", "t2v", "i2v"])) +@click.option("--first_image_path", default="", help="Required for i2v") +@click.option( + "--refiner_gpu_num", default=PPL_CONFIG["refiner_parallelism"], type=int, help="Refiner GPUs; 0 inherits gpu_num" +) +@click.option("--refiner_batch_cfg/--no-refiner_batch_cfg", default=PPL_CONFIG["refiner_batch_cfg"]) +@click.option( + "--refiner_cfg_parallel_degree", + default=PPL_CONFIG["refiner_cfg_parallel_degree"], + type=click.Choice([1, 2]), +) +@click.option("--refiner_co_resident/--no-refiner_co_resident", default=PPL_CONFIG["refiner_co_resident"]) +@click.option("--refine/--no-refine", default=None) +@click.option("--output_path", default="output.mp4") +def main( + gpu_num: int, + cfg_parallel_degree: int, + model_root: str, + prompt: str, + negative_prompt: str | None, + seed: int, + expert_backend: str, + resolution: str, + aspect_ratio: str, + target_video_length: float | None, + task: str, + first_image_path: str, + refiner_gpu_num: int, + refiner_batch_cfg: bool, + refiner_cfg_parallel_degree: int, + refiner_co_resident: bool, + refine: bool | None, + output_path: str, +) -> None: + """Generate with the LingBot-Video MoE 30B checkpoint and refiner.""" + pipeline = get_pipeline( + gpu_num, + model_root, + refiner_parallelism=refiner_gpu_num, + refiner_batch_cfg=refiner_batch_cfg, + cfg_parallel_degree=cfg_parallel_degree, + refiner_cfg_parallel_degree=refiner_cfg_parallel_degree, + refiner_co_resident=refiner_co_resident, + expert_backend=expert_backend, + ) + try: + result = run_with_file( + pipeline, + prompt, + negative_prompt, + seed, + output_path, + target_video_length, + resolution, + aspect_ratio, + task, + first_image_path, + refine, + model_root=model_root, + ) + click.echo(f"Output saved to {result['output_path']}") + finally: + stop = getattr(pipeline, "stop", None) + if callable(stop): + stop() + + +if __name__ == "__main__": + main() diff --git a/examples/wan_video/wan21_14b_image_to_video_h100.py b/examples/wan_video/wan21_14b_image_to_video_h100.py index 95b6037..2dff26c 100755 --- a/examples/wan_video/wan21_14b_image_to_video_h100.py +++ b/examples/wan_video/wan21_14b_image_to_video_h100.py @@ -198,18 +198,19 @@ def main( image = Image.open(image_path).convert("RGB") # Run inference - start = time.time() - video = run( - pipe, - image, - prompt, - negative_prompt, - seed=seed, - resolution=resolution, - ) - elapsed_time = time.time() - start - - print(f"Video generation time: {elapsed_time:.2f} seconds") + while True: + start = time.time() + video = run( + pipe, + image, + prompt, + negative_prompt, + seed=seed, + resolution=resolution, + ) + elapsed_time = time.time() - start + + print(f"Video generation time: {elapsed_time:.2f} seconds") # Save results output_dir = os.getenv("TELEAI_EXAMPLE_OUTPUT_DIR", "./") diff --git a/telefuser/core/config.py b/telefuser/core/config.py index 04f6e99..a66a8b5 100644 --- a/telefuser/core/config.py +++ b/telefuser/core/config.py @@ -46,6 +46,7 @@ class ParallelConfig: enable_fsdp: bool = False timeout: int = 600 # Seconds queue_with_cpu: bool = False + worker_intra_op_threads: int = 1 # Per-process PyTorch CPU intra-op threads @property def world_size(self) -> int: @@ -57,6 +58,8 @@ def world_size(self) -> int: def validate(self) -> None: """Validate that device count matches parallelism degrees.""" + if self.worker_intra_op_threads < 1: + raise ValueError("worker_intra_op_threads must be positive") device_num = 1 if self.device_ids is None else len(self.device_ids) degree_sum = ( self.cfg_degree diff --git a/telefuser/core/module_manager.py b/telefuser/core/module_manager.py index 2bb39d4..0659754 100755 --- a/telefuser/core/module_manager.py +++ b/telefuser/core/module_manager.py @@ -3,6 +3,7 @@ from __future__ import annotations import glob +import json import os from typing import Any @@ -226,6 +227,19 @@ def load_model( device = device or self.device torch_dtype = torch_dtype or self.torch_dtype + # Resolve Diffusers-style checkpoint directories to their safetensors + # payloads. Sharded checkpoints are merged below, just like an + # explicitly supplied list of checkpoint files. + if isinstance(file_path, str) and os.path.isdir(file_path): + index_path = os.path.join(file_path, "diffusion_pytorch_model.safetensors.index.json") + single_file = os.path.join(file_path, "diffusion_pytorch_model.safetensors") + if os.path.isfile(index_path): + with open(index_path, encoding="utf-8") as handle: + weight_map = json.load(handle)["weight_map"] + file_path = [os.path.join(file_path, shard) for shard in sorted(set(weight_map.values()))] + elif os.path.isfile(single_file): + file_path = single_file + # Expand wildcards if present if isinstance(file_path, str) and ("*" in file_path or "?" in file_path): expanded_paths = sorted(glob.glob(file_path)) diff --git a/telefuser/distributed/fsdp.py b/telefuser/distributed/fsdp.py index 866e0dd..69ec0a4 100644 --- a/telefuser/distributed/fsdp.py +++ b/telefuser/distributed/fsdp.py @@ -6,6 +6,7 @@ from __future__ import annotations +from collections.abc import Iterable from functools import partial import torch @@ -26,6 +27,7 @@ def shard_model( reduce_dtype: torch.dtype = torch.bfloat16, buffer_dtype: torch.dtype = torch.bfloat16, cpu_offload: bool = False, + ignored_states: Iterable[nn.Parameter] | None = None, ) -> FSDP: """Shard model using FSDP1. @@ -38,6 +40,7 @@ def shard_model( reduce_dtype: Gradient reduction dtype buffer_dtype: Buffer dtype cpu_offload: Whether to offload parameters to CPU + ignored_states: Parameters to retain as unsharded, replicated state. Returns: FSDP-wrapped module @@ -70,6 +73,7 @@ def wrap_fn(m: nn.Module) -> bool: forward_prefetch=True, auto_wrap_policy=partial(lambda_auto_wrap_policy, lambda_fn=wrap_fn), cpu_offload=CPUOffload(offload_params=True) if cpu_offload else None, + ignored_states=ignored_states, ) @@ -142,3 +146,33 @@ def _apply_fsdp_recursive(current_module: nn.Module, module_path: str = "") -> N fully_shard(module, mesh=device_mesh, mp_policy=mp_policy) return module + + +def shard_model_fsdp2_inference( + module: nn.Module, + device_mesh: DeviceMesh, + wrap_module_names: list[str], + ignored_states: Iterable[nn.Parameter] | None = None, +) -> nn.Module: + """Apply source-style composable FSDP2 for inference-only model sharding.""" + ignored_params = set(ignored_states or ()) + if torch.cuda.is_available(): + device = torch.device("cuda", torch.cuda.current_device()) + for submodule in module.modules(): + for name, buffer in tuple(submodule.named_buffers(recurse=False)): + if buffer is not None and buffer.device != device: + submodule._buffers[name] = buffer.to(device=device) + for parameter in ignored_params: + if parameter.device != device: + parameter.data = parameter.data.to(device=device) + + wrapped: set[nn.Module] = set() + for name in wrap_module_names: + target = getattr(module, name) + targets = target if isinstance(target, nn.ModuleList) else (target,) + for child in targets: + child_ignored_params = {parameter for parameter in child.parameters() if parameter in ignored_params} + fully_shard(child, mesh=device_mesh, ignored_params=child_ignored_params) + wrapped.add(child) + fully_shard(module, mesh=device_mesh, ignored_params=ignored_params) + return module diff --git a/telefuser/distributed/ulysses_comm.py b/telefuser/distributed/ulysses_comm.py index faf5833..28f3643 100644 --- a/telefuser/distributed/ulysses_comm.py +++ b/telefuser/distributed/ulysses_comm.py @@ -35,6 +35,21 @@ def _wait_async_tensor(tensor: torch.Tensor) -> torch.Tensor: return tensor +def ulysses_all_to_all_split_cat( + tensor: torch.Tensor, + process_group: dist.ProcessGroup, + *, + scatter_dim: int, + gather_dim: int, +) -> torch.Tensor: + """Run the source-style synchronous list all-to-all used by LingBot Ulysses.""" + _, world_size = _get_distributed_info(process_group) + inputs = [part.contiguous() for part in torch.tensor_split(tensor, world_size, scatter_dim)] + outputs = [torch.empty_like(inputs[0]) for _ in range(world_size)] + dist.all_to_all(outputs, inputs, group=process_group) + return torch.cat(outputs, dim=gather_dim).contiguous() + + def ulysses_scatter_heads( tensor: torch.Tensor, process_group: dist.ProcessGroup, diff --git a/telefuser/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py new file mode 100644 index 0000000..9af27de --- /dev/null +++ b/telefuser/models/lingbot_video_dit.py @@ -0,0 +1,560 @@ +"""Checkpoint-compatible Dense LingBot-Video transformer modules. + +Numerical behavior is adapted from the Apache-2.0 licensed upstream +LingBot-Video transformer implementation. +""" + +from __future__ import annotations + +from math import prod +from typing import Any + +import torch +import torch.distributed as dist +import torch.nn.functional as F +from torch import nn + +from telefuser.core.model_registry import register_model_config +from telefuser.distributed.ulysses_comm import ulysses_all_to_all_split_cat +from telefuser.ops.attention import attention +from telefuser.utils.model_weight import hash_state_dict_keys + + +class LingBotVideoPatchEmbed(nn.Module): + """Patchify and restore ``[B,C,F,H,W]`` latent tensors.""" + + def __init__(self, config: Any) -> None: + super().__init__() + self.config = config + self.projection = nn.Linear(config.in_channels * prod(config.patch_size), config.hidden_size) + + def patchify(self, latent: torch.Tensor) -> torch.Tensor: + if latent.ndim != 5: + raise ValueError("LingBot latent must have shape [B,C,F,H,W]") + _, channels, frames, height, width = latent.shape + pt, ph, pw = self.config.patch_size + if channels != self.config.in_channels or frames % pt or height % ph or width % pw: + raise ValueError("latent shape is incompatible with LingBot patch size") + latent = latent.reshape(latent.shape[0], channels, frames // pt, pt, height // ph, ph, width // pw, pw) + latent = latent.permute(0, 2, 4, 6, 3, 5, 7, 1) + return latent.reshape(latent.shape[0], -1, pt * ph * pw * channels) + + def forward(self, latent: torch.Tensor) -> torch.Tensor: + """Return projected patch tokens with shape ``[B,N,hidden_size]``.""" + return self.projection(self.patchify(latent)) + + def unpatchify(self, tokens: torch.Tensor, *, frames: int, height: int, width: int) -> torch.Tensor: + if tokens.ndim != 3: + raise ValueError("tokens must have shape [B,N,D]") + pt, ph, pw = self.config.patch_size + grid = (frames // pt, height // ph, width // pw) + expected = grid[0] * grid[1] * grid[2] + if tokens.shape[1] != expected: + raise ValueError(f"expected {expected} patch tokens, got {tokens.shape[1]}") + values = tokens.reshape(tokens.shape[0], *grid, pt, ph, pw, self.config.in_channels) + values = values.permute(0, 7, 1, 4, 2, 5, 3, 6) + return values.reshape(tokens.shape[0], self.config.in_channels, frames, height, width) + + +class LingBotVideoRMSNorm(nn.Module): + """Checkpoint-compatible RMSNorm with fp32 accumulation.""" + + def __init__(self, dim: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(dim)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + normalized = hidden_states.float() + normalized = normalized * torch.rsqrt(normalized.pow(2).mean(-1, keepdim=True) + self.variance_epsilon) + return (self.weight * normalized).to(input_dtype) + + +class LingBotVideoMLP(nn.Module): + """Checkpoint-compatible SwiGLU MLP.""" + + def __init__(self, hidden_size: int, intermediate_size: int) -> None: + super().__init__() + self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) + self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down_proj(torch.nn.functional.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + + +def apply_lingbot_video_complex_rope(hidden_states: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: + """Apply the official complex64 RoPE representation to ``[B,S,H,D]`` tensors.""" + if hidden_states.ndim != 4 or hidden_states.shape[-1] % 2: + raise ValueError("RoPE inputs must have shape [B,S,H,even_head_dim]") + if freqs_cis.ndim == 2: + freqs_cis = freqs_cis.unsqueeze(0) + if freqs_cis.shape[-1] != hidden_states.shape[-1] // 2: + raise ValueError("RoPE table does not match attention head dimension") + complex_states = torch.view_as_complex(hidden_states.float().reshape(*hidden_states.shape[:-1], -1, 2)) + output = torch.view_as_real(complex_states * freqs_cis.unsqueeze(2)).flatten(3) + return output.to(hidden_states.dtype) + + +class LingBotVideoTextEmbedder(nn.Module): + """Checkpoint-compatible text feature projection.""" + + def __init__(self, text_dim: int, hidden_size: int) -> None: + super().__init__() + self.norm = LingBotVideoRMSNorm(text_dim, eps=1e-6) + self.linear_1 = nn.Linear(text_dim, hidden_size, bias=True) + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=True) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.linear_2(F.silu(self.linear_1(self.norm(hidden_states)))) + + +class LingBotVideoAttention(nn.Module): + """Native-SDPA equivalent of the official LingBot attention module.""" + + def __init__(self, hidden_size: int, num_heads: int, norm_eps: float, qkv_bias: bool, out_bias: bool) -> None: + super().__init__() + if hidden_size % num_heads: + raise ValueError("hidden_size must be divisible by num_heads") + self.num_heads = num_heads + self.head_dim = hidden_size // num_heads + self.to_q = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.to_k = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.to_v = nn.Linear(hidden_size, hidden_size, bias=qkv_bias) + self.norm_q = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.norm_k = LingBotVideoRMSNorm(self.head_dim, norm_eps) + self.to_out = nn.Linear(hidden_size, hidden_size, bias=out_bias) + self.attention_config: object | None = None + self.ulysses_group: dist.ProcessGroup | None = None + + def set_attention_config(self, attention_config: object) -> None: + """Attach a runtime-selected TeleFuser attention implementation.""" + self.attention_config = attention_config + + def set_ulysses_group(self, group: dist.ProcessGroup | None) -> None: + """Attach the Ulysses process group used for sequence-parallel attention.""" + self.ulysses_group = group + + def forward( + self, + hidden_states: torch.Tensor, + rotary_emb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + packed_sequence_lengths: list[int] | None = None, + ) -> torch.Tensor: + batch, sequence, _ = hidden_states.shape + query = self.to_q(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + key = self.to_k(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + value = self.to_v(hidden_states).view(batch, sequence, self.num_heads, self.head_dim) + query = apply_lingbot_video_complex_rope(self.norm_q(query), rotary_emb) + key = apply_lingbot_video_complex_rope(self.norm_k(key), rotary_emb) + group = self.ulysses_group + use_ulysses = ( + group is not None and dist.is_available() and dist.is_initialized() and dist.get_world_size(group) > 1 + ) + if use_ulysses: + world_size = dist.get_world_size(group) + local_heads = self.num_heads // world_size + query = ulysses_all_to_all_split_cat( + query.reshape(batch, sequence, self.num_heads * self.head_dim), + group, + scatter_dim=2, + gather_dim=1, + ).view(batch, sequence * world_size, local_heads, self.head_dim) + key = ulysses_all_to_all_split_cat( + key.reshape(batch, sequence, self.num_heads * self.head_dim), + group, + scatter_dim=2, + gather_dim=1, + ).view(batch, sequence * world_size, local_heads, self.head_dim) + value = ulysses_all_to_all_split_cat( + value.reshape(batch, sequence, self.num_heads * self.head_dim), + group, + scatter_dim=2, + gather_dim=1, + ).view(batch, sequence * world_size, local_heads, self.head_dim) + output = attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attention_config=self.attention_config, + attn_mask=attention_mask, + sequence_lengths=packed_sequence_lengths, + input_layout="BNSD", + output_layout="BNSD", + ) + if not isinstance(output, torch.Tensor): + raise RuntimeError("LingBot attention does not support log-sum-exp outputs") + output = output.transpose(1, 2) + if use_ulysses: + output = ulysses_all_to_all_split_cat( + output.reshape(batch, sequence * world_size, local_heads * self.head_dim), + group, + scatter_dim=1, + gather_dim=2, + ).view(batch, sequence, self.num_heads, self.head_dim) + return self.to_out(output.reshape(batch, sequence, -1).to(hidden_states.dtype)) + + +class LingBotVideoBlock(nn.Module): + """Checkpoint-compatible Dense LingBot transformer block.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + intermediate_size: int, + norm_eps: float = 1e-6, + qkv_bias: bool = False, + out_bias: bool = True, + ) -> None: + super().__init__() + self.scale_shift_table = nn.Parameter(torch.zeros(1, 6 * hidden_size)) + self.norm1 = LingBotVideoRMSNorm(hidden_size, norm_eps) + self.attn = LingBotVideoAttention(hidden_size, num_attention_heads, norm_eps, qkv_bias, out_bias) + self.norm_post_attn = LingBotVideoRMSNorm(hidden_size, norm_eps) + self.norm2 = LingBotVideoRMSNorm(hidden_size, norm_eps) + self.ffn = LingBotVideoMLP(hidden_size, intermediate_size) + self.norm_post_ffn = LingBotVideoRMSNorm(hidden_size, norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + temb6: torch.Tensor, + rotary_emb: torch.Tensor, + attention_mask: torch.Tensor | None = None, + packed_sequence_lengths: list[int] | None = None, + ) -> torch.Tensor: + batch, sequence, hidden_size = hidden_states.shape + if temb6.shape != (batch, sequence, 6 * hidden_size): + raise ValueError("temb6 must have shape [B,S,6*hidden_size]") + modulation = temb6 + self.scale_shift_table.unsqueeze(0) + shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = modulation.chunk(6, dim=-1) + attention_input = self.norm1(hidden_states) * (1.0 + scale_msa) + shift_msa + attention_output = self.attn( + attention_input.to(self.attn.to_q.weight.dtype), + rotary_emb, + attention_mask, + packed_sequence_lengths, + ) + hidden_states = hidden_states + (gate_msa.tanh() * self.norm_post_attn(attention_output)).to( + hidden_states.dtype + ) + mlp_input = self.norm2(hidden_states) * (1.0 + scale_mlp) + shift_mlp + ffn_weight = getattr(getattr(self.ffn, "gate_proj", None), "weight", self.attn.to_q.weight) + mlp_output = self.ffn(mlp_input.to(ffn_weight.dtype)) + return hidden_states + (gate_mlp.tanh() * self.norm_post_ffn(mlp_output)).to(hidden_states.dtype) + + +def make_lingbot_video_joint_position_ids( + text_len: int, grid_t: int, grid_h: int, grid_w: int, device: torch.device +) -> torch.Tensor: + """Create official [video; text] three-axis position IDs.""" + temporal = torch.arange(grid_t, device=device, dtype=torch.int32) + text_len + 1 + height = torch.arange(grid_h, device=device, dtype=torch.int32) + width = torch.arange(grid_w, device=device, dtype=torch.int32) + video = torch.stack(torch.meshgrid(temporal, height, width, indexing="ij"), dim=-1).flatten(0, 2) + text_t = torch.arange(text_len, device=device, dtype=torch.int32) + 1 + text = torch.stack((text_t, torch.zeros_like(text_t), torch.zeros_like(text_t)), dim=-1) + return torch.cat((video, text), dim=0) + + +def lingbot_video_complex_frequencies( + position_ids: torch.Tensor, axes_dims: tuple[int, int, int], theta: float +) -> torch.Tensor: + """Compute official multi-axis complex RoPE frequencies for position IDs.""" + if position_ids.ndim != 2 or position_ids.shape[1] != len(axes_dims): + raise ValueError("position_ids must have shape [S,3]") + position_ids_cpu = position_ids.detach().to(device="cpu") + frequencies = [] + for axis, dimension in enumerate(axes_dims): + values = torch.arange(0, dimension, 2, device="cpu", dtype=torch.float64) + values = 1.0 / (theta ** (values / dimension)) + angles = (position_ids_cpu[:, axis].to(torch.float64).unsqueeze(1) * values.unsqueeze(0)).float() + frequencies.append(torch.polar(torch.ones_like(angles), angles).to(torch.complex64)) + return torch.cat(frequencies, dim=-1).to(position_ids.device) + + +class LingBotVideoTimeEmbedder(nn.Module): + """Parameter names compatible with Diffusers TimestepEmbedding.""" + + def __init__(self, frequency_dim: int, hidden_size: int, bias: bool = True) -> None: + super().__init__() + self.frequency_dim = frequency_dim + self.linear_1 = nn.Linear(frequency_dim, hidden_size, bias=bias) + self.linear_2 = nn.Linear(hidden_size, hidden_size, bias=bias) + + def forward(self, timesteps: torch.Tensor) -> torch.Tensor: + half = self.frequency_dim // 2 + exponent = ( + -torch.log(torch.tensor(10000.0, device=timesteps.device)) + * torch.arange(half, device=timesteps.device, dtype=torch.float32) + / half + ) + angles = timesteps.float().reshape(-1, 1) * exponent.exp().reshape(1, -1) + embedding = torch.cat((angles.cos(), angles.sin()), dim=-1) + if self.frequency_dim % 2: + embedding = torch.nn.functional.pad(embedding, (0, 1)) + return self.linear_2(F.silu(self.linear_1(embedding.to(self.linear_1.weight.dtype)))) + + +class LingBotVideoTransformer3DModel(nn.Module): + """Source-equivalent Dense LingBot-Video transformer native reference path.""" + + def __init__( + self, + patch_size: tuple[int, int, int] = (1, 2, 2), + in_channels: int = 16, + out_channels: int = 16, + hidden_size: int = 2048, + num_attention_heads: int = 16, + depth: int = 24, + intermediate_size: int = 6144, + text_dim: int = 2560, + freq_dim: int = 256, + norm_eps: float = 1e-6, + rope_theta: float = 256.0, + axes_dims: tuple[int, int, int] = (32, 48, 48), + qkv_bias: bool = False, + out_bias: bool = True, + patch_embed_bias: bool = True, + timestep_mlp_bias: bool = True, + ) -> None: + super().__init__() + if hidden_size // num_attention_heads != sum(axes_dims): + raise ValueError("head dimension must equal sum(axes_dims)") + self.patch_size = patch_size + self.out_channels = out_channels + self.axes_dims = axes_dims + self.rope_theta = rope_theta + self.patch_embedder = nn.Linear(in_channels * prod(patch_size), hidden_size, bias=patch_embed_bias) + self.time_embedder = LingBotVideoTimeEmbedder(freq_dim, hidden_size, timestep_mlp_bias) + self.time_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 6 * hidden_size)) + self.text_embedder = LingBotVideoTextEmbedder(text_dim, hidden_size) + self.blocks = nn.ModuleList( + [ + LingBotVideoBlock(hidden_size, num_attention_heads, intermediate_size, norm_eps, qkv_bias, out_bias) + for _ in range(depth) + ] + ) + self.norm_out = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=norm_eps) + self.norm_out_modulation = nn.Sequential(nn.SiLU(), nn.Linear(hidden_size, 2 * hidden_size)) + self.proj_out = nn.Linear(hidden_size, prod(patch_size) * out_channels) + self.ulysses_group: dist.ProcessGroup | None = None + + def set_attention_config(self, attention_config: object) -> None: + """Propagate a shared attention backend configuration to every DiT block.""" + for block in self.blocks: + block.attn.set_attention_config(attention_config) + + @staticmethod + def state_dict_converter() -> "LingBotVideoStateDictConverter": + """Return the Diffusers checkpoint converter used by ModuleManager.""" + return LingBotVideoStateDictConverter() + + def promote_stability_layers_to_fp32(self) -> None: + """Keep source FP32 normalization, modulation, and routing-adjacent layers.""" + fp32_names = { + "time_embedder", + "time_modulation", + "scale_shift_table", + "norm", + "norm1", + "norm2", + "norm_q", + "norm_k", + "norm_post_attn", + "norm_post_ffn", + "norm_out", + "norm_out_modulation", + "router", + } + for name, module in self.named_modules(): + if any(part in fp32_names for part in name.split(".")): + module.float() + for name, parameter in self.named_parameters(): + if any(part in fp32_names for part in name.split(".")): + parameter.data = parameter.data.float() + + def set_ulysses_group(self, group: dist.ProcessGroup | None) -> None: + """Enable source-order Ulysses sequence parallelism for the joint token stream.""" + self.ulysses_group = group + for block in self.blocks: + block.attn.set_ulysses_group(group) + + def get_fsdp_module_names(self) -> list[str]: + """Return block containers suitable for per-block FSDP wrapping.""" + return ["blocks"] + + def _ulysses_world_size(self) -> int: + group = self.ulysses_group + if group is None or not dist.is_available() or not dist.is_initialized(): + return 1 + return dist.get_world_size(group) + + def _ulysses_shard_joint( + self, + joint: torch.Tensor, + rotary: torch.Tensor, + temb_input: torch.Tensor, + valid_mask: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, int]: + """Pad and shard the joint video/text sequence into equal Ulysses slices.""" + world_size = self._ulysses_world_size() + if world_size == 1: + return joint, rotary, temb_input, valid_mask, 0 + sequence = joint.shape[1] + padding = (-sequence) % world_size + if padding: + joint = F.pad(joint, (0, 0, 0, padding)) + rotary = F.pad(rotary, (0, 0, 0, padding)) + temb_input = F.pad(temb_input, (0, 0, 0, padding)) + valid_mask = F.pad(valid_mask, (0, padding), value=False) + local_sequence = joint.shape[1] // world_size + rank = dist.get_rank(self.ulysses_group) + start = rank * local_sequence + end = start + local_sequence + return joint[:, start:end], rotary[:, start:end], temb_input[:, start:end], valid_mask[:, start:end], padding + + def _ulysses_gather_sequence(self, local: torch.Tensor) -> torch.Tensor: + """Gather equal local token slices in rank order without changing their layout.""" + if self._ulysses_world_size() == 1: + return local + gathered = [torch.empty_like(local) for _ in range(self._ulysses_world_size())] + dist.all_gather(gathered, local.contiguous(), group=self.ulysses_group) + return torch.cat(gathered, dim=1) + + def forward( + self, + hidden_states: torch.Tensor, + timestep: torch.Tensor, + encoder_hidden_states: torch.Tensor, + encoder_attention_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + batch, channels, frames, height, width = hidden_states.shape + patch_t, patch_h, patch_w = self.patch_size + if channels * prod(self.patch_size) != self.patch_embedder.in_features: + raise ValueError("latent channels do not match checkpoint configuration") + if frames % patch_t or height % patch_h or width % patch_w: + raise ValueError("latent geometry must be divisible by the checkpoint patch size") + grid_t, grid_h, grid_w = frames // patch_t, height // patch_h, width // patch_w + video_tokens = grid_t * grid_h * grid_w + patches = hidden_states.reshape(batch, channels, grid_t, patch_t, grid_h, patch_h, grid_w, patch_w) + patches = patches.permute(0, 2, 4, 6, 3, 5, 7, 1).reshape(batch, video_tokens, -1) + video = self.patch_embedder(patches) + encoded_text_length = encoder_hidden_states.shape[1] + text_mask = ( + encoder_attention_mask.bool() + if encoder_attention_mask is not None + else torch.ones(batch, encoded_text_length, dtype=torch.bool, device=hidden_states.device) + ) + text_lengths = [int(length) for length in text_mask.sum(dim=-1).detach().cpu().tolist()] + packed_batch = batch > 1 + packed_attention = packed_batch or self._ulysses_world_size() > 1 + time_embedding = self.time_embedder(timestep) + + if packed_attention: + joint_parts: list[torch.Tensor] = [] + rotary_parts: list[torch.Tensor] = [] + temb_parts: list[torch.Tensor] = [] + packed_sequence_lengths = [video_tokens + text_length for text_length in text_lengths] + for index, text_length in enumerate(text_lengths): + text = self.text_embedder(encoder_hidden_states[index : index + 1, :text_length]) + joint_parts.append(torch.cat((video[index : index + 1], text), dim=1)) + positions = make_lingbot_video_joint_position_ids( + text_length, grid_t, grid_h, grid_w, hidden_states.device + ) + rotary_parts.append(lingbot_video_complex_frequencies(positions, self.axes_dims, self.rope_theta)) + temb_parts.append( + time_embedding[index : index + 1].unsqueeze(1).expand(1, video_tokens + text_length, -1) + ) + joint = torch.cat(joint_parts, dim=1) + rotary = torch.cat(rotary_parts).unsqueeze(0) + temb_input = torch.cat(temb_parts, dim=1) + valid_mask = torch.ones(1, joint.shape[1], dtype=torch.bool, device=hidden_states.device) + else: + text = self.text_embedder(encoder_hidden_states) + joint = torch.cat((video, text), dim=1) + positions = make_lingbot_video_joint_position_ids( + encoded_text_length, grid_t, grid_h, grid_w, hidden_states.device + ) + rotary = lingbot_video_complex_frequencies(positions, self.axes_dims, self.rope_theta).unsqueeze(0) + temb_input = time_embedding.unsqueeze(1).expand(-1, joint.shape[1], -1) + valid_mask = torch.cat( + (torch.ones(batch, video_tokens, dtype=torch.bool, device=hidden_states.device), text_mask), dim=1 + ) + packed_sequence_lengths = None + + joint, rotary, temb_input, local_valid_mask, padding = self._ulysses_shard_joint( + joint, rotary, temb_input, valid_mask + ) + if packed_sequence_lengths is not None and padding: + packed_sequence_lengths = [*packed_sequence_lengths, padding] + temb6 = self.time_modulation(temb_input) + attention_mask = None + if packed_sequence_lengths is None and not bool(local_valid_mask.all()): + attention_mask = local_valid_mask[:, None, None, :] + for block in self.blocks: + joint = block(joint, temb6, rotary, attention_mask, packed_sequence_lengths) + final_modulation = self.norm_out_modulation(temb_input) + shift, scale = final_modulation.chunk(2, dim=-1) + projected = self.proj_out((self.norm_out(joint) * (1.0 + scale) + shift).to(self.proj_out.weight.dtype)) + if self._ulysses_world_size() > 1: + projected = self._ulysses_gather_sequence(projected) + if padding: + projected = projected[:, :-padding] + if packed_batch: + projected = torch.cat( + [part[:, :video_tokens] for part in torch.split(projected, packed_sequence_lengths[:batch], dim=1)] + ) + else: + projected = projected[:, :video_tokens] + output = projected.reshape(batch, grid_t, grid_h, grid_w, patch_t, patch_h, patch_w, self.out_channels) + return output.permute(0, 7, 1, 4, 2, 5, 3, 6).reshape(batch, self.out_channels, frames, height, width) + + +class LingBotVideoStateDictConverter: + """Instantiate the registered Dense LingBot-Video checkpoint from Diffusers weights.""" + + _CONFIGS = { + "2bcf511fe5e0000519394d242b4d8abd": { + "patch_size": (1, 2, 2), + "in_channels": 16, + "out_channels": 16, + "hidden_size": 2048, + "num_attention_heads": 16, + "depth": 24, + "intermediate_size": 6144, + "text_dim": 2560, + "freq_dim": 256, + "norm_eps": 1e-6, + "rope_theta": 256.0, + "axes_dims": (32, 48, 48), + "qkv_bias": False, + "out_bias": True, + "patch_embed_bias": True, + "timestep_mlp_bias": True, + } + } + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Return unchanged Diffusers weights and their registered architecture.""" + state_dict_hash = hash_state_dict_keys(state_dict, with_shape=True) + try: + return state_dict, self._CONFIGS[state_dict_hash] + except KeyError as exc: + raise ValueError(f"Unsupported LingBot-Video Dense checkpoint hash: {state_dict_hash}") from exc + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """LingBot checkpoints use the same parameter names in both supported layouts.""" + return self.from_diffusers(state_dict) + + +register_model_config( + None, + "2bcf511fe5e0000519394d242b4d8abd", + ["lingbot_video_transformer"], + [LingBotVideoTransformer3DModel], + "diffusers", +) diff --git a/telefuser/models/lingbot_video_moe.py b/telefuser/models/lingbot_video_moe.py new file mode 100644 index 0000000..b1193dd --- /dev/null +++ b/telefuser/models/lingbot_video_moe.py @@ -0,0 +1,349 @@ +"""Checkpoint-compatible LingBot-Video MoE transformer modules. + +Numerical behavior is adapted from the Apache-2.0 licensed upstream +LingBot-Video transformer implementation. +""" + +from __future__ import annotations + +from typing import Any + +import torch +import torch.nn.functional as F +from torch import nn + +from telefuser.core.model_registry import register_model_config +from telefuser.ops.moe import fp8_expert_forward, grouped_expert_forward, quantize_expert_weight_fp8, route_topk +from telefuser.utils.model_weight import hash_state_dict_keys + +from .lingbot_video_dit import LingBotVideoBlock, LingBotVideoMLP, LingBotVideoTransformer3DModel + + +class LingBotVideoRouter(nn.Module): + """Checkpoint-compatible group-limited sigmoid MoE router.""" + + def __init__( + self, hidden_size: int, num_experts: int, top_k: int, n_group: int, topk_group: int, route_scale: float + ) -> None: + super().__init__() + self.weight = nn.Parameter(torch.empty(num_experts, hidden_size)) + self.register_buffer("e_score_correction_bias", torch.zeros(num_experts), persistent=True) + self.num_experts = num_experts + self.top_k = top_k + self.n_group = n_group + self.topk_group = topk_group + self.route_scale = route_scale + + def forward(self, tokens: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + logits = F.linear(tokens.float(), self.weight.float()) + indices, weights = route_topk( + logits, + num_expert_groups=self.n_group, + top_k_groups=self.topk_group, + top_k=self.top_k, + correction_bias=self.e_score_correction_bias, + routing_scale=self.route_scale, + ) + return indices, weights.to(tokens.dtype) + + +class LingBotVideoGroupedExperts(nn.Module): + """Official grouped expert layout with sorted eager execution.""" + + def __init__( + self, + num_experts: int, + hidden_size: int, + intermediate_size: int, + *, + execution_backend: str = "sorted", + ) -> None: + super().__init__() + self.w1 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.w2 = nn.Parameter(torch.empty(num_experts, hidden_size, intermediate_size)) + self.w3 = nn.Parameter(torch.empty(num_experts, intermediate_size, hidden_size)) + self.register_buffer("w1_scale", None, persistent=False) + self.register_buffer("w2_scale", None, persistent=False) + self.register_buffer("w3_scale", None, persistent=False) + self.set_execution_backend(execution_backend) + + def set_execution_backend(self, backend: str) -> None: + """Select the source-style sorted path or the diagnostic where fallback.""" + if backend not in {"fp8", "grouped_mm", "sorted", "where"}: + raise ValueError("execution backend must be 'fp8', 'grouped_mm', 'sorted', or 'where'") + self.execution_backend = backend + + def quantize_fp8_(self) -> None: + """Replace BF16 expert weights with per-channel FP8 parameters in place.""" + if self.w1.dtype == torch.float8_e4m3fn: + return + self.w1, self.w1_scale = self._quantized_parameter(self.w1) + self.w2, self.w2_scale = self._quantized_parameter(self.w2) + self.w3, self.w3_scale = self._quantized_parameter(self.w3) + + @staticmethod + def _quantized_parameter(weight: nn.Parameter) -> tuple[nn.Parameter, torch.Tensor]: + quantized, scale = quantize_expert_weight_fp8(weight.detach()) + return nn.Parameter(quantized, requires_grad=False), scale + + def forward(self, tokens: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + if self.execution_backend == "fp8": + if self.w1_scale is None or self.w2_scale is None or self.w3_scale is None: + raise RuntimeError("FP8 expert backend requires quantize_fp8_() before execution") + return fp8_expert_forward( + tokens, + indices, + weights, + self.w1, + self.w2, + self.w3, + self.w1_scale, + self.w2_scale, + self.w3_scale, + ) + if self.execution_backend == "grouped_mm": + return grouped_expert_forward(tokens, indices, weights, self.w1, self.w2, self.w3) + if self.execution_backend == "where": + return self._forward_where(tokens, indices, weights) + return self._forward_sorted(tokens, indices, weights) + + def _forward_sorted(self, tokens: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + top_k = indices.shape[1] + flat_weights = weights.reshape(-1) + flat_indices = indices.reshape(-1) + active_positions = torch.where(flat_weights != 0)[0] + if active_positions.numel() == 0: + return tokens.new_zeros(tokens.shape) + active_experts = flat_indices[active_positions] + counts = torch.zeros(self.w1.shape[0], device=tokens.device, dtype=torch.int64) + counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64)) + sort_order = torch.argsort(active_experts, stable=True) + sorted_positions = active_positions[sort_order] + sorted_tokens = tokens[sorted_positions // top_k] + + outputs: list[torch.Tensor] = [] + for expert_index, selected in enumerate(torch.split(sorted_tokens, counts.tolist(), dim=0)): + if selected.numel() == 0: + continue + activation = F.silu(F.linear(selected, self.w1[expert_index])) * F.linear(selected, self.w3[expert_index]) + outputs.append(F.linear(activation, self.w2[expert_index])) + expert_output = torch.cat(outputs, dim=0) + routed = torch.zeros(tokens.shape[0] * top_k, tokens.shape[-1], dtype=tokens.dtype, device=tokens.device) + routed[sorted_positions] = expert_output + return ( + (routed.reshape(tokens.shape[0], top_k, -1).float() * weights.float().unsqueeze(-1)) + .sum(dim=1) + .to(tokens.dtype) + ) + + def _forward_where(self, tokens: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + top_k = indices.shape[1] + output = torch.zeros( + tokens.shape[0] * top_k, + tokens.shape[-1], + dtype=tokens.dtype, + device=tokens.device, + ) + for expert_index in range(self.w1.shape[0]): + token_indices, topk_indices = torch.where(indices == expert_index) + if token_indices.numel() == 0: + continue + selected = tokens[token_indices] + activation = F.silu(F.linear(selected, self.w1[expert_index])) * F.linear(selected, self.w3[expert_index]) + output[token_indices * top_k + topk_indices] = F.linear(activation, self.w2[expert_index]) + return ( + (output.reshape(tokens.shape[0], top_k, -1).float() * weights.float().unsqueeze(-1)) + .sum(dim=1) + .to(tokens.dtype) + ) + + +class LingBotVideoSparseMoeBlock(nn.Module): + """Correctness-first MoE block with routed and optional shared experts.""" + + def __init__( + self, + hidden_size: int, + num_experts: int, + top_k: int, + intermediate_size: int, + n_group: int, + topk_group: int, + route_scale: float, + n_shared_experts: int = 1, + ) -> None: + super().__init__() + self.router = LingBotVideoRouter(hidden_size, num_experts, top_k, n_group, topk_group, route_scale) + self.experts = LingBotVideoGroupedExperts(num_experts, hidden_size, intermediate_size) + self.shared_experts = ( + LingBotVideoMLP(hidden_size, intermediate_size * n_shared_experts) if n_shared_experts > 0 else None + ) + + def forward(self, hidden_states: torch.Tensor, padding_mask: torch.Tensor | None = None) -> torch.Tensor: + batch, sequence, hidden_size = hidden_states.shape + tokens = hidden_states.reshape(-1, hidden_size) + indices, weights = self.router(tokens) + if padding_mask is not None: + weights = weights * padding_mask.reshape(-1, 1).to(weights.dtype) + output = self.experts(tokens, indices, weights).reshape(batch, sequence, hidden_size) + if self.shared_experts is not None: + output = output + self.shared_experts(hidden_states) + return output + + +class LingBotVideoMoeBlock(LingBotVideoBlock): + """LingBot block selecting official sparse or dense FFN by layer index.""" + + def __init__( + self, + hidden_size: int, + num_attention_heads: int, + intermediate_size: int, + *, + norm_eps: float, + qkv_bias: bool, + out_bias: bool, + layer_index: int, + num_experts: int, + top_k: int, + moe_intermediate_size: int, + decoder_sparse_step: int, + mlp_only_layers: tuple[int, ...], + n_group: int, + topk_group: int, + route_scale: float, + n_shared_experts: int, + ) -> None: + super().__init__(hidden_size, num_attention_heads, intermediate_size, norm_eps, qkv_bias, out_bias) + if layer_index not in mlp_only_layers and num_experts > 0 and (layer_index + 1) % decoder_sparse_step == 0: + self.ffn = LingBotVideoSparseMoeBlock( + hidden_size, + num_experts, + top_k, + moe_intermediate_size, + n_group, + topk_group, + route_scale, + n_shared_experts, + ) + + +class LingBotVideoMoeTransformer3DModel(LingBotVideoTransformer3DModel): + """Official MoE/refiner transformer using eager expert execution for correctness.""" + + def __init__( + self, + *, + num_experts: int, + num_experts_per_tok: int, + moe_intermediate_size: int, + decoder_sparse_step: int = 1, + mlp_only_layers: tuple[int, ...] = (), + n_group: int = 1, + topk_group: int = 1, + routed_scaling_factor: float = 1.0, + n_shared_experts: int = 1, + **kwargs: Any, + ) -> None: + super().__init__(**kwargs) + hidden_size = self.patch_embedder.out_features + num_attention_heads = self.blocks[0].attn.num_heads + intermediate_size = self.blocks[0].ffn.gate_proj.out_features + norm_eps = self.blocks[0].norm1.variance_epsilon + qkv_bias = self.blocks[0].attn.to_q.bias is not None + out_bias = self.blocks[0].attn.to_out.bias is not None + self.blocks = nn.ModuleList( + [ + LingBotVideoMoeBlock( + hidden_size, + num_attention_heads, + intermediate_size, + norm_eps=norm_eps, + qkv_bias=qkv_bias, + out_bias=out_bias, + layer_index=index, + num_experts=num_experts, + top_k=num_experts_per_tok, + moe_intermediate_size=moe_intermediate_size, + decoder_sparse_step=decoder_sparse_step, + mlp_only_layers=tuple(mlp_only_layers), + n_group=n_group, + topk_group=topk_group, + route_scale=routed_scaling_factor, + n_shared_experts=n_shared_experts, + ) + for index in range(len(self.blocks)) + ] + ) + + def set_expert_execution_backend(self, backend: str) -> None: + """Set one eager expert backend consistently across all sparse blocks.""" + for module in self.modules(): + if isinstance(module, LingBotVideoGroupedExperts): + module.set_execution_backend(backend) + + @staticmethod + def state_dict_converter() -> "LingBotVideoMoeStateDictConverter": + """Return the Diffusers checkpoint converter used by ModuleManager.""" + return LingBotVideoMoeStateDictConverter() + + def quantize_experts_fp8_(self) -> None: + """Quantize every routed expert while keeping dense/shared layers in BF16.""" + for module in self.modules(): + if isinstance(module, LingBotVideoGroupedExperts): + module.quantize_fp8_() + + +class LingBotVideoMoeStateDictConverter: + """Instantiate the registered MoE LingBot-Video checkpoint from Diffusers weights.""" + + _CONFIGS = { + "65b83aa625cd362ff5ff3409fb367a6f": { + "patch_size": (1, 2, 2), + "in_channels": 16, + "out_channels": 16, + "hidden_size": 2048, + "num_attention_heads": 16, + "depth": 48, + "intermediate_size": 6144, + "text_dim": 2560, + "freq_dim": 256, + "norm_eps": 1e-6, + "rope_theta": 256.0, + "axes_dims": (32, 48, 48), + "qkv_bias": False, + "out_bias": True, + "patch_embed_bias": True, + "timestep_mlp_bias": True, + "num_experts": 128, + "num_experts_per_tok": 8, + "moe_intermediate_size": 768, + "decoder_sparse_step": 1, + "mlp_only_layers": (), + "n_group": 4, + "topk_group": 2, + "routed_scaling_factor": 2.5, + "n_shared_experts": 1, + } + } + + def from_diffusers(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """Return unchanged Diffusers weights and their registered architecture.""" + state_dict_hash = hash_state_dict_keys(state_dict, with_shape=True) + try: + return state_dict, self._CONFIGS[state_dict_hash] + except KeyError as exc: + raise ValueError(f"Unsupported LingBot-Video MoE checkpoint hash: {state_dict_hash}") from exc + + def from_official(self, state_dict: dict[str, torch.Tensor]) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + """LingBot checkpoints use the same parameter names in both supported layouts.""" + return self.from_diffusers(state_dict) + + +register_model_config( + None, + "65b83aa625cd362ff5ff3409fb367a6f", + ["lingbot_video_moe_transformer"], + [LingBotVideoMoeTransformer3DModel], + "diffusers", +) diff --git a/telefuser/ops/__init__.py b/telefuser/ops/__init__.py index 10d060f..2b76b31 100644 --- a/telefuser/ops/__init__.py +++ b/telefuser/ops/__init__.py @@ -12,6 +12,7 @@ from .base import CustomOp, CustomOpFunction from .custom_op import TritonKernelWrapper, register_custom_op +from .moe import grouped_expert_forward, route_topk from .normalization import AdaLayerNormContinuous, LayerNorm, RMSNorm, fused_scale_shift, modulate from .rotary import apply_rotary_emb @@ -28,6 +29,8 @@ "AdaLayerNormContinuous", "fused_scale_shift", "modulate", + "route_topk", + "grouped_expert_forward", # Rotary "apply_rotary_emb", ] diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index 0612e08..d2bc4d2 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -48,6 +48,31 @@ _warned_attn_fallback: set[str] = set() +def _packed_sdpa( + q: Tensor, + k: Tensor, + v: Tensor, + sequence_lengths: list[int], + *, + scale: float | None, + is_causal: bool, +) -> Tensor: + """Run SDPA independently for sequences packed along the BNSD sequence axis.""" + if q.shape != k.shape or q.shape != v.shape or q.ndim != 4 or q.shape[0] != 1: + raise ValueError("packed SDPA requires matching Q/K/V tensors with shape [1, heads, sequence, dim]") + if any(length <= 0 for length in sequence_lengths) or sum(sequence_lengths) != q.shape[2]: + raise ValueError("packed SDPA sequence lengths must be positive and sum to the packed sequence dimension") + outputs = [ + F.scaled_dot_product_attention(q_part, k_part, v_part, scale=scale, is_causal=is_causal) + for q_part, k_part, v_part in zip( + torch.split(q, sequence_lengths, dim=2), + torch.split(k, sequence_lengths, dim=2), + torch.split(v, sequence_lengths, dim=2), + ) + ] + return torch.cat(outputs, dim=2) + + class SparseAttentionState: """Runtime state for sparse attention computation. @@ -97,6 +122,7 @@ def attention( output_layout: Literal["BSND", "BNSD"] = "BSND", is_causal: bool = False, return_lse: bool = False, + sequence_lengths: list[int] | None = None, **kwargs: Any, ) -> Tensor | tuple[Tensor, Tensor]: """Unified attention function. @@ -131,6 +157,8 @@ def attention( attn_impl = attention_config.attn_impl scale = scale if scale is not None else attention_config.scale is_causal = is_causal or attention_config.is_causal + if sequence_lengths is not None and attn_impl != AttnImplType.TORCH_SDPA: + raise ValueError("packed sequence attention currently requires TORCH_SDPA") # Handle sparse attention if attn_impl in (AttnImplType.RADIAL_ATTN, AttnImplType.LOCAL_SPARSE_ATTN): @@ -206,7 +234,12 @@ def attention( elif attn_impl == AttnImplType.TORCH_CUDNN and SDPA_AVAILABLE: output = sdpa_attn_cudnn(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) elif attn_impl == AttnImplType.TORCH_SDPA and SDPA_AVAILABLE: - output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) + if sequence_lengths is None: + output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) + else: + if attn_mask is not None or current_layout != "BNSD": + raise ValueError("packed SDPA requires BNSD layout without an attention mask") + output = _packed_sdpa(q, k, v, sequence_lengths, scale=scale, is_causal=is_causal) # Sage Attention variants elif SAGE_ATTN_AVAILABLE and sageattention is not None: @@ -265,7 +298,12 @@ def attention( v = v.transpose(1, 2).contiguous() current_layout = "BNSD" - output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) + if sequence_lengths is None: + output = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, scale=scale, is_causal=is_causal) + else: + if attn_mask is not None or current_layout != "BNSD": + raise ValueError("packed SDPA requires BNSD layout without an attention mask") + output = _packed_sdpa(q, k, v, sequence_lengths, scale=scale, is_causal=is_causal) attn_impl = AttnImplType.TORCH_SDPA # Handle output layout conversion - output matches current_layout, may need to convert to output_layout diff --git a/telefuser/ops/moe.py b/telefuser/ops/moe.py new file mode 100644 index 0000000..ce69555 --- /dev/null +++ b/telefuser/ops/moe.py @@ -0,0 +1,272 @@ +"""Mixture-of-Experts routing and grouped expert primitives.""" + +from __future__ import annotations + +import torch +import torch.nn.functional as F + + +def route_topk( + logits: torch.Tensor, + *, + num_expert_groups: int = 1, + experts_per_group: int | None = None, + top_k_groups: int | None = None, + top_k: int = 1, + correction_bias: torch.Tensor | None = None, + routing_scale: float = 1.0, +) -> tuple[torch.Tensor, torch.Tensor]: + """Select experts using LingBot's group-limited sigmoid router semantics. + + Correction bias participates only in discrete selection. Returned weights + are gathered from the unbiased sigmoid scores and normalized over top-k. + """ + if logits.ndim != 2: + raise ValueError("router logits must have shape [tokens, experts]") + tokens, experts = logits.shape + if top_k < 1 or top_k > experts: + raise ValueError("top_k must be within the number of experts") + if experts_per_group is None: + if experts % num_expert_groups: + raise ValueError("experts must divide evenly across groups") + experts_per_group = experts // num_expert_groups + if num_expert_groups * experts_per_group != experts: + raise ValueError("group dimensions do not match expert count") + if top_k_groups is None: + top_k_groups = num_expert_groups + if not 1 <= top_k_groups <= num_expert_groups: + raise ValueError("top_k_groups must be within the number of groups") + scores = logits.sigmoid() + selection_scores = scores if correction_bias is None else scores + correction_bias.reshape(1, -1) + grouped = selection_scores.reshape(tokens, num_expert_groups, experts_per_group) + group_values = grouped.topk(min(2, experts_per_group), dim=-1, sorted=False).values.sum(dim=-1) + selected_groups = group_values.topk(top_k_groups, dim=-1, sorted=False).indices + group_mask = torch.zeros_like(group_values, dtype=torch.bool) + group_mask.scatter_(1, selected_groups, True) + score_mask = group_mask.unsqueeze(-1).expand(tokens, num_expert_groups, experts_per_group).reshape(tokens, experts) + masked = selection_scores.masked_fill(~score_mask, float("-inf")) + indices = masked.topk(top_k, dim=-1, sorted=False).indices + weights = scores.gather(-1, indices) + weights = weights / weights.sum(dim=-1, keepdim=True).clamp_min(torch.finfo(weights.dtype).eps) + return indices, weights * routing_scale + + +def grouped_expert_forward( + tokens: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + *, + alignment: int = 8, +) -> torch.Tensor: + """Run LingBot routed experts with PyTorch's native grouped GEMM. + + The token packing, alignment, BF16 GEMMs, and FP32 route reduction mirror + the optimized upstream LingBot-Video implementation. Callers should retain + an eager fallback because torch._grouped_mm is only available in recent + CUDA-enabled PyTorch builds. + """ + if not hasattr(torch, "_grouped_mm"): + raise RuntimeError("grouped expert execution requires torch._grouped_mm") + if tokens.device.type != "cuda": + raise RuntimeError("grouped expert execution requires CUDA tensors") + if indices.shape != weights.shape or indices.ndim != 2: + raise ValueError("expert indices and weights must have matching [tokens, top_k] shapes") + if tokens.ndim != 2: + raise ValueError("expert tokens must have shape [tokens, hidden_size]") + if alignment <= 0: + raise ValueError("grouped expert alignment must be positive") + + num_tokens, hidden_size = tokens.shape + top_k = indices.shape[1] + num_experts = w1.shape[0] + flat_weights = weights.reshape(-1) + flat_indices = indices.reshape(-1) + active_positions = torch.where(flat_weights != 0)[0] + if active_positions.numel() == 0: + return tokens.new_zeros(tokens.shape) + + active_experts = flat_indices[active_positions] + counts = torch.zeros(num_experts, device=tokens.device, dtype=torch.int64) + counts.scatter_add_(0, active_experts, torch.ones_like(active_experts, dtype=torch.int64)) + sort_order = torch.argsort(active_experts, stable=True) + sorted_positions = active_positions[sort_order] + sorted_tokens = tokens[sorted_positions // top_k] + + padded_tokens, padded_indices, aligned_counts = _pad_grouped_tokens( + sorted_tokens, + counts, + alignment=alignment, + ) + offsets = torch.cumsum(aligned_counts, dim=0, dtype=torch.int32) + grouped_mm = torch._grouped_mm + gate = grouped_mm( + padded_tokens.bfloat16(), + w1.bfloat16().transpose(-2, -1), + offs=offsets, + ) + up = grouped_mm( + padded_tokens.bfloat16(), + w3.bfloat16().transpose(-2, -1), + offs=offsets, + ) + expert_output = grouped_mm( + F.silu(gate) * up, + w2.bfloat16().transpose(-2, -1), + offs=offsets, + ).to(tokens.dtype) + expert_output = _unpad_grouped_tokens(expert_output, padded_indices, sorted_tokens.shape[0]) + + routed = torch.zeros( + num_tokens * top_k, + hidden_size, + dtype=expert_output.dtype, + device=expert_output.device, + ) + routed[sorted_positions] = expert_output + return ( + (routed.reshape(num_tokens, top_k, hidden_size).float() * weights.float().unsqueeze(-1)) + .sum(dim=1) + .to(tokens.dtype) + ) + + +def quantize_expert_weight_fp8(weight: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Quantize [experts, out, in] weights with per-output-channel FP8 scales.""" + if weight.ndim != 3: + raise ValueError("expert weight must have shape [experts, out_features, in_features]") + fp8_dtype = torch.float8_e4m3fn + fp8_info = torch.finfo(fp8_dtype) + source = weight.float() + scales = source.abs().amax(dim=2).clamp_min(torch.finfo(torch.float32).tiny) / fp8_info.max + quantized = (source / scales.unsqueeze(2)).clamp(min=fp8_info.min, max=fp8_info.max).to(fp8_dtype) + return quantized, scales + + +def fp8_expert_forward( + tokens: torch.Tensor, + indices: torch.Tensor, + weights: torch.Tensor, + w1: torch.Tensor, + w2: torch.Tensor, + w3: torch.Tensor, + w1_scale: torch.Tensor, + w2_scale: torch.Tensor, + w3_scale: torch.Tensor, +) -> torch.Tensor: + """Run sorted LingBot experts with native dynamic W8A8 scaled GEMMs.""" + if not hasattr(torch, "_scaled_mm"): + raise RuntimeError("FP8 expert execution requires torch._scaled_mm") + if tokens.device.type != "cuda": + raise RuntimeError("FP8 expert execution requires CUDA tensors") + if w1.dtype != torch.float8_e4m3fn or w2.dtype != torch.float8_e4m3fn or w3.dtype != torch.float8_e4m3fn: + raise RuntimeError("FP8 expert weights must be quantized before execution") + if indices.shape != weights.shape or indices.ndim != 2: + raise ValueError("expert indices and weights must have matching [tokens, top_k] shapes") + + top_k = indices.shape[1] + flat_weights = weights.reshape(-1) + flat_indices = indices.reshape(-1) + active_positions = torch.where(flat_weights != 0)[0] + if active_positions.numel() == 0: + return tokens.new_zeros(tokens.shape) + + active_experts = flat_indices[active_positions] + sort_order = torch.argsort(active_experts, stable=True) + sorted_positions = active_positions[sort_order] + sorted_experts = active_experts[sort_order] + sorted_tokens = tokens[sorted_positions // top_k] + counts = torch.bincount(sorted_experts, minlength=w1.shape[0]) + + outputs: list[torch.Tensor] = [] + offset = 0 + for expert_index, count in enumerate(counts.tolist()): + if count == 0: + continue + selected = sorted_tokens[offset : offset + count] + offset += count + selected_fp8, selected_scale = _quantize_rows_fp8(selected) + gate = _scaled_mm_fp8(selected_fp8, w1[expert_index], selected_scale, w1_scale[expert_index]) + up = _scaled_mm_fp8(selected_fp8, w3[expert_index], selected_scale, w3_scale[expert_index]) + activation_fp8, activation_scale = _quantize_rows_fp8(F.silu(gate) * up) + outputs.append(_scaled_mm_fp8(activation_fp8, w2[expert_index], activation_scale, w2_scale[expert_index])) + + expert_output = torch.cat(outputs, dim=0).to(tokens.dtype) + routed = torch.zeros( + tokens.shape[0] * top_k, + tokens.shape[-1], + dtype=tokens.dtype, + device=tokens.device, + ) + routed[sorted_positions] = expert_output + return ( + (routed.reshape(tokens.shape[0], top_k, -1).float() * weights.float().unsqueeze(-1)).sum(dim=1).to(tokens.dtype) + ) + + +def _quantize_rows_fp8(tensor: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + """Dynamically quantize matrix rows for torch._scaled_mm row-wise scaling.""" + fp8_dtype = torch.float8_e4m3fn + fp8_info = torch.finfo(fp8_dtype) + scale = tensor.float().abs().amax(dim=1, keepdim=True).clamp_min(torch.finfo(torch.float32).tiny) + scale = (scale / fp8_info.max).contiguous() + quantized = (tensor.float() / scale).clamp(min=fp8_info.min, max=fp8_info.max).to(fp8_dtype) + return quantized, scale + + +def _scaled_mm_fp8( + activation: torch.Tensor, + weight: torch.Tensor, + activation_scale: torch.Tensor, + weight_scale: torch.Tensor, +) -> torch.Tensor: + """Apply one row-wise scaled FP8 matrix multiplication.""" + return torch._scaled_mm( + activation, + weight.transpose(0, 1), + scale_a=activation_scale, + scale_b=weight_scale.unsqueeze(0).contiguous(), + out_dtype=torch.bfloat16, + ) + + +def _pad_grouped_tokens( + tokens: torch.Tensor, + counts: torch.Tensor, + *, + alignment: int, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Pad each expert group without synchronizing counts back to the CPU.""" + num_tokens = tokens.shape[0] + num_experts = counts.shape[0] + max_length = ((num_tokens + num_experts * alignment + alignment - 1) // alignment) * alignment + counts_i64 = counts.to(torch.int64) + aligned_counts_i64 = (torch.clamp_min(counts_i64, alignment) + alignment - 1) // alignment * alignment + write_offsets = torch.cumsum(aligned_counts_i64, dim=0) - aligned_counts_i64 + end_offsets = torch.cumsum(aligned_counts_i64, dim=0) + start_indices = torch.cumsum(counts_i64, dim=0) - counts_i64 + + slots = torch.arange(max_length, dtype=torch.int64, device=tokens.device) + expert_indices = torch.bucketize(slots, end_offsets, right=True) + valid_expert = expert_indices < num_experts + safe_expert_indices = expert_indices.clamp(max=num_experts - 1) + local_indices = slots - write_offsets[safe_expert_indices] + source_indices = start_indices[safe_expert_indices] + local_indices + valid = valid_expert & (local_indices < counts_i64[safe_expert_indices]) + padded_indices = torch.where(valid, source_indices, torch.full_like(source_indices, num_tokens)) + + tokens_with_sentinel = torch.vstack((tokens, tokens.new_zeros((1, tokens.shape[-1])))) + return tokens_with_sentinel[padded_indices], padded_indices, aligned_counts_i64.to(torch.int32) + + +def _unpad_grouped_tokens( + output: torch.Tensor, + padded_indices: torch.Tensor, + num_tokens: int, +) -> torch.Tensor: + """Restore expert-sorted rows from the aligned grouped-MM output.""" + unpadded = output.new_empty((num_tokens + 1, output.shape[-1])) + unpadded[padded_indices] = output + return unpadded[:num_tokens] diff --git a/telefuser/pipelines/lingbot_video/__init__.py b/telefuser/pipelines/lingbot_video/__init__.py new file mode 100644 index 0000000..aa6b48b --- /dev/null +++ b/telefuser/pipelines/lingbot_video/__init__.py @@ -0,0 +1,79 @@ +"""LingBot-Video reusable pipeline contracts and components.""" + +from .data import ( + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_NEGATIVE_PROMPT_IMAGE, + LingBotVideoModelConfig, + LingBotVideoRequest, + default_negative_caption, + load_lingbot_video_model_config, + load_lingbot_video_prompt, + num_frames_from_duration, + parse_lingbot_video_prompt, + preprocess_ti2v_image, + smart_resize, + validate_frame_count, +) +from .denoising import LingBotVideoDenoisingStage, denoise_lingbot_video, reinject_ti2v_condition +from .pipeline import ( + LingBotVideoGeneration, + LingBotVideoPipeline, + LingBotVideoPipelineConfig, + LingBotVideoPromptConditions, +) +from .refiner import ( + LingBotVideoRefinerStage, + compute_refiner_sigmas, + compute_training_aligned_indices, + compute_training_frame_budget, + load_refiner_first_frame, + load_refiner_video_file, + prepare_refiner_latent, + prepare_refiner_video, +) +from .text_encoding import LingBotVideoTextEncodingStage +from .vae import ( + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + denormalize_latent, + first_frame_condition_mask, + latent_shape, + normalize_latent, +) + +__all__ = [ + "LingBotVideoPipeline", + "LingBotVideoGeneration", + "DEFAULT_NEGATIVE_PROMPT", + "DEFAULT_NEGATIVE_PROMPT_IMAGE", + "LingBotVideoPromptConditions", + "LingBotVideoRefinerStage", + "compute_refiner_sigmas", + "load_refiner_first_frame", + "load_refiner_video_file", + "prepare_refiner_latent", + "compute_training_frame_budget", + "compute_training_aligned_indices", + "prepare_refiner_video", + "LingBotVideoVAEEncodeStage", + "LingBotVideoVAEDecodeStage", + "LingBotVideoTextEncodingStage", + "LingBotVideoDenoisingStage", + "latent_shape", + "normalize_latent", + "denormalize_latent", + "first_frame_condition_mask", + "denoise_lingbot_video", + "default_negative_caption", + "reinject_ti2v_condition", + "LingBotVideoPipelineConfig", + "LingBotVideoModelConfig", + "preprocess_ti2v_image", + "smart_resize", + "load_lingbot_video_model_config", + "LingBotVideoRequest", + "load_lingbot_video_prompt", + "num_frames_from_duration", + "parse_lingbot_video_prompt", + "validate_frame_count", +] diff --git a/telefuser/pipelines/lingbot_video/data.py b/telefuser/pipelines/lingbot_video/data.py new file mode 100644 index 0000000..df8225d --- /dev/null +++ b/telefuser/pipelines/lingbot_video/data.py @@ -0,0 +1,206 @@ +"""Data contracts shared by LingBot-Video pipeline stages. + +TI2V geometry and structured-caption behavior are adapted from the +Apache-2.0 licensed upstream LingBot-Video implementation. +""" + +from __future__ import annotations + +import json +from dataclasses import dataclass +from pathlib import Path + +import torch + + +def validate_frame_count(num_frames: int) -> int: + """Validate LingBot's temporal frame contract (1 or ``4n + 1``).""" + if num_frames < 1 or (num_frames - 1) % 4: + raise ValueError("LingBot-Video requires num_frames to be 1 or 4n+1") + return num_frames + + +_RUNTIME_PROMPT_FIELDS = frozenset({"duration", "fps", "height", "width", "num_frames", "resolution", "ratio"}) + +# Source-compatible default CFG conditions, adapted from the Apache-2.0 +# licensed upstream LingBot-Video pipeline. Keep the strings serialized exactly +# so Qwen3-VL receives the same structured negative caption by default. +DEFAULT_NEGATIVE_PROMPT = ( + '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", "pixelated", ' + '"jpeg artifacts", "low resolution", "unstable color", "color flicker", "underexposed", "overexposed", ' + '"invisible subject", "subject hidden in darkness"], "artistic_style": ["painting", "illustration", ' + '"drawing", "cartoon", "3d render", "cgi", "sketch", "digital art"], "composition_and_content": ' + '["text", "watermark", "signature", "logo", "subtitles", "pillarboxed", "side bars", "portrait image ' + 'in landscape frame"], "temporal_and_motion_stability": ["flickering", "jittery", "motion blur", ' + '"temporal inconsistency", "warping", "morphing", "incoherent motion", "unnatural movement", "static ' + 'object with sudden jump", "frame-to-frame inconsistency"], "material_and_structure": ["plastic-like glass", ' + '"unrealistic texture", "deformed bottle", "liquid freezing improperly", "distorted reflections"]}}' +) +DEFAULT_NEGATIVE_PROMPT_IMAGE = ( + '{"universal_negative": {"visual_quality": ["low quality", "worst quality", "blurry", "pixelated", ' + '"jpeg artifacts", "low resolution", "underexposed", "overexposed", "invisible subject", "subject hidden ' + 'in darkness"], "artistic_style": ["painting", "illustration", "drawing", "cartoon", "3d render", ' + '"cgi", "sketch", "digital art"], "composition_and_content": ["text", "watermark", "signature", ' + '"logo", "pillarboxed", "side bars", "portrait image in landscape frame"], "material_and_structure": ' + '["plastic-like glass", "unrealistic texture", "deformed bottle", "distorted reflections"]}}' +) + + +def default_negative_caption(num_frames: int) -> str: + """Return the source-compatible T2I or video default CFG caption.""" + return DEFAULT_NEGATIVE_PROMPT_IMAGE if num_frames == 1 else DEFAULT_NEGATIVE_PROMPT + + +def parse_lingbot_video_prompt(payload: object) -> tuple[str, float | None]: + """Extract the source-compatible caption and optional duration from a prompt sample.""" + if isinstance(payload, list): + if not payload: + raise ValueError("LingBot-Video prompt sample list must not be empty") + payload = payload[0] + if not isinstance(payload, dict): + raise ValueError("LingBot-Video prompt sample must be a dictionary or a non-empty list of dictionaries") + caption = ( + payload["caption"] + if "caption" in payload + else {key: value for key, value in payload.items() if key not in _RUNTIME_PROMPT_FIELDS} + ) + if isinstance(caption, (dict, list)): + caption_text = json.dumps(caption, ensure_ascii=False, separators=(",", ":")) + else: + caption_text = str(caption) + if not caption_text: + raise ValueError("LingBot-Video prompt caption must not be empty") + duration = payload.get("duration") + return caption_text, None if duration is None else float(duration) + + +def load_lingbot_video_prompt(path: str | Path) -> tuple[str, float | None]: + """Load an upstream rewriter prompt file and return its caption and duration.""" + return parse_lingbot_video_prompt(json.loads(Path(path).read_text(encoding="utf-8"))) + + +def num_frames_from_duration(duration: float, fps: int = 24) -> int: + """Match upstream duration conversion for a ``4n+1`` LingBot video length.""" + if duration < 0 or fps <= 0: + raise ValueError("duration must be non-negative and fps must be positive") + frame_count = int(float(duration) * fps) + return ((frame_count - 1) // 4 + 1) * 4 + 1 + + +@dataclass(frozen=True) +class LingBotVideoModelConfig: + """Architecture values encoded by the official Dense/MoE checkpoints.""" + + variant: str = "dense" + num_layers: int = 24 + hidden_size: int = 2048 + num_heads: int = 16 + in_channels: int = 16 + patch_size: tuple[int, int, int] = (1, 2, 2) + text_dim: int = 2560 + intermediate_size: int = 6144 + num_experts: int = 0 + top_k: int = 0 + + def __post_init__(self) -> None: + if self.variant not in {"dense", "moe", "refiner"}: + raise ValueError(f"Unsupported LingBot-Video variant: {self.variant}") + if self.num_layers <= 0 or self.hidden_size <= 0 or self.num_heads <= 0: + raise ValueError("LingBot-Video architecture dimensions must be positive") + if self.variant in {"moe", "refiner"} and (self.num_experts <= 0 or self.top_k <= 0): + raise ValueError("MoE and refiner variants require num_experts and top_k") + + +@dataclass(frozen=True) +class LingBotVideoRequest: + """Structured request consumed by the numerical pipeline.""" + + caption: str + height: int + width: int + num_frames: int + image: object | None = None + + def __post_init__(self) -> None: + validate_frame_count(self.num_frames) + if self.height <= 0 or self.width <= 0: + raise ValueError("height and width must be positive") + if self.height % 16 or self.width % 16: + raise ValueError("LingBot-Video height and width must be divisible by 16") + if self.image is not None and self.num_frames == 1: + raise ValueError("TI2V image conditioning requires a video frame count") + + +def load_lingbot_video_model_config(path: str | Path, *, variant: str = "dense") -> LingBotVideoModelConfig: + """Load the architecture subset from a Diffusers transformer config.json.""" + config_path = Path(path) + if config_path.is_dir(): + config_path = config_path / "config.json" + payload = json.loads(config_path.read_text(encoding="utf-8")) + required = ( + "depth", + "hidden_size", + "num_attention_heads", + "in_channels", + "patch_size", + "text_dim", + "intermediate_size", + ) + missing = [key for key in required if key not in payload] + if missing: + raise ValueError(f"LingBot transformer config is missing: {missing}") + return LingBotVideoModelConfig( + variant=variant, + num_layers=int(payload["depth"]), + hidden_size=int(payload["hidden_size"]), + num_heads=int(payload["num_attention_heads"]), + in_channels=int(payload["in_channels"]), + patch_size=tuple(int(v) for v in payload["patch_size"]), + text_dim=int(payload["text_dim"]), + intermediate_size=int(payload["intermediate_size"]), + num_experts=int(payload.get("num_experts", 0)), + top_k=int(payload.get("num_experts_per_tok", 0)) if int(payload.get("num_experts", 0)) else 0, + ) + + +def smart_resize( + height: int, width: int, *, factor: int, min_pixels: int = 4 * 28**2, max_pixels: int = 16384 * 28**2 +) -> tuple[int, int]: + """Match upstream Qwen3-VL image smart-resize geometry.""" + import math + + if min_pixels > max_pixels: + raise ValueError("max_pixels must be greater than or equal to min_pixels") + if min(height, width) <= 0 or max(height, width) / min(height, width) > 200: + raise ValueError("image dimensions have an unsupported aspect ratio") + resize_height = max(factor, round(height / factor) * factor) + resize_width = max(factor, round(width / factor) * factor) + if resize_height * resize_width > max_pixels: + scale = math.sqrt(height * width / max_pixels) + resize_height = math.floor(height / scale / factor) * factor + resize_width = math.floor(width / scale / factor) * factor + elif resize_height * resize_width < min_pixels: + scale = math.sqrt(min_pixels / (height * width)) + resize_height = math.ceil(height * scale / factor) * factor + resize_width = math.ceil(width * scale / factor) * factor + return resize_height, resize_width + + +def preprocess_ti2v_image(image: "torch.Tensor", *, height: int, width: int) -> "torch.Tensor": + """Resize-short-side then center-crop an RGB image to the condition frame.""" + import math + + import torch.nn.functional as F + + if image.ndim != 4 or image.shape[1] != 3: + raise ValueError("image must have shape [B,3,H,W]") + old_height, old_width = image.shape[-2:] + scale = max(height / old_height, width / old_width) + resized_height = max(math.ceil(old_height * scale), height) + resized_width = max(math.ceil(old_width * scale), width) + resized = F.interpolate( + image.to(torch.uint8), size=(resized_height, resized_width), mode="bilinear", align_corners=False + ) + top = int(round((resized_height - height) / 2.0)) + left = int(round((resized_width - width) / 2.0)) + return resized[:, :, top : top + height, left : left + width].float().div(255.0).unsqueeze(2) diff --git a/telefuser/pipelines/lingbot_video/denoising.py b/telefuser/pipelines/lingbot_video/denoising.py new file mode 100644 index 0000000..68d97c9 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -0,0 +1,286 @@ +"""Reference LingBot-Video denoising loop primitives. + +Sampling order and TI2V condition behavior are adapted from the Apache-2.0 +licensed upstream LingBot-Video implementation. +""" + +from __future__ import annotations + +from collections.abc import Callable +from functools import partial + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.distributed.device_mesh import ( + create_device_mesh_from_config, + get_cfg_rank, + get_cfg_world_size, + get_ulysses_group, +) +from telefuser.distributed.fsdp import shard_model_fsdp2_inference +from telefuser.distributed.parallel_shard import cfg_parallel_unshard +from telefuser.platforms import current_platform +from telefuser.utils.logging import logger + + +def reinject_ti2v_condition(latent: torch.Tensor, condition: torch.Tensor, mask: torch.Tensor) -> torch.Tensor: + """Replace the masked temporal prefix with the clean VAE condition latent.""" + if latent.ndim != 5 or condition.ndim != 5 or mask.ndim not in {1, latent.ndim}: + raise ValueError("latent, condition, and mask shapes are incompatible") + if latent.shape[:2] != condition.shape[:2] or latent.shape[3:] != condition.shape[3:]: + raise ValueError("condition must match latent batch, channel, and spatial dimensions") + if condition.shape[2] > latent.shape[2]: + raise ValueError("condition temporal prefix is longer than the latent") + if condition.shape != latent.shape: + if mask.ndim != 1 or not bool(mask[: condition.shape[2]].all()) or bool(mask[condition.shape[2] :].any()): + raise ValueError("prefix condition requires a matching prefix-only temporal mask") + output = latent.clone() + output[:, :, : condition.shape[2]] = condition.to(dtype=latent.dtype) + return output + expanded = mask.to(device=latent.device, dtype=torch.bool) + if expanded.ndim == 1 and latent.ndim >= 3: + expanded = expanded.reshape(1, 1, -1, *([1] * (latent.ndim - 3))) + return torch.where(expanded, condition, latent) + + +def transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.dtype) -> torch.Tensor: + """Match the upstream BF16/FP16 sigma rounding before DiT time embedding.""" + sigma = timestep.float() / 1000.0 + if transformer_dtype in {torch.bfloat16, torch.float16}: + sigma = sigma.to(transformer_dtype) + return (sigma * 1000.0).float() + + +def _batch_cfg_prompt_inputs( + positive_embeds: torch.Tensor, + negative_embeds: torch.Tensor, + positive_mask: torch.Tensor | None, + negative_mask: torch.Tensor | None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Pad positive and negative prompt conditions for one batched CFG forward.""" + if positive_embeds.ndim != 3 or negative_embeds.ndim != 3: + raise ValueError("CFG prompt embeddings must have shape [batch, sequence, hidden_size]") + if positive_embeds.shape[0] != negative_embeds.shape[0] or positive_embeds.shape[2] != negative_embeds.shape[2]: + raise ValueError("positive and negative CFG embeddings must have matching batch and hidden dimensions") + + def pad( + embeds: torch.Tensor, + mask: torch.Tensor | None, + target_length: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + if mask is None: + mask = torch.ones(embeds.shape[:2], dtype=torch.bool, device=embeds.device) + if mask.shape != embeds.shape[:2]: + raise ValueError("CFG prompt attention mask must match the embedding batch and sequence dimensions") + padding = target_length - embeds.shape[1] + if padding: + embeds = torch.nn.functional.pad(embeds, (0, 0, 0, padding)) + mask = torch.nn.functional.pad(mask, (0, padding), value=False) + return embeds, mask + + target_length = max(positive_embeds.shape[1], negative_embeds.shape[1]) + positive_embeds, positive_mask = pad(positive_embeds, positive_mask, target_length) + negative_embeds, negative_mask = pad(negative_embeds, negative_mask, target_length) + return torch.cat((positive_embeds, negative_embeds)), torch.cat((positive_mask, negative_mask)) + + +def denoise_lingbot_video( + latent: torch.Tensor, + timesteps: torch.Tensor, + predict: Callable[[torch.Tensor, torch.Tensor], torch.Tensor], + step: Callable[[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor], + *, + condition: torch.Tensor | None = None, + condition_mask: torch.Tensor | None = None, +) -> torch.Tensor: + """Run a scheduler-agnostic denoising loop with optional TI2V reinjection.""" + current = latent + for timestep in timesteps: + if condition is not None: + if condition_mask is None: + raise ValueError("condition_mask is required when condition is provided") + current = reinject_ti2v_condition(current, condition, condition_mask) + prediction = predict(current, timestep) + current = step(prediction, timestep, current) + if condition is not None and condition_mask is not None: + current = reinject_ti2v_condition(current, condition, condition_mask) + return current + + +class LingBotVideoDenoisingStage(BaseStage): + """Dense LingBot-Video denoising stage with source-equivalent CFG order.""" + + def __init__( + self, + name: str, + module_manager: ModuleManager, + model_runtime_config: ModelRuntimeConfig, + *, + batch_cfg: bool = False, + ) -> None: + super().__init__(name, model_runtime_config) + transformer = module_manager.fetch_module("transformer") + self.transformer = transformer + self.scheduler = module_manager.fetch_module("scheduler") + set_attention_config = getattr(transformer, "set_attention_config", None) + if callable(set_attention_config): + set_attention_config(model_runtime_config.attention_config) + self.model_names = ["transformer"] + self.batch_cfg = batch_cfg + # Denoising repeatedly invokes the same FSDP graph. Releasing the CUDA + # allocator cache after every step forces costly weight-buffer + # reallocations that the source runner keeps cached for the full loop. + self.empty_cache_after_call = False + + def parallel_models(self) -> None: + """Attach Ulysses SP and optional per-block FSDP to the DiT stage.""" + if self.device.type == "cuda": + # Match the source runner's FP32 matmul policy. This primarily + # affects the MoE router projections, which intentionally run in + # FP32 and are otherwise much slower on Ampere-and-newer GPUs. + torch.backends.cuda.matmul.allow_tf32 = True + torch.set_float32_matmul_precision("high") + + parallel_config = self.model_runtime_config.parallel_config + if parallel_config.world_size == 1: + return + self.transformer.device_mesh = create_device_mesh_from_config(parallel_config) + cfg_world_size = get_cfg_world_size(self.transformer.device_mesh) + if cfg_world_size not in {1, 2}: + raise ValueError(f"LingBot CFG parallel requires degree 1 or 2, got {cfg_world_size}") + if cfg_world_size > 1 and self.batch_cfg: + raise ValueError("LingBot CFG parallel and batch CFG are mutually exclusive") + if parallel_config.sp_ulysses_degree > 1: + self.transformer.set_ulysses_group(get_ulysses_group(self.transformer.device_mesh)) + logger.info(f"enabled LingBot Ulysses SP degree={parallel_config.sp_ulysses_degree}") + if parallel_config.enable_fsdp: + if self.model_runtime_config.offload_config.offload_type != WeightOffloadType.NO_CPU_OFFLOAD: + raise ValueError("LingBot FSDP inference cannot be combined with model CPU offload") + ignored_states = [ + parameter for parameter in self.transformer.parameters() if parameter.dtype == torch.float32 + ] + if ignored_states: + logger.info(f"retaining {len(ignored_states)} LingBot parameters with a non-runtime dtype outside FSDP") + logger.info(f"enabled LingBot block FSDP2 for {self.name}") + self.transformer = shard_model_fsdp2_inference( + module=self.transformer, + device_mesh=self.transformer.device_mesh, + wrap_module_names=self.transformer.get_fsdp_module_names(), + ignored_states=ignored_states, + ) + self.onload_models_flag = True + current_platform.empty_cache() + + @with_model_offload(["transformer"]) + @torch.no_grad() + def predict_noise_with_cfg( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + positive_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None = None, + positive_attention_mask: torch.Tensor | None = None, + negative_attention_mask: torch.Tensor | None = None, + guidance_scale: float = 1.0, + ) -> torch.Tensor: + """Run positive then negative CFG under transformer compute autocast.""" + return self._predict_noise_with_cfg( + latents, + timestep, + positive_prompt_embeds, + negative_prompt_embeds, + positive_attention_mask, + negative_attention_mask, + guidance_scale, + ) + + def _predict_noise_with_cfg( + self, + latents: torch.Tensor, + timestep: torch.Tensor, + positive_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + positive_attention_mask: torch.Tensor | None, + negative_attention_mask: torch.Tensor | None, + guidance_scale: float, + ) -> torch.Tensor: + """Run CFG without crossing the stage lifecycle boundary.""" + transformer_dtype = self.transformer.patch_embedder.weight.dtype + model_timestep = transformer_timestep(timestep, transformer_dtype) + autocast_enabled = latents.device.type == "cuda" and transformer_dtype in {torch.bfloat16, torch.float16} + with torch.autocast(device_type=latents.device.type, dtype=transformer_dtype, enabled=autocast_enabled): + if guidance_scale <= 1.0: + return self.transformer( + latents, model_timestep, positive_prompt_embeds, positive_attention_mask + ).float() + if negative_prompt_embeds is None: + raise ValueError("negative_prompt_embeds is required when guidance_scale is greater than 1") + device_mesh = getattr(self.transformer, "device_mesh", None) + cfg_world_size = get_cfg_world_size(device_mesh) + if cfg_world_size > 1: + cfg_rank = get_cfg_rank(device_mesh) + prompt_embeds = positive_prompt_embeds if cfg_rank == 0 else negative_prompt_embeds + attention_mask = positive_attention_mask if cfg_rank == 0 else negative_attention_mask + prediction = self.transformer(latents, model_timestep, prompt_embeds, attention_mask).float() + positive, negative = cfg_parallel_unshard(device_mesh, [prediction])[0].chunk(2) + elif self.batch_cfg: + combined_latents = torch.cat((latents, latents)) + combined_timestep = torch.cat((model_timestep, model_timestep)) + combined_embeds, combined_mask = _batch_cfg_prompt_inputs( + positive_prompt_embeds, + negative_prompt_embeds, + positive_attention_mask, + negative_attention_mask, + ) + positive, negative = ( + self.transformer(combined_latents, combined_timestep, combined_embeds, combined_mask) + .float() + .chunk(2) + ) + else: + positive = self.transformer( + latents, model_timestep, positive_prompt_embeds, positive_attention_mask + ).float() + negative = self.transformer( + latents, model_timestep, negative_prompt_embeds, negative_attention_mask + ).float() + return negative + guidance_scale * (positive - negative) + + @with_model_offload(["transformer"]) + @torch.no_grad() + def denoise( + self, + latent: torch.Tensor, + positive_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + positive_attention_mask: torch.Tensor | None, + negative_attention_mask: torch.Tensor | None, + guidance_scale: float, + num_inference_steps: int, + shift: float, + condition: torch.Tensor | None = None, + condition_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """Run the complete sampling loop in one worker invocation.""" + self.scheduler.set_timesteps(num_inference_steps, device=latent.device, shift=shift) + for timestep in self.scheduler.timesteps: + if condition is not None: + if condition_mask is None: + raise ValueError("condition_mask is required when condition is provided") + latent = reinject_ti2v_condition(latent, condition, condition_mask) + prediction = self._predict_noise_with_cfg( + latent, + timestep.expand(latent.shape[0]), + positive_prompt_embeds, + negative_prompt_embeds, + positive_attention_mask, + negative_attention_mask, + guidance_scale, + ) + latent = self.scheduler.step(prediction, timestep, latent) + if condition is not None and condition_mask is not None: + latent = reinject_ti2v_condition(latent, condition, condition_mask) + return latent diff --git a/telefuser/pipelines/lingbot_video/pipeline.py b/telefuser/pipelines/lingbot_video/pipeline.py new file mode 100644 index 0000000..0b359bb --- /dev/null +++ b/telefuser/pipelines/lingbot_video/pipeline.py @@ -0,0 +1,272 @@ +"""Precision-first LingBot-Video sampling pipeline. + +Sampling order and TI2V conditioning behavior are adapted from the +Apache-2.0 licensed upstream LingBot-Video implementation. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch +import torch.distributed as dist + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler +from telefuser.worker.parallel_worker import ParallelWorker + +from .data import ( + LingBotVideoModelConfig, + LingBotVideoRequest, + default_negative_caption, + preprocess_ti2v_image, + validate_frame_count, +) +from .denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from .text_encoding import LingBotVideoTextEncodingStage +from .vae import ( + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + first_frame_condition_mask, + latent_shape, +) + + +@dataclass +class LingBotVideoPipelineConfig: + """Configuration for LingBot-Video model stages and sampling.""" + + model: LingBotVideoModelConfig = field(default_factory=LingBotVideoModelConfig) + text_encoding_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + dit_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + vae_config: ModelRuntimeConfig = field(default_factory=ModelRuntimeConfig) + guidance_scale: float = 3.0 + num_inference_steps: int = 50 + shift: float = 3.0 + batch_cfg: bool = False + enable_denoising_parallel: bool = False + + def __post_init__(self) -> None: + if self.guidance_scale < 0 or self.num_inference_steps <= 0 or self.shift <= 0: + raise ValueError("guidance_scale must be non-negative and steps must be positive") + + +@dataclass(frozen=True) +class LingBotVideoPromptConditions: + """Text conditions produced for one source-order CFG generation.""" + + positive_prompt_embeds: torch.Tensor + negative_prompt_embeds: torch.Tensor | None + positive_attention_mask: torch.Tensor + negative_attention_mask: torch.Tensor | None + has_visual_condition: bool + + +@dataclass(frozen=True) +class LingBotVideoGeneration: + """Generated output together with prompt conditions that produced it.""" + + output: torch.Tensor + prompt_conditions: LingBotVideoPromptConditions + + +class LingBotVideoPipeline(BasePipeline): + """LingBot-Video pipeline initialized from ModuleManager-owned components.""" + + def __init__(self, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat16) -> None: + super().__init__(device=device, torch_dtype=torch_dtype) + self.config: LingBotVideoPipelineConfig | None = None + self.text_stage: LingBotVideoTextEncodingStage | None = None + self.denoising_stage: LingBotVideoDenoisingStage | None = None + self.vae_encode_stage: LingBotVideoVAEEncodeStage | None = None + self.vae_decode_stage: LingBotVideoVAEDecodeStage | None = None + self.scheduler: FlowUniPCMultistepScheduler | None = None + self.variant: str = "dense" + self._parallel_denoising_stage_template: LingBotVideoDenoisingStage | None = None + self._stopped = False + + def init(self, module_manager: ModuleManager, config: LingBotVideoPipelineConfig) -> None: + """Initialize all LingBot-Video stages from a module manager.""" + self._model_info = module_manager.get_model_info() + self.config = config + self.variant = config.model.variant + self._stopped = False + + self.text_stage = LingBotVideoTextEncodingStage("text_encoder", module_manager, config.text_encoding_config) + denoising_stage = LingBotVideoDenoisingStage( + "transformer", module_manager, config.dit_config, batch_cfg=config.batch_cfg + ) + self.denoising_stage = denoising_stage + self.vae_encode_stage = LingBotVideoVAEEncodeStage("vae_encode", module_manager, config.vae_config) + self.vae_decode_stage = LingBotVideoVAEDecodeStage("vae_decode", module_manager, config.vae_config) + self.scheduler = module_manager.fetch_module("scheduler") + + if config.enable_denoising_parallel and not dist.is_initialized(): + self._parallel_denoising_stage_template = denoising_stage + self.denoising_stage = ParallelWorker(denoising_stage) + else: + self._parallel_denoising_stage_template = None + denoising_stage.parallel_models() + + def _get_stages(self) -> list[object]: + """Return independently managed stages, including a distributed DiT worker.""" + return [ + stage + for stage in (self.text_stage, self.denoising_stage, self.vae_encode_stage, self.vae_decode_stage) + if stage is not None + ] + + def stop(self) -> None: + """Close any owned distributed stage workers during service shutdown.""" + self._stopped = True + for stage in self._get_stages(): + close = getattr(stage, "close", None) + if callable(close): + close() + + @staticmethod + def _resolve_stage_result(value: object) -> torch.Tensor: + """Resolve a ParallelWorker deferred result while preserving direct-stage outputs.""" + result = value() if callable(value) else value + if not isinstance(result, torch.Tensor): + raise TypeError(f"LingBot stage returned {type(result).__name__}, expected a tensor") + return result + + @staticmethod + def validate_request(request: LingBotVideoRequest) -> LingBotVideoRequest: + """Validate and return a request before allocating model resources.""" + validate_frame_count(request.num_frames) + return request + + def release_gpu_resources(self) -> None: + """Offload attached stages so a separately loaded refiner can use the GPU.""" + for stage in (self.text_stage, self.denoising_stage, self.vae_encode_stage, self.vae_decode_stage): + close = getattr(stage, "close", None) + if callable(close): + close() + continue + offload_models = getattr(stage, "offload_models", None) + if callable(offload_models): + offload_models() + stage.onload_models_flag = False + if torch.cuda.is_available(): + torch.cuda.empty_cache() + + def _restore_released_denoising_worker(self) -> None: + """Recreate a distributed base worker released before loading the refiner.""" + if self._stopped: + raise RuntimeError("LingBot-Video runtime has been stopped") + if not isinstance(self.denoising_stage, ParallelWorker) or not self.denoising_stage.closed: + return + if self._parallel_denoising_stage_template is None: + raise RuntimeError("LingBot-Video distributed denoising stage cannot be restored") + self.denoising_stage = ParallelWorker(self._parallel_denoising_stage_template) + + def __call__( + self, + request: LingBotVideoRequest, + *, + negative_caption: str | None = None, + generator: torch.Generator | None = None, + decode: bool = True, + ) -> torch.Tensor: + """Run source-order sampling and return only the generated output.""" + return self.generate(request, negative_caption=negative_caption, generator=generator, decode=decode).output + + def generate( + self, + request: LingBotVideoRequest, + *, + negative_caption: str | None = None, + generator: torch.Generator | None = None, + decode: bool = True, + ) -> LingBotVideoGeneration: + """Run source-order sampling and retain its exact CFG prompt conditions.""" + self.validate_request(request) + self._restore_released_denoising_worker() + if self.config is None or self.text_stage is None or self.denoising_stage is None or self.scheduler is None: + raise RuntimeError("LingBot-Video runtime has not been configured") + if negative_caption is None: + negative_caption = default_negative_caption(request.num_frames) + condition_pixels: torch.Tensor | None = None + vision_images: list[object] | None = None + if request.image is not None: + if self.vae_encode_stage is None or not isinstance(request.image, torch.Tensor): + raise ValueError("TI2V requires a VAE encode stage and an RGB tensor image") + source_image = request.image + if source_image.ndim == 5: + if source_image.shape[2] < 1: + raise ValueError("TI2V source image must include a frame") + source_image = source_image[:, :, 0] + condition_pixels = preprocess_ti2v_image(source_image, height=request.height, width=request.width) + vision_images = [self.text_stage.prepare_ti2v_vlm_image(condition_pixels)] + positive, positive_mask = self.text_stage.encode(request.caption, images=vision_images) + if self.config.guidance_scale > 1.0: + negative, negative_mask = self.text_stage.encode(negative_caption, images=vision_images) + else: + negative, negative_mask = None, None + prompt_conditions = LingBotVideoPromptConditions( + positive_prompt_embeds=positive, + negative_prompt_embeds=negative, + positive_attention_mask=positive_mask, + negative_attention_mask=negative_mask, + has_visual_condition=condition_pixels is not None, + ) + frames, latent_height, latent_width = latent_shape(request.num_frames, request.height, request.width) + condition = None + condition_mask = None + if condition_pixels is not None: + condition = self.vae_encode_stage.encode(condition_pixels, generator=generator) + condition_mask = first_frame_condition_mask(frames, device=self.denoising_stage.device) + latent = torch.randn( + 1, + self.config.model.in_channels, + frames, + latent_height, + latent_width, + device=self.denoising_stage.device, + dtype=torch.float32, + generator=generator, + ) + if self.config.enable_denoising_parallel: + latent = self._resolve_stage_result( + self.denoising_stage.denoise( + latent, + positive, + negative, + positive_mask, + negative_mask, + self.config.guidance_scale, + self.config.num_inference_steps, + self.config.shift, + condition, + condition_mask, + ) + ) + else: + self.scheduler.set_timesteps(self.config.num_inference_steps, device=latent.device, shift=self.config.shift) + for step in self.scheduler.timesteps: + if condition is not None: + latent = reinject_ti2v_condition(latent, condition, condition_mask) + prediction = self._resolve_stage_result( + self.denoising_stage.predict_noise_with_cfg( + latent, + step.expand(latent.shape[0]), + positive, + negative, + positive_mask, + negative_mask, + self.config.guidance_scale, + ) + ) + latent = self.scheduler.step(prediction, step, latent) + if condition is not None: + latent = reinject_ti2v_condition(latent, condition, condition_mask) + if not decode: + return LingBotVideoGeneration(latent, prompt_conditions) + if self.vae_decode_stage is None: + raise RuntimeError("decode=True requires a configured VAE decode stage") + output = self._resolve_stage_result(self.vae_decode_stage.decode(latent)) + return LingBotVideoGeneration(output, prompt_conditions) diff --git a/telefuser/pipelines/lingbot_video/refiner.py b/telefuser/pipelines/lingbot_video/refiner.py new file mode 100644 index 0000000..88f1bc6 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/refiner.py @@ -0,0 +1,316 @@ +"""Low-noise refiner schedule and latent handoff primitives.""" +# Training-aligned frame selection below is adapted from the Apache-2.0 +# licensed upstream LingBot-Video ``utils.py`` implementation. + +from __future__ import annotations + +from collections.abc import Callable +from pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +from telefuser.platforms import current_platform + +from .denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from .vae import ( + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + first_frame_condition_mask, +) + + +def compute_refiner_sigmas( + *, + sigma_max: float, + sigma_min: float, + num_inference_steps: int, + shift: float, + t_thresh: float | None, + tail_steps: int = 0, +) -> np.ndarray | None: + """Match the official refiner's low-noise sigma-tail construction.""" + if t_thresh is None: + return None + if not 0 < t_thresh <= 1 or num_inference_steps < 1 or tail_steps < 0: + raise ValueError("invalid refiner sigma configuration") + base = np.linspace(sigma_max, sigma_min, num_inference_steps + 1)[:-1] + shifted = shift * base / (1.0 + (shift - 1.0) * base) + sigmas = shifted[shifted <= t_thresh + 1e-6] + if not len(sigmas) or abs(float(sigmas[0]) - t_thresh) > 1e-6: + sigmas = np.concatenate([[t_thresh], sigmas]) + if tail_steps: + tail = np.linspace(float(sigmas[-1]), min(sigma_min, float(sigmas[-1])), tail_steps + 2)[1:-1] + sigmas = np.concatenate([sigmas, tail]) + if np.any(~np.isfinite(sigmas)) or np.any(sigmas < 0) or np.any(sigmas > 1) or np.any(np.diff(sigmas) >= 0): + raise ValueError("refiner sigma schedule must be finite, descending, and within [0,1]") + return sigmas.astype(np.float32) + + +def prepare_refiner_latent(x_up: torch.Tensor, noise: torch.Tensor, t_thresh: float | torch.Tensor) -> torch.Tensor: + """Mix encoded low-resolution video with noise at the refiner threshold.""" + if x_up.shape != noise.shape: + raise ValueError("x_up and noise must have matching shapes") + threshold = torch.as_tensor(t_thresh, device=x_up.device, dtype=x_up.dtype) + while threshold.ndim < x_up.ndim: + threshold = threshold.unsqueeze(-1) + return (1.0 - threshold) * x_up + threshold * noise + + +def compute_training_frame_budget( + num_source_frames: int, source_fps: float, *, sample_fps: int = 24, vae_tc: int = 4 +) -> tuple[int, float, int]: + """Return the upstream training-aligned frame and temporal-latent budgets.""" + if num_source_frames <= 0: + return 1, 0.0, 1 + if source_fps <= 0 or sample_fps <= 0 or vae_tc <= 0: + raise ValueError("source_fps, sample_fps, and vae_tc must be positive") + raw_frames = int(num_source_frames / source_fps * sample_fps) if source_fps > sample_fps else num_source_frames + sample_frames = max(((raw_frames - 1) // vae_tc) * vae_tc + 1, 1) + vae_fps = sample_frames / num_source_frames * source_fps + return int(sample_frames), float(vae_fps), (sample_frames - 1) // vae_tc + 1 + + +def compute_training_aligned_indices(num_source_frames: int, sample_frames: int) -> torch.Tensor: + """Select training-aligned indices and pad short inputs with their last frame.""" + if sample_frames <= 0: + return torch.empty(0, dtype=torch.long) + if num_source_frames <= 0: + return torch.zeros(sample_frames, dtype=torch.long) + if num_source_frames >= sample_frames: + return torch.from_numpy(np.linspace(0, num_source_frames - 1, sample_frames, dtype=int)).long() + return torch.cat( + (torch.arange(num_source_frames), torch.full((sample_frames - num_source_frames,), num_source_frames - 1)) + ) + + +def prepare_refiner_video( + video: torch.Tensor, + *, + source_fps: float, + height: int, + width: int, + sample_fps: int = 24, + vae_tc: int = 4, + max_frames: int | None = None, +) -> tuple[torch.Tensor, dict[str, int | float | bool | None]]: + """Sample and bicubically resize an in-memory RGB video for the refiner VAE.""" + if video.ndim != 5 or video.shape[1] != 3: + raise ValueError("video must have shape [B,3,F,H,W]") + if height <= 0 or width <= 0: + raise ValueError("height and width must be positive") + if max_frames is not None and max_frames < 1: + raise ValueError("max_frames must be positive when set") + total_frames = video.shape[2] + sample_frames, vae_fps, temporal_latents = compute_training_frame_budget( + total_frames, source_fps, sample_fps=sample_fps, vae_tc=vae_tc + ) + uncapped_frames = sample_frames + truncated = max_frames is not None and sample_frames > max_frames + if truncated: + sample_frames = int(max_frames) + vae_fps = sample_frames / total_frames * source_fps + temporal_latents = (sample_frames - 1) // vae_tc + 1 + indices = compute_training_aligned_indices(total_frames, sample_frames).to(video.device) + sampled = video.index_select(2, indices) + batch, channels, frames, input_height, input_width = sampled.shape + flat = sampled.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, input_height, input_width) + resized = F.interpolate(flat, size=(height, width), mode="bicubic", align_corners=False).clamp(0.0, 1.0) + prepared = resized.reshape(batch, frames, channels, height, width).permute(0, 2, 1, 3, 4).contiguous() + return prepared, { + "src_fps": float(source_fps), + "sample_frame": int(sample_frames), + "sample_frame_uncapped": int(uncapped_frames), + "max_frames": max_frames, + "truncated_by_max_frames": bool(truncated), + "vae_fps": float(vae_fps), + "t_vae": int(temporal_latents), + "num_source_frames": int(total_frames), + "align_to_training": True, + } + + +def load_refiner_video_file( + path: str | Path, + *, + height: int, + width: int, + sample_fps: int = 24, + vae_tc: int = 4, + max_frames: int | None = None, +) -> tuple[torch.Tensor, dict[str, int | float | bool | None]]: + """Decode an MP4 with PyAV and apply the upstream refiner sampling contract.""" + try: + import av + except ImportError as exc: + raise RuntimeError("MP4 refiner handoff requires the optional PyAV dependency") from exc + + container = av.open(str(path)) + try: + stream = next(iter(container.streams.video), None) + if stream is None: + raise ValueError(f"video has no video stream: {path}") + source_fps = float(stream.average_rate) if stream.average_rate is not None else 0.0 + frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)] + finally: + container.close() + if not frames: + raise ValueError(f"video has no decoded frames: {path}") + video = torch.from_numpy(np.stack(frames)).permute(3, 0, 1, 2).unsqueeze(0).float().div(255.0) + return prepare_refiner_video( + video, + source_fps=source_fps, + height=height, + width=width, + sample_fps=sample_fps, + vae_tc=vae_tc, + max_frames=max_frames, + ) + + +def load_refiner_first_frame( + path: str | Path, + *, + target_height: int, + target_width: int, + geometry_height: int, + geometry_width: int, +) -> torch.Tensor: + """Load the TI2V refiner frame zero with the upstream low-resolution geometry.""" + from PIL import Image + + image = Image.open(path).convert("RGB") + image_width, image_height = image.size + geometry_aspect = float(geometry_width) / float(geometry_height) + image_aspect = float(image_width) / float(image_height) + if image_aspect > geometry_aspect: + crop_height = image_height + crop_width = max(1, int(round(crop_height * geometry_aspect))) + left, top = int(round((image_width - crop_width) / 2.0)), 0 + else: + crop_width = image_width + crop_height = max(1, int(round(crop_width / geometry_aspect))) + left, top = 0, int(round((image_height - crop_height) / 2.0)) + crop = image.crop((left, top, left + crop_width, top + crop_height)) + resized = crop.resize((target_width, target_height), resample=Image.BICUBIC) + pixels = torch.from_numpy(np.asarray(resized, dtype=np.float32) / 255.0).permute(2, 0, 1).unsqueeze(0) + return pixels.permute(1, 0, 2, 3).unsqueeze(0).contiguous() + + +class LingBotVideoRefinerStage: + """In-memory low-noise refiner runtime independent of the base DiT stage.""" + + def __init__( + self, + *, + denoising_stage: LingBotVideoDenoisingStage, + vae_encode_stage: LingBotVideoVAEEncodeStage, + vae_decode_stage: LingBotVideoVAEDecodeStage, + scheduler: object, + ) -> None: + self.denoising_stage = denoising_stage + self.vae_encode_stage = vae_encode_stage + self.vae_decode_stage = vae_decode_stage + self.scheduler = scheduler + + @staticmethod + def _resolve_stage_result(value: object) -> torch.Tensor: + """Resolve deferred ParallelWorker outputs while accepting direct stage tensors.""" + result = value() if callable(value) else value + if not isinstance(result, torch.Tensor): + raise TypeError(f"LingBot refiner stage returned {type(result).__name__}, expected a tensor") + return result + + def close(self) -> None: + """Close an owned distributed denoising worker after the refiner stage completes.""" + close = getattr(self.denoising_stage, "close", None) + if callable(close): + close() + + def _offload_vae_between_stages(self) -> None: + """Release VAE weights and allocator blocks before high-resolution DiT execution.""" + for stage in (self.vae_encode_stage, self.vae_decode_stage): + offload_models = getattr(stage, "offload_models", None) + if callable(offload_models): + offload_models() + current_platform.empty_cache() + + def refine( + self, + lowres_video: torch.Tensor, + positive_prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor, + positive_attention_mask: torch.Tensor | None, + negative_attention_mask: torch.Tensor | None, + *, + num_inference_steps: int, + guidance_scale: float, + shift: float, + t_thresh: float, + tail_steps: int = 2, + clean_first_frame: torch.Tensor | None = None, + generator: torch.Generator | None = None, + noise: torch.Tensor | None = None, + before_decode: Callable[[], None] | None = None, + ) -> torch.Tensor: + """Refine an in-memory RGB video using official low-noise initialization.""" + x_up = self.vae_encode_stage.encode(lowres_video, generator=generator) + condition = None + if clean_first_frame is not None: + if clean_first_frame.ndim == 4: + clean_first_frame = clean_first_frame.unsqueeze(2) + clean_latent = self.vae_encode_stage.encode(clean_first_frame, generator=generator) + condition = clean_latent[:, :, :1].contiguous() + x_up = x_up.clone() + x_up[:, :, :1] = condition.to(x_up.dtype) + # A 1088p VAE encode may leave several GiB of allocator blocks on + # rank zero. The SP refiner is a separate stage, so it must not share + # that residency with the long-sequence MoE denoising worker. + self._offload_vae_between_stages() + if noise is None: + noise = torch.randn(x_up.shape, dtype=x_up.dtype, device=x_up.device, generator=generator) + elif noise.shape != x_up.shape: + raise ValueError("injected refiner noise must match the encoded low-resolution latent") + else: + noise = noise.to(device=x_up.device, dtype=x_up.dtype) + latent = prepare_refiner_latent(x_up, noise, t_thresh) + sigmas = compute_refiner_sigmas( + sigma_max=float(self.scheduler.sigma_max), + sigma_min=float(self.scheduler.sigma_min), + num_inference_steps=num_inference_steps, + shift=shift, + t_thresh=t_thresh, + tail_steps=tail_steps, + ) + if sigmas is None: + self.scheduler.set_timesteps(num_inference_steps, device=latent.device, shift=shift) + else: + self.scheduler.set_timesteps(len(sigmas), device=latent.device, sigmas=sigmas, shift=1.0) + condition_mask = ( + first_frame_condition_mask(latent.shape[2], device=latent.device) if condition is not None else None + ) + for timestep in self.scheduler.timesteps: + if condition is not None: + latent = reinject_ti2v_condition(latent, condition, condition_mask) + prediction = self._resolve_stage_result( + self.denoising_stage.predict_noise_with_cfg( + latent, + timestep.expand(latent.shape[0]), + positive_prompt_embeds, + negative_prompt_embeds, + positive_attention_mask, + negative_attention_mask, + guidance_scale, + ) + ) + latent = self.scheduler.step(prediction, timestep, latent) + if condition is not None: + latent = reinject_ti2v_condition(latent, condition, condition_mask) + # The native SP worker owns the refiner weights. Release it before + # loading the VAE decoder, otherwise the two high-memory stages can + # overlap on rank zero. + self.close() + if before_decode is not None: + before_decode() + return self._resolve_stage_result(self.vae_decode_stage.decode(latent)) diff --git a/telefuser/pipelines/lingbot_video/text_encoding.py b/telefuser/pipelines/lingbot_video/text_encoding.py new file mode 100644 index 0000000..aabe86c --- /dev/null +++ b/telefuser/pipelines/lingbot_video/text_encoding.py @@ -0,0 +1,140 @@ +"""Qwen3-VL prompt encoding for LingBot-Video. + +Prompt templates and Qwen3-VL image preparation behavior are adapted from the +Apache-2.0 licensed upstream LingBot-Video implementation. +""" + +from __future__ import annotations + +from typing import Any + +import torch +from PIL import Image + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + +from .data import smart_resize + +PROMPT_TEMPLATE = ( + "<|im_start|>system\nGiven a user input that may include a text prompt alone, " + "a text prompt with an image reference, or a text prompt with a video reference " + 'or a video reference alone, generate an "Enhanced prompt" that provides detailed ' + "visual descriptions suitable for video generation. Evaluate the level of detail " + "in the user's input: if it is simple, enrich it by adding specifics about colors, " + "shapes, sizes, textures, lighting, motion dynamics, camera movement, temporal " + "progression, and spatial relationships to create vivid, concrete, and temporally " + "coherent scenes to create vivid and concrete scenes. Please generate only the " + "enhanced description for the prompt below and avoid including any additional " + "commentary or evaluations:<|im_end|>\n<|im_start|>user\n{}<|im_end|>\n" + "<|im_start|>assistant\n" +) +IMAGE_PROMPT_TEMPLATE = "<|vision_start|><|image_pad|><|vision_end|>" +_VISION_SPATIAL_MERGE_SIZE = 2 +_VISION_MIN_TOKEN_COUNT = 4 +_VISION_MAX_TOKEN_COUNT = 16384 + + +class LingBotVideoTextEncodingStage(BaseStage): + """Encode structured captions and optional TI2V images with Qwen3-VL.""" + + def __init__( + self, + name: str, + module_manager: ModuleManager, + model_runtime_config: ModelRuntimeConfig, + *, + token_length: int = 37698, + hidden_state_skip_layer: int | None = 0, + ) -> None: + super().__init__(name, model_runtime_config) + self.text_encoder = module_manager.fetch_module("text_encoder") + self.processor = module_manager.fetch_module("processor") + self.token_length = token_length + self.hidden_state_skip_layer = hidden_state_skip_layer + self._crop_start: int | None = None + self.model_names = ["text_encoder"] + + def _crop_prefix_length(self) -> int: + if self._crop_start is None: + marker = "<|USER_INPUT_MARKER|>" + prefix = PROMPT_TEMPLATE.format(marker).split(marker, maxsplit=1)[0] + inputs = self.processor(text=prefix, images=None, videos=None, return_tensors="pt") + self._crop_start = int(inputs["input_ids"].shape[1]) + return self._crop_start + + def _vision_patch_size(self) -> int: + """Return the Qwen vision patch size using the source lookup order.""" + for obj in ( + getattr(getattr(self.text_encoder, "config", None), "vision_config", None), + getattr(getattr(self.processor, "image_processor", None), "config", None), + getattr(self.processor, "image_processor", None), + ): + patch_size = getattr(obj, "patch_size", None) + if patch_size is not None: + return int(patch_size) + return 16 + + def prepare_ti2v_vlm_image(self, pixel_values: torch.Tensor) -> Image.Image: + """Convert a condition frame to the source-equivalent Qwen3-VL image input.""" + if ( + pixel_values.ndim != 5 + or pixel_values.shape[0] != 1 + or pixel_values.shape[1] != 3 + or pixel_values.shape[2] < 1 + ): + raise ValueError("TI2V pixels must have shape [1,3,F,H,W]") + frame = pixel_values[0, :, 0].detach().cpu().clamp(0.0, 1.0) + image = Image.fromarray(frame.permute(1, 2, 0).mul(255).byte().numpy(), mode="RGB") + patch_factor = self._vision_patch_size() * _VISION_SPATIAL_MERGE_SIZE + width, height = image.size + resized_height, resized_width = smart_resize( + height, + width, + factor=patch_factor, + min_pixels=_VISION_MIN_TOKEN_COUNT * patch_factor**2, + max_pixels=_VISION_MAX_TOKEN_COUNT * patch_factor**2, + ) + return image.resize((resized_width, resized_height)) + + @with_model_offload(["text_encoder"]) + @torch.no_grad() + def encode( + self, + prompt: str | list[str], + *, + images: Any | None = None, + video_metadata: Any | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Return cropped prompt embeddings and their attention mask.""" + prompts = [prompt] if isinstance(prompt, str) else prompt + visual_prefix = IMAGE_PROMPT_TEMPLATE if images is not None else "" + texts = [PROMPT_TEMPLATE.format(visual_prefix + item) for item in prompts] + inputs = self.processor( + text=texts, + images=images, + videos=None, + video_metadata=video_metadata, + do_resize=False, + truncation=True, + max_length=self.token_length, + padding="longest", + return_tensors="pt", + ).to(self.device) + output = self.text_encoder( + **inputs, + output_hidden_states=self.hidden_state_skip_layer is not None, + ) + embeddings = ( + output.hidden_states[-(self.hidden_state_skip_layer + 1)] + if self.hidden_state_skip_layer is not None + else output.last_hidden_state + ) + mask = inputs["attention_mask"] + crop_start = self._crop_prefix_length() + embeddings, mask = embeddings[:, crop_start:], mask[:, crop_start:] + if embeddings.shape[0] == 1: + length = int(mask[0].sum().item()) + embeddings, mask = embeddings[:, :length], mask[:, :length] + return embeddings, mask diff --git a/telefuser/pipelines/lingbot_video/vae.py b/telefuser/pipelines/lingbot_video/vae.py new file mode 100644 index 0000000..98fd7b8 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/vae.py @@ -0,0 +1,103 @@ +"""VAE geometry and normalization helpers for LingBot-Video. + +Latent normalization and condition-encoding behavior are adapted from the +Apache-2.0 licensed upstream LingBot-Video implementation. +""" + +from __future__ import annotations + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager + + +def latent_shape( + num_frames: int, height: int, width: int, *, temporal_factor: int = 4, spatial_factor: int = 8 +) -> tuple[int, int, int]: + """Return ``(latent_frames, latent_height, latent_width)`` for video input.""" + if num_frames < 1 or (num_frames - 1) % temporal_factor: + raise ValueError("num_frames must satisfy the VAE temporal contract") + if height % spatial_factor or width % spatial_factor: + raise ValueError("height and width must be divisible by VAE spatial factor") + return ((num_frames - 1) // temporal_factor + 1, height // spatial_factor, width // spatial_factor) + + +def normalize_latent(latent: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: + """Apply checkpoint-specific latent normalization with channel broadcasting.""" + if latent.ndim < 2: + raise ValueError("latent must include a channel dimension") + mean = mean.to(device=latent.device, dtype=latent.dtype).reshape(1, -1, *([1] * (latent.ndim - 2))) + std = std.to(device=latent.device, dtype=latent.dtype).reshape(1, -1, *([1] * (latent.ndim - 2))) + if mean.shape[1] != latent.shape[1] or torch.any(std == 0): + raise ValueError("latent normalization statistics do not match channels") + return (latent - mean) * std.reciprocal() + + +def denormalize_latent(latent: torch.Tensor, mean: torch.Tensor, std: torch.Tensor) -> torch.Tensor: + """Undo checkpoint-specific latent normalization.""" + if latent.ndim < 2: + raise ValueError("latent must include a channel dimension") + mean = mean.to(device=latent.device, dtype=latent.dtype).reshape(1, -1, *([1] * (latent.ndim - 2))) + std = std.to(device=latent.device, dtype=latent.dtype).reshape(1, -1, *([1] * (latent.ndim - 2))) + if mean.shape[1] != latent.shape[1]: + raise ValueError("latent normalization statistics do not match channels") + # Preserve the upstream division-by-inverse operation order. Multiplying + # by ``std`` is mathematically equivalent but does not have identical FP32 + # rounding before a temporally coupled VAE decode. + return latent / std.reciprocal() + mean + + +def first_frame_condition_mask(latent_frames: int, *, device: torch.device | str | None = None) -> torch.Tensor: + """Return a temporal mask selecting only latent frame zero.""" + if latent_frames < 1: + raise ValueError("latent_frames must be positive") + mask = torch.zeros(latent_frames, dtype=torch.bool, device=device) + mask[0] = True + return mask + + +class LingBotVideoVAEEncodeStage(BaseStage): + """Encode RGB video tensors using official LingBot latent normalization.""" + + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.vae = module_manager.fetch_module("vae") + self.model_names = ["vae"] + + @with_model_offload(["vae"]) + @torch.no_grad() + def encode(self, pixel_values: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + """Encode ``[B,3,F,H,W]`` pixels in [0,1] to normalized diffusion latents.""" + if pixel_values.ndim != 5 or pixel_values.shape[1] != 3: + raise ValueError("pixel_values must have shape [B,3,F,H,W]") + autocast_enabled = self.device.type == "cuda" + with torch.autocast(device_type=self.device.type, dtype=torch.bfloat16, enabled=autocast_enabled): + encoded = self.vae.encode((pixel_values.to(self.device, torch.float32) - 0.5) / 0.5) + posterior = encoded.latent_dist + latent = posterior.sample(generator=generator) if generator is not None else posterior.sample() + mean = torch.tensor(self.vae.config.latents_mean, device=latent.device) + std = torch.tensor(self.vae.config.latents_std, device=latent.device) + return normalize_latent(latent.float(), mean, std) + + +class LingBotVideoVAEDecodeStage(BaseStage): + """Decode normalized LingBot diffusion latents to RGB video tensors.""" + + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.vae = module_manager.fetch_module("vae") + self.model_names = ["vae"] + + @with_model_offload(["vae"]) + @torch.no_grad() + def decode(self, latents: torch.Tensor) -> torch.Tensor: + """Decode normalized latents and map output from [-1,1] to [0,1].""" + mean = torch.tensor(self.vae.config.latents_mean, device=self.device) + std = torch.tensor(self.vae.config.latents_std, device=self.device) + raw_latent = denormalize_latent(latents.to(self.device, torch.float32), mean, std) + if raw_latent.ndim == 5: + raw_latent = raw_latent.contiguous(memory_format=torch.channels_last_3d) + decoded = self.vae.decode(raw_latent).sample + return decoded.float().add(1.0).div(2.0).clamp(0.0, 1.0) diff --git a/telefuser/platforms/cpu.py b/telefuser/platforms/cpu.py index 5ec38ae..30370d7 100644 --- a/telefuser/platforms/cpu.py +++ b/telefuser/platforms/cpu.py @@ -17,6 +17,10 @@ class CpuPlatform(BasePlatform): dist_backend: str = "gloo" full_dist_backend: str = "cpu:gloo" + @staticmethod + def empty_cache() -> None: + """CPU has no device allocator cache to release.""" + @staticmethod def default_device(): return torch.device("cpu") diff --git a/telefuser/worker/parallel_worker.py b/telefuser/worker/parallel_worker.py index de9dff5..17b01c3 100644 --- a/telefuser/worker/parallel_worker.py +++ b/telefuser/worker/parallel_worker.py @@ -69,11 +69,15 @@ def _worker_loop( try: device = stage.device if world_size > 1: + parallel_config = stage.model_runtime_config.parallel_config + # Match torchrun's per-rank default. Letting every spawned worker + # inherit the host-wide intra-op pool oversubscribes CPU launch + # threads and can leave accelerators idle between eager kernels. + torch.set_num_threads(parallel_config.worker_intra_op_threads) os.environ["RANK"] = str(rank) os.environ["WORLD_SIZE"] = str(world_size) os.environ["MASTER_ADDR"] = "localhost" os.environ["MASTER_PORT"] = str(master_port) - parallel_config = stage.model_runtime_config.parallel_config device_ids = parallel_config.device_ids device_id = rank if device_ids is not None: @@ -118,7 +122,8 @@ def _worker_loop( args = [args] y = getattr(stage, name)(*args, **kwargs) del kwargs, args - current_platform.empty_cache() + if getattr(stage, "empty_cache_after_call", True): + current_platform.empty_cache() # Always output results when world_size=1 if world_size == 1 or rank == 0: queue_out.put(y) diff --git a/tests/unit/core/test_config.py b/tests/unit/core/test_config.py index 206533c..ea34f80 100644 --- a/tests/unit/core/test_config.py +++ b/tests/unit/core/test_config.py @@ -66,6 +66,13 @@ def test_default_world_size(self): """Default world_size should be 1.""" config = ParallelConfig() assert config.world_size == 1 + assert config.worker_intra_op_threads == 1 + + def test_worker_intra_op_threads_must_be_positive(self): + """Distributed worker CPU thread count must be positive.""" + config = ParallelConfig(worker_intra_op_threads=0) + with pytest.raises(ValueError, match="worker_intra_op_threads must be positive"): + config.validate() @pytest.mark.parametrize( "device_ids,degrees,expected_world_size", diff --git a/tests/unit/core/test_config_serializer.py b/tests/unit/core/test_config_serializer.py index d1e1326..a9825bc 100644 --- a/tests/unit/core/test_config_serializer.py +++ b/tests/unit/core/test_config_serializer.py @@ -46,6 +46,7 @@ def test_serialize_config_parallel_config(): assert result["dp_degree"] == 2 assert result["sp_ulysses_degree"] == 1 assert result["enable_fsdp"] is False + assert result["worker_intra_op_threads"] == 1 def test_serialize_config_offload_config(): diff --git a/tests/unit/models/test_lingbot_video_dit.py b/tests/unit/models/test_lingbot_video_dit.py new file mode 100644 index 0000000..255edfb --- /dev/null +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -0,0 +1,120 @@ +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from telefuser.core.config import AttentionConfig, AttnImplType +from telefuser.models.lingbot_video_dit import ( + LingBotVideoAttention, + LingBotVideoPatchEmbed, + LingBotVideoTransformer3DModel, + apply_lingbot_video_complex_rope, +) +from telefuser.pipelines.lingbot_video.data import LingBotVideoModelConfig + + +def test_patchify_matches_official_feature_order_and_round_trips() -> None: + config = LingBotVideoModelConfig(in_channels=2, hidden_size=8) + patch_embed = LingBotVideoPatchEmbed(config) + latent = torch.arange(8).reshape(1, 2, 1, 2, 2) + + tokens = patch_embed.patchify(latent) + + assert tokens.tolist() == [[[0, 4, 1, 5, 2, 6, 3, 7]]] + assert torch.equal(patch_embed.unpatchify(tokens, frames=1, height=2, width=2), latent) + + +def test_complex_rope_identity_table_preserves_values() -> None: + states = torch.randn(1, 4, 2, 8) + identity = torch.ones(4, 4, dtype=torch.complex64) + + assert torch.equal(apply_lingbot_video_complex_rope(states, identity), states) + + +def test_attention_dispatcher_sdpa_preserves_native_attention_result() -> None: + torch.manual_seed(7) + module = LingBotVideoAttention(hidden_size=8, num_heads=2, norm_eps=1e-6, qkv_bias=True, out_bias=True) + hidden_states = torch.randn(1, 3, 8) + rotary = torch.ones(3, 2, dtype=torch.complex64) + module.set_attention_config(AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA)) + + actual = module(hidden_states, rotary) + query = apply_lingbot_video_complex_rope(module.norm_q(module.to_q(hidden_states).view(1, 3, 2, 4)), rotary) + key = apply_lingbot_video_complex_rope(module.norm_k(module.to_k(hidden_states).view(1, 3, 2, 4)), rotary) + value = module.to_v(hidden_states).view(1, 3, 2, 4) + expected = module.to_out( + F.scaled_dot_product_attention(query.transpose(1, 2), key.transpose(1, 2), value.transpose(1, 2)) + .transpose(1, 2) + .reshape(1, 3, 8) + ) + + assert torch.equal(actual, expected) + + +def test_transformer_rejects_non_patch_aligned_latent_geometry() -> None: + model = LingBotVideoTransformer3DModel( + in_channels=16, + hidden_size=16, + num_attention_heads=2, + depth=1, + intermediate_size=32, + text_dim=12, + freq_dim=8, + axes_dims=(2, 2, 4), + ) + + with pytest.raises(ValueError, match="divisible by the checkpoint patch size"): + model(torch.randn(1, 16, 1, 3, 4), torch.tensor([1]), torch.randn(1, 3, 12)) + + +def test_transformer_packed_batch_matches_independent_samples_with_different_text_lengths() -> None: + torch.manual_seed(11) + model = LingBotVideoTransformer3DModel( + in_channels=2, + out_channels=2, + hidden_size=16, + num_attention_heads=2, + depth=1, + intermediate_size=32, + text_dim=12, + freq_dim=8, + axes_dims=(2, 2, 4), + ).eval() + latents = torch.randn(2, 2, 1, 2, 2) + timesteps = torch.tensor([700.0, 300.0]) + text = torch.randn(2, 3, 12) + mask = torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.bool) + + actual = model(latents, timesteps, text, mask) + expected = torch.cat( + [ + model(latents[index : index + 1], timesteps[index : index + 1], text[index : index + 1, :length]) + for index, length in enumerate((3, 1)) + ] + ) + + torch.testing.assert_close(actual, expected, rtol=1e-5, atol=1e-5) + + +def test_transformer_rejects_packed_batch_with_non_sdpa_attention() -> None: + model = LingBotVideoTransformer3DModel( + in_channels=2, + out_channels=2, + hidden_size=16, + num_attention_heads=2, + depth=1, + intermediate_size=32, + text_dim=12, + freq_dim=8, + axes_dims=(2, 2, 4), + ).eval() + model.set_attention_config(AttentionConfig.dense_attention(AttnImplType.FLASH_ATTN_2)) + + with pytest.raises(ValueError, match="packed sequence attention currently requires TORCH_SDPA"): + model( + torch.randn(2, 2, 1, 2, 2), + torch.tensor([700.0, 300.0]), + torch.randn(2, 3, 12), + torch.tensor([[1, 1, 1], [1, 0, 0]], dtype=torch.bool), + ) diff --git a/tests/unit/models/test_lingbot_video_moe.py b/tests/unit/models/test_lingbot_video_moe.py new file mode 100644 index 0000000..3ed6d8f --- /dev/null +++ b/tests/unit/models/test_lingbot_video_moe.py @@ -0,0 +1,182 @@ +from __future__ import annotations + +import pytest +import torch +import torch.nn.functional as F + +from telefuser.models.lingbot_video_moe import LingBotVideoGroupedExperts, LingBotVideoMoeTransformer3DModel + + +def test_moe_transformer_uses_official_router_and_expert_key_layout() -> None: + model = LingBotVideoMoeTransformer3DModel( + hidden_size=16, + num_attention_heads=2, + depth=1, + intermediate_size=32, + text_dim=12, + freq_dim=8, + axes_dims=(2, 2, 4), + num_experts=4, + num_experts_per_tok=2, + moe_intermediate_size=8, + n_group=2, + topk_group=1, + routed_scaling_factor=2.5, + n_shared_experts=1, + ) + + output = model(torch.randn(1, 16, 1, 4, 4), torch.tensor([1]), torch.randn(1, 3, 12)) + + assert output.shape == (1, 16, 1, 4, 4) + assert "blocks.0.ffn.router.weight" in model.state_dict() + assert "blocks.0.ffn.experts.w1" in model.state_dict() + assert "blocks.0.ffn.shared_experts.up_proj.weight" in model.state_dict() + + +def test_grouped_experts_restores_route_outputs_with_fp32_weighted_sum() -> None: + torch.manual_seed(7) + experts = LingBotVideoGroupedExperts(num_experts=2, hidden_size=4, intermediate_size=6).bfloat16() + with torch.no_grad(): + experts.w1.uniform_(-0.1, 0.1) + experts.w2.uniform_(-0.1, 0.1) + experts.w3.uniform_(-0.1, 0.1) + tokens = torch.randn(3, 4, dtype=torch.bfloat16) + indices = torch.tensor([[0, 1], [1, 0], [0, 1]]) + weights = torch.tensor([[0.625, 0.375], [0.25, 0.75], [0.5, 0.5]], dtype=torch.bfloat16) + + actual = experts(tokens, indices, weights) + routed = torch.zeros(tokens.shape[0] * indices.shape[1], tokens.shape[1], dtype=tokens.dtype) + for expert_index in range(experts.w1.shape[0]): + token_indices, topk_indices = torch.where(indices == expert_index) + selected = tokens[token_indices] + activation = F.silu(F.linear(selected, experts.w1[expert_index])) * F.linear(selected, experts.w3[expert_index]) + routed[token_indices * indices.shape[1] + topk_indices] = F.linear(activation, experts.w2[expert_index]) + expected = (routed.reshape(tokens.shape[0], indices.shape[1], -1).float() * weights.float().unsqueeze(-1)).sum(1) + + assert torch.equal(actual, expected.to(torch.bfloat16)) + + +def test_grouped_experts_skips_zero_weight_routes_without_changing_output() -> None: + torch.manual_seed(11) + experts = LingBotVideoGroupedExperts(num_experts=3, hidden_size=4, intermediate_size=6).bfloat16() + with torch.no_grad(): + experts.w1.uniform_(-0.1, 0.1) + experts.w2.uniform_(-0.1, 0.1) + experts.w3.uniform_(-0.1, 0.1) + tokens = torch.randn(3, 4, dtype=torch.bfloat16) + indices = torch.tensor([[0, 1], [2, 0], [1, 2]]) + weights = torch.tensor([[0.625, 0.375], [0.0, 0.0], [0.5, 0.5]], dtype=torch.bfloat16) + + actual = experts(tokens, indices, weights) + routed = torch.zeros(tokens.shape[0] * indices.shape[1], tokens.shape[1], dtype=tokens.dtype) + for expert_index in range(experts.w1.shape[0]): + token_indices, topk_indices = torch.where(indices == expert_index) + selected = tokens[token_indices] + activation = F.silu(F.linear(selected, experts.w1[expert_index])) * F.linear(selected, experts.w3[expert_index]) + routed[token_indices * indices.shape[1] + topk_indices] = F.linear(activation, experts.w2[expert_index]) + expected = (routed.reshape(tokens.shape[0], indices.shape[1], -1).float() * weights.float().unsqueeze(-1)).sum(1) + + assert torch.equal(actual, expected.to(torch.bfloat16)) + + +def test_grouped_expert_backends_are_exactly_equivalent() -> None: + torch.manual_seed(13) + experts = LingBotVideoGroupedExperts(num_experts=4, hidden_size=4, intermediate_size=6).bfloat16() + with torch.no_grad(): + experts.w1.uniform_(-0.1, 0.1) + experts.w2.uniform_(-0.1, 0.1) + experts.w3.uniform_(-0.1, 0.1) + tokens = torch.randn(5, 4, dtype=torch.bfloat16) + indices = torch.tensor([[0, 1], [2, 3], [3, 0], [1, 2], [0, 3]]) + weights = torch.tensor([[0.625, 0.375], [0.5, 0.5], [0.25, 0.75], [0.75, 0.25], [0.5, 0.5]], dtype=torch.bfloat16) + + sorted_output = experts(tokens, indices, weights) + experts.set_execution_backend("where") + where_output = experts(tokens, indices, weights) + + assert torch.equal(sorted_output, where_output) + + +def test_grouped_mm_backend_requires_cuda() -> None: + experts = LingBotVideoGroupedExperts(num_experts=2, hidden_size=4, intermediate_size=8).bfloat16() + experts.set_execution_backend("grouped_mm") + + with pytest.raises(RuntimeError, match="requires CUDA"): + experts( + torch.randn(2, 4, dtype=torch.bfloat16), + torch.tensor([[0, 1], [1, 0]]), + torch.full((2, 2), 0.5, dtype=torch.bfloat16), + ) + + +def test_fp8_expert_quantization_preserves_checkpoint_keys() -> None: + experts = LingBotVideoGroupedExperts(num_experts=2, hidden_size=16, intermediate_size=32).bfloat16() + expected_keys = set(experts.state_dict()) + + experts.quantize_fp8_() + + assert experts.w1.dtype == torch.float8_e4m3fn + assert experts.w2.dtype == torch.float8_e4m3fn + assert experts.w3.dtype == torch.float8_e4m3fn + assert experts.w1_scale.shape == (2, 32) + assert experts.w2_scale.shape == (2, 16) + assert experts.w3_scale.shape == (2, 32) + assert set(experts.state_dict()) == expected_keys + + +def test_fp8_backend_requires_cuda() -> None: + experts = LingBotVideoGroupedExperts(num_experts=2, hidden_size=16, intermediate_size=32).bfloat16() + experts.quantize_fp8_() + experts.set_execution_backend("fp8") + + with pytest.raises(RuntimeError, match="requires CUDA"): + experts( + torch.randn(2, 16, dtype=torch.bfloat16), + torch.tensor([[0, 1], [1, 0]]), + torch.full((2, 2), 0.5, dtype=torch.bfloat16), + ) + + +@pytest.mark.gpu +def test_grouped_mm_backend_matches_sorted_backend_on_cuda() -> None: + if not torch.cuda.is_available() or not hasattr(torch, "_grouped_mm"): + pytest.skip("native CUDA grouped_mm is unavailable") + torch.manual_seed(17) + experts = LingBotVideoGroupedExperts(num_experts=4, hidden_size=32, intermediate_size=24).cuda().bfloat16() + with torch.no_grad(): + experts.w1.uniform_(-0.1, 0.1) + experts.w2.uniform_(-0.1, 0.1) + experts.w3.uniform_(-0.1, 0.1) + tokens = torch.randn(32, 32, device="cuda", dtype=torch.bfloat16) + indices = torch.randint(0, 4, (32, 2), device="cuda") + weights = torch.rand(32, 2, device="cuda", dtype=torch.bfloat16) + weights /= weights.sum(dim=-1, keepdim=True) + + sorted_output = experts(tokens, indices, weights) + experts.set_execution_backend("grouped_mm") + grouped_output = experts(tokens, indices, weights) + + assert torch.equal(grouped_output, sorted_output) + + +@pytest.mark.gpu +def test_fp8_backend_matches_sorted_backend_on_cuda() -> None: + if not torch.cuda.is_available() or not hasattr(torch, "_scaled_mm"): + pytest.skip("native CUDA scaled_mm is unavailable") + torch.manual_seed(19) + experts = LingBotVideoGroupedExperts(num_experts=4, hidden_size=64, intermediate_size=128).cuda().bfloat16() + with torch.no_grad(): + experts.w1.normal_(std=0.02) + experts.w2.normal_(std=0.02) + experts.w3.normal_(std=0.02) + tokens = torch.randn(64, 64, device="cuda", dtype=torch.bfloat16) + indices = torch.randint(0, 4, (64, 2), device="cuda") + weights = torch.rand(64, 2, device="cuda", dtype=torch.bfloat16) + weights /= weights.sum(dim=-1, keepdim=True) + + sorted_output = experts(tokens, indices, weights) + experts.quantize_fp8_() + experts.set_execution_backend("fp8") + fp8_output = experts(tokens, indices, weights) + + torch.testing.assert_close(fp8_output, sorted_output, rtol=0.12, atol=2e-3) diff --git a/tests/unit/pipelines/lingbot_video/test_data.py b/tests/unit/pipelines/lingbot_video/test_data.py new file mode 100644 index 0000000..1f455c3 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_data.py @@ -0,0 +1,67 @@ +from __future__ import annotations + +from pathlib import Path + +import pytest +import torch +import torch.nn.functional as F + +from telefuser.pipelines.lingbot_video.data import ( + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_NEGATIVE_PROMPT_IMAGE, + LingBotVideoRequest, + default_negative_caption, + load_lingbot_video_model_config, + num_frames_from_duration, + parse_lingbot_video_prompt, + preprocess_ti2v_image, +) + + +def test_load_dense_transformer_config() -> None: + config_path = Path("/hhb-data/aigc/model_zoo/lingbot/lingbot-video-dense-1.3b/transformer") + if not config_path.is_dir(): + return + + config = load_lingbot_video_model_config(config_path) + + assert config.num_layers == 24 + assert config.hidden_size == 2048 + assert config.patch_size == (1, 2, 2) + + +def test_preprocess_ti2v_image_preserves_official_uint8_interpolation() -> None: + image = torch.tensor([[[[0.0, 10.0], [20.0, 30.0]]]]).repeat(1, 3, 1, 1) + + result = preprocess_ti2v_image(image, height=3, width=3) + expected = F.interpolate(image.to(torch.uint8), size=(3, 3), mode="bilinear", align_corners=False).float() / 255.0 + + assert torch.equal(result, expected.unsqueeze(2)) + + +def test_rewriter_prompt_envelope_matches_upstream_caption_and_duration_contract() -> None: + caption, duration = parse_lingbot_video_prompt({"caption": {"scene": "a test scene"}, "duration": 5, "fps": 24}) + + assert caption == '{"scene":"a test scene"}' + assert duration == 5.0 + assert num_frames_from_duration(duration) == 121 + + +def test_raw_prompt_sample_excludes_runtime_metadata() -> None: + caption, duration = parse_lingbot_video_prompt({"scene": "test", "duration": 2, "fps": 24}) + + assert caption == '{"scene":"test"}' + assert duration == 2.0 + assert num_frames_from_duration(duration) == 49 + + +def test_request_requires_dit_patch_aligned_spatial_dimensions() -> None: + with pytest.raises(ValueError, match="divisible by 16"): + LingBotVideoRequest(caption="{}", height=480, width=856, num_frames=1) + + +def test_default_negative_captions_match_mode_contract() -> None: + assert default_negative_caption(1) == DEFAULT_NEGATIVE_PROMPT_IMAGE + assert default_negative_caption(5) == DEFAULT_NEGATIVE_PROMPT + assert "temporal_and_motion_stability" not in DEFAULT_NEGATIVE_PROMPT_IMAGE + assert "temporal_and_motion_stability" in DEFAULT_NEGATIVE_PROMPT diff --git a/tests/unit/pipelines/lingbot_video/test_pipeline.py b/tests/unit/pipelines/lingbot_video/test_pipeline.py new file mode 100644 index 0000000..19e7325 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -0,0 +1,429 @@ +from __future__ import annotations + +import torch +import torch.nn as nn + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.pipelines.lingbot_video import denoising as denoising_module +from telefuser.pipelines.lingbot_video import pipeline as pipeline_module +from telefuser.pipelines.lingbot_video.data import DEFAULT_NEGATIVE_PROMPT, LingBotVideoRequest +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, transformer_timestep +from telefuser.pipelines.lingbot_video.pipeline import LingBotVideoPipeline, LingBotVideoPipelineConfig + + +class _TextStage: + def prepare_ti2v_vlm_image(self, pixel_values: torch.Tensor) -> object: + return object() + + def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tensor, torch.Tensor]: + del caption, images + return torch.zeros(1, 1, 4), torch.ones(1, 1, dtype=torch.long) + + +class _DenoisingStage: + device = torch.device("cpu") + + def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object) -> torch.Tensor: + return torch.ones_like(latents) + + +class _LoopDenoisingStage(_DenoisingStage): + def __init__(self) -> None: + self.calls = 0 + + def denoise(self, latent: torch.Tensor, *_: object) -> torch.Tensor: + self.calls += 1 + return torch.full_like(latent, 7.0) + + +class _PipelineTransformer(nn.Module): + def set_attention_config(self, attention_config: object) -> None: + self.attention_config = attention_config + + +class _PipelineModuleManager: + def __init__(self) -> None: + self.modules = { + "transformer": _PipelineTransformer(), + "text_encoder": nn.Linear(1, 1), + "processor": object(), + "vae": nn.Linear(1, 1), + "scheduler": _Scheduler(), + } + + def get_model_info(self) -> list[dict[str, object]]: + return [] + + def fetch_module(self, name: str) -> object: + return self.modules[name] + + +def _transformer_manager(transformer: nn.Module) -> _PipelineModuleManager: + manager = _PipelineModuleManager() + manager.modules["transformer"] = transformer + return manager + + +def _initialized_pipeline() -> LingBotVideoPipeline: + pipeline = LingBotVideoPipeline(device="cpu") + pipeline.init(_PipelineModuleManager(), LingBotVideoPipelineConfig(num_inference_steps=2)) + return pipeline + + +class _Scheduler: + timesteps = torch.tensor([2, 1]) + + def set_timesteps(self, *_: object, **__: object) -> None: + return None + + def step(self, prediction: torch.Tensor, timestep: torch.Tensor, latent: torch.Tensor) -> torch.Tensor: + del timestep + return latent - prediction + + +def test_initialized_pipeline_runs_t2v_sampling() -> None: + pipeline = _initialized_pipeline() + pipeline.text_stage = _TextStage() + pipeline.denoising_stage = _DenoisingStage() + pipeline.scheduler = _Scheduler() + + generation = pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert generation.output.shape == (1, 16, 2, 2, 2) + assert not generation.prompt_conditions.has_visual_condition + assert torch.equal(generation.prompt_conditions.positive_prompt_embeds, torch.zeros(1, 1, 4)) + assert pipeline(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False).shape == ( + 1, + 16, + 2, + 2, + 2, + ) + + +def test_parallel_pipeline_runs_the_complete_denoise_loop_in_one_stage_call() -> None: + pipeline = _initialized_pipeline() + stage = _LoopDenoisingStage() + pipeline.config.enable_denoising_parallel = True + pipeline.text_stage = _TextStage() + pipeline.denoising_stage = stage + pipeline.scheduler = _Scheduler() + + generation = pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert stage.calls == 1 + assert torch.equal(generation.output, torch.full_like(generation.output, 7.0)) + + +class _DtypeProbeTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.patch_embedder = nn.Linear(1, 1, bias=False) + + def forward( + self, latents: torch.Tensor, timestep: torch.Tensor, embeds: torch.Tensor, mask: torch.Tensor | None + ) -> torch.Tensor: + del timestep, mask + return (latents + embeds.mean()).to(torch.bfloat16) + + +def test_denoising_stage_returns_fp32_scheduler_prediction() -> None: + stage = LingBotVideoDenoisingStage( + "denoising", + _transformer_manager(_DtypeProbeTransformer()), + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + ) + latents = torch.zeros(1, 1, 1, 1, 1) + positive = torch.full((1, 1, 1), 3.0) + negative = torch.full((1, 1, 1), 1.0) + + prediction = stage.predict_noise_with_cfg(latents, torch.tensor([1]), positive, negative, guidance_scale=2.0) + + assert prediction.dtype is torch.float32 + assert torch.equal(prediction, torch.full_like(latents, 5.0)) + + +def test_denoising_stage_disables_cfg_when_guidance_is_below_one() -> None: + transformer = _BatchCfgProbeTransformer() + stage = LingBotVideoDenoisingStage( + "denoising", + _transformer_manager(transformer), + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + ) + + prediction = stage.predict_noise_with_cfg( + torch.zeros(1, 1, 1, 1, 1), + torch.tensor([1]), + torch.full((1, 1, 1), 3.0), + None, + guidance_scale=0.5, + ) + + assert torch.equal(prediction, torch.full_like(prediction, 3.0)) + assert transformer.calls == 1 + + +class _BatchCfgProbeTransformer(nn.Module): + def __init__(self) -> None: + super().__init__() + self.patch_embedder = nn.Linear(1, 1, bias=False) + self.calls = 0 + + def forward( + self, latents: torch.Tensor, timestep: torch.Tensor, embeds: torch.Tensor, mask: torch.Tensor | None + ) -> torch.Tensor: + del timestep + self.calls += 1 + if mask is None: + offset = embeds.mean(dim=(1, 2)) + else: + weights = mask.to(embeds.dtype).unsqueeze(-1) + offset = (embeds * weights).sum(dim=(1, 2)) / weights.sum(dim=(1, 2)) + offset = offset.reshape(-1, 1, 1, 1, 1) + return latents + offset + + +def test_denoising_stage_batched_cfg_matches_two_source_order_forwards() -> None: + latents = torch.zeros(1, 1, 1, 1, 1) + positive = torch.full((1, 2, 1), 3.0) + negative = torch.full((1, 2, 1), 1.0) + runtime_config = ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32) + serial_transformer = _BatchCfgProbeTransformer() + batched_transformer = _BatchCfgProbeTransformer() + serial = LingBotVideoDenoisingStage("serial", _transformer_manager(serial_transformer), runtime_config) + batched = LingBotVideoDenoisingStage( + "batched", _transformer_manager(batched_transformer), runtime_config, batch_cfg=True + ) + + expected = serial.predict_noise_with_cfg(latents, torch.tensor([1]), positive, negative, guidance_scale=2.0) + actual = batched.predict_noise_with_cfg(latents, torch.tensor([1]), positive, negative, guidance_scale=2.0) + + assert torch.equal(actual, expected) + assert serial_transformer.calls == 2 + assert batched_transformer.calls == 1 + + +def test_denoising_stage_cfg_parallel_gathers_positive_then_negative(monkeypatch: object) -> None: + latents = torch.zeros(1, 1, 1, 1, 1) + transformer = _BatchCfgProbeTransformer() + transformer.device_mesh = object() + stage = LingBotVideoDenoisingStage( + "cfg-parallel", + _transformer_manager(transformer), + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + ) + + monkeypatch.setattr(denoising_module, "get_cfg_world_size", lambda _: 2) + monkeypatch.setattr(denoising_module, "get_cfg_rank", lambda _: 0) + monkeypatch.setattr( + denoising_module, + "cfg_parallel_unshard", + lambda _, predictions: [torch.cat((predictions[0], predictions[0] - 2.0))], + ) + + prediction = stage.predict_noise_with_cfg( + latents, + torch.tensor([1]), + torch.full((1, 2, 1), 3.0), + torch.full((1, 2, 1), 1.0), + guidance_scale=2.0, + ) + + assert torch.equal(prediction, torch.full_like(prediction, 5.0)) + assert transformer.calls == 1 + + +def test_denoising_stage_batched_cfg_pads_different_prompt_lengths() -> None: + latents = torch.zeros(1, 1, 1, 1, 1) + positive = torch.full((1, 3, 1), 3.0) + negative = torch.full((1, 1, 1), 1.0) + runtime_config = ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32) + serial_transformer = _BatchCfgProbeTransformer() + batched_transformer = _BatchCfgProbeTransformer() + serial = LingBotVideoDenoisingStage("serial", _transformer_manager(serial_transformer), runtime_config) + batched = LingBotVideoDenoisingStage( + "batched", _transformer_manager(batched_transformer), runtime_config, batch_cfg=True + ) + + expected = serial.predict_noise_with_cfg(latents, torch.tensor([1]), positive, negative, guidance_scale=2.0) + actual = batched.predict_noise_with_cfg(latents, torch.tensor([1]), positive, negative, guidance_scale=2.0) + + assert torch.equal(actual, expected) + assert serial_transformer.calls == 2 + assert batched_transformer.calls == 1 + + +def test_denoising_stage_can_own_the_complete_scheduler_loop() -> None: + transformer = _BatchCfgProbeTransformer() + stage = LingBotVideoDenoisingStage( + "batched", + _transformer_manager(transformer), + ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + batch_cfg=True, + ) + latent = stage.denoise( + torch.zeros(1, 1, 1, 1, 1), + torch.full((1, 1, 1), 3.0), + torch.full((1, 1, 1), 1.0), + torch.ones(1, 1, dtype=torch.bool), + torch.ones(1, 1, dtype=torch.bool), + 2.0, + 2, + 3.0, + ) + + assert transformer.calls == 2 + assert torch.equal(latent, torch.full_like(latent, -5.0)) + + +class _RecordingTextStage(_TextStage): + def __init__(self) -> None: + self.images: list[object | None] = [] + + def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tensor, torch.Tensor]: + self.images.append(images) + return super().encode(caption, images) + + +def test_pipeline_uses_source_negative_caption_when_not_overridden() -> None: + class _CapturingTextStage(_TextStage): + def __init__(self) -> None: + self.captions: list[str] = [] + + def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tensor, torch.Tensor]: + self.captions.append(caption) + return super().encode(caption, images) + + text_stage = _CapturingTextStage() + pipeline = _initialized_pipeline() + pipeline.text_stage = text_stage + pipeline.denoising_stage = _DenoisingStage() + pipeline.scheduler = _Scheduler() + + pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert text_stage.captions == ["{}", DEFAULT_NEGATIVE_PROMPT] + + +def test_pipeline_skips_negative_encoding_when_cfg_is_disabled() -> None: + class _CapturingTextStage(_TextStage): + def __init__(self) -> None: + self.captions: list[str] = [] + + def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tensor, torch.Tensor]: + self.captions.append(caption) + return super().encode(caption, images) + + text_stage = _CapturingTextStage() + pipeline = _initialized_pipeline() + pipeline.config.guidance_scale = 0.5 + pipeline.text_stage = text_stage + pipeline.denoising_stage = _DenoisingStage() + pipeline.scheduler = _Scheduler() + + generation = pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert text_stage.captions == ["{}"] + assert generation.prompt_conditions.negative_prompt_embeds is None + assert generation.prompt_conditions.negative_attention_mask is None + + +class _VAEEncodeStage: + def __init__(self) -> None: + self.pixel_values: torch.Tensor | None = None + + def encode(self, pixel_values: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + del generator + self.pixel_values = pixel_values + return torch.zeros(1, 16, 1, 2, 2) + + +def test_pipeline_ti2v_reuses_one_visual_condition_for_both_cfg_branches() -> None: + text_stage = _RecordingTextStage() + vae_encode_stage = _VAEEncodeStage() + pipeline = _initialized_pipeline() + pipeline.text_stage = text_stage + pipeline.denoising_stage = _DenoisingStage() + pipeline.scheduler = _Scheduler() + pipeline.vae_encode_stage = vae_encode_stage + + generation = pipeline.generate( + LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5, image=torch.full((1, 3, 8, 8), 128.0)), + decode=False, + ) + + assert generation.prompt_conditions.has_visual_condition + assert len(text_stage.images) == 2 + assert all(images is not None for images in text_stage.images) + assert vae_encode_stage.pixel_values is not None + assert vae_encode_stage.pixel_values.shape == (1, 3, 1, 16, 16) + assert torch.equal(generation.output[:, :, :1], torch.zeros_like(generation.output[:, :, :1])) + + +def test_transformer_timestep_preserves_upstream_bfloat16_rounding() -> None: + timestep = torch.tensor([999], dtype=torch.int64) + + actual = transformer_timestep(timestep, torch.bfloat16) + + expected = ((timestep.float() / 1000).bfloat16() * 1000).float() + assert torch.equal(actual, expected) + + +class _OffloadStage: + def __init__(self) -> None: + self.calls = 0 + self.onload_models_flag = True + + def offload_models(self) -> None: + self.calls += 1 + + +def test_pipeline_releases_stages_before_a_separate_refiner_is_loaded() -> None: + text_stage = _OffloadStage() + denoising_stage = _OffloadStage() + vae_encode_stage = _OffloadStage() + vae_decode_stage = _OffloadStage() + pipeline = _initialized_pipeline() + pipeline.text_stage = text_stage + pipeline.denoising_stage = denoising_stage + pipeline.scheduler = _Scheduler() + pipeline.vae_encode_stage = vae_encode_stage + pipeline.vae_decode_stage = vae_decode_stage + + pipeline.release_gpu_resources() + + for stage in (text_stage, denoising_stage, vae_encode_stage, vae_decode_stage): + assert stage.calls == 1 + assert not stage.onload_models_flag + + +def test_pipeline_recreates_released_parallel_worker_for_next_request(monkeypatch) -> None: + class _FakeParallelWorker: + created: list[object] = [] + + def __init__(self, stage: object) -> None: + self.created.append(stage) + self.closed = False + self.device = torch.device("cpu") + + def close(self) -> None: + self.closed = True + + def denoise(self, latent: torch.Tensor, *_: object) -> torch.Tensor: + return torch.full_like(latent, 7.0) + + monkeypatch.setattr(pipeline_module, "ParallelWorker", _FakeParallelWorker) + pipeline = _initialized_pipeline() + template = object() + pipeline.config.enable_denoising_parallel = True + pipeline._parallel_denoising_stage_template = template + pipeline.text_stage = _TextStage() + pipeline.denoising_stage = _FakeParallelWorker(template) + pipeline.vae_encode_stage = None + pipeline.vae_decode_stage = None + + pipeline.release_gpu_resources() + generation = pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert len(_FakeParallelWorker.created) == 2 + assert torch.equal(generation.output, torch.full_like(generation.output, 7.0)) diff --git a/tests/unit/pipelines/lingbot_video/test_refiner.py b/tests/unit/pipelines/lingbot_video/test_refiner.py new file mode 100644 index 0000000..0f529b4 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_refiner.py @@ -0,0 +1,167 @@ +from __future__ import annotations + +import importlib +import sys +import types +from pathlib import Path + +import numpy as np +import pytest +import torch + +from telefuser.pipelines.lingbot_video.refiner import ( + compute_refiner_sigmas, + compute_training_aligned_indices, + compute_training_frame_budget, + load_refiner_first_frame, + load_refiner_video_file, + prepare_refiner_latent, + prepare_refiner_video, +) +from tools.validation.capture_lingbot_video_reference import _upstream_import_path + +UPSTREAM_ROOT = Path("work_dirs/lingbot-video-master") + + +def _require_upstream_checkout() -> Path: + if not (UPSTREAM_ROOT / "lingbot_video" / "__init__.py").is_file(): + pytest.skip("LingBot-Video upstream checkout is not available") + return UPSTREAM_ROOT + + +def test_refiner_schedule_starts_at_threshold_and_descends() -> None: + sigmas = compute_refiner_sigmas( + sigma_max=1.0, + sigma_min=0.0, + num_inference_steps=8, + shift=3.0, + t_thresh=0.85, + tail_steps=2, + ) + + assert sigmas is not None + assert sigmas[0] == 0.85 + assert all(left > right for left, right in zip(sigmas, sigmas[1:])) + + +def test_refiner_latent_mixing() -> None: + mixed = prepare_refiner_latent(torch.ones(1, 1, 1), torch.zeros(1, 1, 1), 0.25) + + assert torch.allclose(mixed, torch.full_like(mixed, 0.75)) + + +def test_training_aligned_refiner_handoff_matches_upstream_sampling_contract() -> None: + frames, vae_fps, temporal_latents = compute_training_frame_budget(10, 48.0, sample_fps=24, vae_tc=4) + + assert (frames, vae_fps, temporal_latents) == (5, 24.0, 2) + assert torch.equal(compute_training_aligned_indices(10, frames), torch.tensor([0, 2, 4, 6, 9])) + + video = torch.linspace(0.0, 1.0, 10).reshape(1, 1, 10, 1, 1).repeat(1, 3, 1, 2, 2) + prepared, metadata = prepare_refiner_video(video, source_fps=48.0, height=2, width=2) + + assert prepared.shape == (1, 3, 5, 2, 2) + assert torch.equal(prepared[:, :, :, 0, 0], video[:, :, [0, 2, 4, 6, 9], 0, 0]) + assert metadata == { + "src_fps": 48.0, + "sample_frame": 5, + "sample_frame_uncapped": 5, + "max_frames": None, + "truncated_by_max_frames": False, + "vae_fps": 24.0, + "t_vae": 2, + "num_source_frames": 10, + "align_to_training": True, + } + + +def test_pyav_mp4_refiner_handoff_uses_training_aligned_sampling(tmp_path: Path) -> None: + av = pytest.importorskip("av") + path = tmp_path / "base.mp4" + container = av.open(str(path), mode="w") + stream = container.add_stream("mpeg4", rate=48) + stream.width = 16 + stream.height = 16 + stream.pix_fmt = "yuv420p" + try: + for value in range(10): + image = np.full((16, 16, 3), value * 20, dtype=np.uint8) + frame = av.VideoFrame.from_ndarray(image, format="rgb24") + for packet in stream.encode(frame): + container.mux(packet) + for packet in stream.encode(): + container.mux(packet) + finally: + container.close() + + video, metadata = load_refiner_video_file(path, height=16, width=16, sample_fps=24, vae_tc=4) + + assert video.shape == (1, 3, 5, 16, 16) + assert metadata["sample_frame"] == 5 + assert metadata["t_vae"] == 2 + assert metadata["src_fps"] == 48.0 + + +def test_mp4_refiner_handoff_matches_upstream_with_pyav_decord_adapter( + tmp_path: Path, monkeypatch: pytest.MonkeyPatch +) -> None: + av = pytest.importorskip("av") + export_to_video = pytest.importorskip("diffusers.utils").export_to_video + upstream_root = _require_upstream_checkout() + path = tmp_path / "base.mp4" + + frames = [np.full((16, 16, 3), value * 20, dtype=np.uint8) for value in range(10)] + export_to_video(frames, str(path), fps=48) + + class _Batch: + def __init__(self, values: np.ndarray) -> None: + self.values = values + + def asnumpy(self) -> np.ndarray: + return self.values + + class _VideoReader: + def __init__(self, source: str, ctx: object) -> None: + del ctx + decoded = av.open(source) + try: + video_stream = next(iter(decoded.streams.video)) + self.fps = float(video_stream.average_rate) + self.frames = [frame.to_ndarray(format="rgb24") for frame in decoded.decode(video_stream)] + finally: + decoded.close() + + def __len__(self) -> int: + return len(self.frames) + + def get_avg_fps(self) -> float: + return self.fps + + def get_batch(self, indices: np.ndarray) -> _Batch: + return _Batch(np.stack([self.frames[int(index)] for index in indices])) + + decord = types.ModuleType("decord") + decord.VideoReader = _VideoReader + decord.cpu = lambda _: object() + monkeypatch.setitem(sys.modules, "decord", decord) + with _upstream_import_path(upstream_root): + upstream_utils = importlib.import_module("lingbot_video.utils") + reference, reference_metadata = upstream_utils.load_refiner_video_tensor(path, 16, 16, sample_fps=24, vae_tc=4) + candidate, candidate_metadata = load_refiner_video_file(path, height=16, width=16, sample_fps=24, vae_tc=4) + + assert reference_metadata == candidate_metadata + assert torch.equal(reference, candidate) + + +def test_ti2v_refiner_first_frame_matches_upstream_geometry(tmp_path: Path) -> None: + from PIL import Image + + upstream_root = _require_upstream_checkout() + image = np.arange(9 * 15 * 3, dtype=np.uint8).reshape(9, 15, 3) + path = tmp_path / "first.png" + Image.fromarray(image).save(path) + with _upstream_import_path(upstream_root): + upstream_utils = importlib.import_module("lingbot_video.utils") + reference = upstream_utils.load_first_frame_condition_tensor(path, 8, 12, 6, 10) + candidate = load_refiner_first_frame(path, target_height=8, target_width=12, geometry_height=6, geometry_width=10) + + assert torch.equal(reference, candidate) diff --git a/tests/unit/pipelines/lingbot_video/test_refiner_stage.py b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py new file mode 100644 index 0000000..fa06c52 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py @@ -0,0 +1,211 @@ +from __future__ import annotations + +import torch + +from telefuser.pipelines.lingbot_video.refiner import LingBotVideoRefinerStage + + +class _EncodeStage: + def encode(self, values: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + del generator + return values[:, :1] + + +class _DecodeStage: + def decode(self, values: torch.Tensor) -> torch.Tensor: + return values + + +class _DenoiseStage: + def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object) -> torch.Tensor: + return torch.ones_like(latents) + + +class _DeferredDenoiseStage(_DenoiseStage): + def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object): + return lambda: torch.ones_like(latents) + + +class _LifecycleEncodeStage(_EncodeStage): + def __init__(self) -> None: + self.offloaded = False + + def offload_models(self) -> None: + self.offloaded = True + + +class _LifecycleDecodeStage(_DecodeStage): + def __init__(self) -> None: + self.offloaded = False + + def offload_models(self) -> None: + self.offloaded = True + + +class _LifecycleDenoiseStage(_DenoiseStage): + def __init__(self) -> None: + self.closed = False + + def close(self) -> None: + self.closed = True + + +class _Scheduler: + sigma_max = 1.0 + sigma_min = 0.0 + timesteps = torch.tensor([1]) + + def set_timesteps(self, *_: object, **__: object) -> None: + return None + + def step(self, prediction: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor) -> torch.Tensor: + del timestep + return sample - prediction + + +def test_refiner_reinjects_clean_first_frame_after_scheduler_step() -> None: + stage = LingBotVideoRefinerStage( + denoising_stage=_DenoiseStage(), + vae_encode_stage=_EncodeStage(), + vae_decode_stage=_DecodeStage(), + scheduler=_Scheduler(), + ) + video = torch.zeros(1, 3, 2, 1, 1) + first_frame = torch.ones(1, 3, 1, 1) + + result = stage.refine( + video, + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + None, + None, + num_inference_steps=1, + guidance_scale=1.0, + shift=1.0, + t_thresh=0.5, + tail_steps=0, + clean_first_frame=first_frame, + generator=torch.Generator().manual_seed(0), + ) + + assert torch.equal(result[:, :, :1], torch.ones_like(result[:, :, :1])) + + +def test_refiner_accepts_captured_noise() -> None: + stage = LingBotVideoRefinerStage( + denoising_stage=_DenoiseStage(), + vae_encode_stage=_EncodeStage(), + vae_decode_stage=_DecodeStage(), + scheduler=_Scheduler(), + ) + video = torch.zeros(1, 3, 2, 1, 1) + noise = torch.zeros(1, 1, 2, 1, 1) + + result = stage.refine( + video, + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + None, + None, + num_inference_steps=1, + guidance_scale=1.0, + shift=1.0, + t_thresh=0.5, + tail_steps=0, + noise=noise, + ) + + assert torch.equal(result, -torch.ones_like(result)) + + +def test_refiner_resolves_distributed_denoising_result() -> None: + stage = LingBotVideoRefinerStage( + denoising_stage=_DeferredDenoiseStage(), + vae_encode_stage=_EncodeStage(), + vae_decode_stage=_DecodeStage(), + scheduler=_Scheduler(), + ) + + result = stage.refine( + torch.zeros(1, 3, 2, 1, 1), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + None, + None, + num_inference_steps=1, + guidance_scale=1.0, + shift=1.0, + t_thresh=0.5, + tail_steps=0, + noise=torch.zeros(1, 1, 2, 1, 1), + ) + + assert torch.equal(result, -torch.ones_like(result)) + + +def test_refiner_releases_vae_before_denoising_and_worker_before_decode() -> None: + encode = _LifecycleEncodeStage() + decode = _LifecycleDecodeStage() + denoise = _LifecycleDenoiseStage() + stage = LingBotVideoRefinerStage( + denoising_stage=denoise, + vae_encode_stage=encode, + vae_decode_stage=decode, + scheduler=_Scheduler(), + ) + + stage.refine( + torch.zeros(1, 3, 2, 1, 1), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + None, + None, + num_inference_steps=1, + guidance_scale=1.0, + shift=1.0, + t_thresh=0.5, + tail_steps=0, + noise=torch.zeros(1, 1, 2, 1, 1), + ) + + assert encode.offloaded + assert decode.offloaded + assert denoise.closed + + +def test_refiner_runs_release_callback_after_denoising_and_before_decode() -> None: + events: list[str] = [] + + class RecordingDecodeStage(_DecodeStage): + def decode(self, values: torch.Tensor) -> torch.Tensor: + events.append("decode") + return values + + denoise = _LifecycleDenoiseStage() + stage = LingBotVideoRefinerStage( + denoising_stage=denoise, + vae_encode_stage=_EncodeStage(), + vae_decode_stage=RecordingDecodeStage(), + scheduler=_Scheduler(), + ) + + def release_base() -> None: + assert denoise.closed + events.append("release_base") + + stage.refine( + torch.zeros(1, 3, 2, 1, 1), + torch.zeros(1, 1, 1), + torch.zeros(1, 1, 1), + None, + None, + num_inference_steps=1, + guidance_scale=1.0, + shift=1.0, + t_thresh=0.5, + tail_steps=0, + noise=torch.zeros(1, 1, 2, 1, 1), + before_decode=release_base, + ) + + assert events == ["release_base", "decode"] diff --git a/tests/unit/pipelines/lingbot_video/test_runtime.py b/tests/unit/pipelines/lingbot_video/test_runtime.py new file mode 100644 index 0000000..a4412fd --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -0,0 +1,110 @@ +"""Tests for LingBot checkpoint assembly and memory lifecycle defaults.""" + +from __future__ import annotations + +import sys +import types +from pathlib import Path + +from examples.lingbot_video import lingbot_video_moe_30b as moe_example +from telefuser.core.config import WeightOffloadType +from telefuser.pipelines.lingbot_video import LingBotVideoModelConfig + + +class _FakeModule: + def __init__(self) -> None: + self.device: str | None = None + self.expert_execution_backend: str | None = None + + def to(self, device: str) -> _FakeModule: + self.device = device + return self + + def eval(self) -> _FakeModule: + return self + + def set_attention_config(self, attention_config) -> None: + self.attention_config = attention_config + + def set_expert_execution_backend(self, backend: str) -> None: + self.expert_execution_backend = backend + + def promote_stability_layers_to_fp32(self) -> None: + self.promoted_stability_layers = True + + +def test_moe_example_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: + loaded_transformers: list[tuple[str, str | None]] = [] + + class _Processor: + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return object() + + class _TextEncoder(_FakeModule): + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return cls() + + class _VAE(_FakeModule): + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return cls() + + class _Scheduler: + @classmethod + def from_pretrained(cls, *args, **kwargs): + del args, kwargs + return cls() + + diffusers = types.ModuleType("diffusers") + diffusers.AutoencoderKLWan = _VAE + transformers = types.ModuleType("transformers") + transformers.AutoProcessor = _Processor + transformers.Qwen3VLForConditionalGeneration = _TextEncoder + monkeypatch.setitem(sys.modules, "diffusers", diffusers) + monkeypatch.setitem(sys.modules, "transformers", transformers) + monkeypatch.setattr( + moe_example, + "load_lingbot_video_model_config", + lambda *args, **kwargs: LingBotVideoModelConfig(variant="moe", num_experts=2, top_k=1), + ) + monkeypatch.setattr(moe_example.FlowUniPCMultistepScheduler, "from_pretrained", _Scheduler.from_pretrained) + + def load_model(self, file_path, *, name=None, **kwargs) -> None: + del kwargs + loaded_transformers.append((file_path, name)) + self.add_module(_FakeModule(), name=name or "transformer", path=file_path) + + monkeypatch.setattr(moe_example.ModuleManager, "load_model", load_model) + + def load_from_huggingface(self, module_path, *, module_name=None, module_class=None, **kwargs) -> None: + del kwargs + assert module_class is not None + self.add_module( + module_class.from_pretrained(module_path), + name=module_name or Path(module_path).name, + path=str(module_path), + ) + + monkeypatch.setattr(moe_example.ModuleManager, "load_from_huggingface", load_from_huggingface) + + base = moe_example.build_pipeline(tmp_path, device="cuda") + refiner = moe_example.build_refiner(tmp_path, device="cuda") + + assert loaded_transformers == [ + (str(tmp_path / "transformer"), "transformer"), + (str(tmp_path / "refiner"), "transformer"), + ] + assert base.variant == "moe" + assert base.denoising_stage.transformer.expert_execution_backend == "sorted" + assert ( + base.denoising_stage.transformer.attention_config is base.denoising_stage.model_runtime_config.attention_config + ) + assert base.denoising_stage.model_runtime_config.offload_config.offload_type is WeightOffloadType.MODEL_CPU_OFFLOAD + assert ( + refiner.denoising_stage.model_runtime_config.offload_config.offload_type is WeightOffloadType.MODEL_CPU_OFFLOAD + ) diff --git a/tests/unit/pipelines/lingbot_video/test_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py new file mode 100644 index 0000000..4a67157 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -0,0 +1,340 @@ +"""Tests for the split LingBot-Video CLI and service examples.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import ModuleType, SimpleNamespace + +import numpy as np +import pytest +import torch +from click.testing import CliRunner + +from telefuser.pipelines.lingbot_video import DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT_IMAGE +from telefuser.service.core.pipeline_contract import load_pipeline_contract + +EXAMPLE_PATHS = { + "dense": Path("examples/lingbot_video/lingbot_video_dense_1_3b.py"), + "moe": Path("examples/lingbot_video/lingbot_video_moe_30b.py"), +} + + +def _load_example(variant: str) -> ModuleType: + path = EXAMPLE_PATHS[variant] + module_name = f"lingbot_video_{variant}_example_test" + spec = importlib.util.spec_from_file_location(module_name, path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[module_name] = module + spec.loader.exec_module(module) + return module + + +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_examples_expose_cli_and_service_entrypoints(variant: str) -> None: + module = _load_example(variant) + + for name in ("PPL_CONFIG", "PIPELINE_CONTRACT", "get_pipeline", "run", "run_with_file"): + assert hasattr(module, name) + assert module.PPL_CONFIG["variant"] == variant + + prompt_payload = json.loads(module.PPL_CONFIG["prompt"]) + assert prompt_payload["duration"] == 5 + assert "comprehensive_description" in prompt_payload["caption"] + contract, is_explicit = load_pipeline_contract( + module, + ppl_file=EXAMPLE_PATHS[variant].name, + default_task="t2v", + ) + assert is_explicit + assert contract.pipeline_name == module.PPL_CONFIG["name"] + assert contract.supported_tasks == ("t2i", "t2v", "i2v") + assert contract.get_task_contract("t2i").parameters["negative_prompt"].default == DEFAULT_NEGATIVE_PROMPT_IMAGE + assert contract.get_task_contract("t2v").parameters["negative_prompt"].default == DEFAULT_NEGATIVE_PROMPT + assert contract.get_task_contract("i2v").required_inputs == ("first_image_path",) + + +def test_only_moe_contract_exposes_refiner() -> None: + dense = _load_example("dense") + moe = _load_example("moe") + + dense_contract, _ = load_pipeline_contract(dense, ppl_file=EXAMPLE_PATHS["dense"].name, default_task="t2v") + moe_contract, _ = load_pipeline_contract(moe, ppl_file=EXAMPLE_PATHS["moe"].name, default_task="t2v") + + assert "refine" not in dense_contract.get_task_contract("t2v").parameters + assert "refine" not in moe_contract.get_task_contract("t2i").parameters + assert moe_contract.get_task_contract("t2v").parameters["refine"].default is True + + +@pytest.mark.parametrize( + ("variant", "checkpoint"), + [ + ("dense", "lingbot-video-dense-1.3b"), + ("moe", "lingbot-video-moe-30b-a3b"), + ], +) +def test_default_model_root_uses_tf_model_zoo_path( + variant: str, checkpoint: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setenv("TF_MODEL_ZOO_PATH", "/models") + + module = _load_example(variant) + + assert module.PPL_CONFIG["model_root"] == f"/models/lingbot/{checkpoint}" + + +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_get_pipeline_uses_fixed_variant(variant: str, monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example(variant) + calls: list[tuple[object, dict[str, object]]] = [] + pipeline = SimpleNamespace() + + def capture_build(model_root: str, **kwargs: object) -> object: + calls.append((model_root, kwargs)) + return pipeline + + monkeypatch.setattr(module, "build_pipeline", capture_build) + result = module.get_pipeline(parallelism=4, model_root="checkpoint") + + assert result is pipeline + assert calls[0][0] == "checkpoint" + parallel_config = calls[0][1]["parallel_config"] + assert parallel_config.device_ids == [0, 1, 2, 3] + assert parallel_config.sp_ulysses_degree == 4 + assert parallel_config.enable_fsdp + assert module.PPL_CONFIG["variant"] == variant + + +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_get_pipeline_supports_cfg_parallel_with_sequence_parallel( + variant: str, monkeypatch: pytest.MonkeyPatch +) -> None: + module = _load_example(variant) + calls: list[dict[str, object]] = [] + + def capture_build(_: str, **kwargs: object) -> SimpleNamespace: + calls.append(kwargs) + return SimpleNamespace() + + monkeypatch.setattr(module, "build_pipeline", capture_build) + pipeline = module.get_pipeline(parallelism=4, cfg_parallel_degree=2) + + parallel_config = calls[0]["parallel_config"] + assert parallel_config.cfg_degree == 2 + assert parallel_config.sp_ulysses_degree == 2 + assert parallel_config.enable_fsdp + if variant == "moe": + assert pipeline.refiner_co_resident + + +def test_dense_run_with_file_supports_service_t2i(tmp_path: Path) -> None: + module = _load_example("dense") + + class FakePipeline: + device = "cpu" + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.num_frames == 1 + assert request.caption == '{"scene":"test"}' + assert (request.height, request.width) == (480, 832) + return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2)) + + output_path = tmp_path / "dense.png" + result = module.run_with_file( + FakePipeline(), + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2i", + output_path=str(output_path), + ) + + assert result == {"output_path": str(output_path)} + assert output_path.is_file() + + +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_t2i_replaces_video_output_suffix_with_png(variant: str, tmp_path: Path) -> None: + module = _load_example(variant) + + result = module._save_output(torch.zeros(1, 3, 1, 2, 2), str(tmp_path / "result.mp4")) + + expected = tmp_path / "result.png" + assert result == {"output_path": str(expected)} + assert expected.is_file() + + +def test_dense_run_with_file_supports_service_t2v(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example("dense") + exported: list[np.ndarray] = [] + + def capture_export(frames: list[np.ndarray], output_path: str, fps: int) -> None: + del output_path + assert fps == 24 + exported.extend(frames) + + monkeypatch.setattr(module, "export_to_video", capture_export) + + class FakePipeline: + device = "cpu" + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.num_frames == 49 + return SimpleNamespace(output=torch.full((1, 3, 5, 2, 2), 0.5)) + + result = module.run_with_file( + FakePipeline(), + json.dumps({"caption": {"scene": "test"}, "duration": 2}), + task="t2v", + output_path=str(tmp_path / "dense.mp4"), + ) + + assert result == {"output_path": str(tmp_path / "dense.mp4")} + assert len(exported) == 5 + assert exported[0].dtype == np.float32 + assert float(exported[0][0, 0, 0]) == 0.5 + + +def test_moe_run_releases_base_then_runs_refiner(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example("moe") + + class FakePipeline: + device = "cpu" + text_stage = object() + + def __init__(self) -> None: + self.released = False + + def generate(self, request, **_: object) -> SimpleNamespace: + del request + conditions = SimpleNamespace( + has_visual_condition=False, + positive_prompt_embeds=torch.full((1, 2, 2), 4.0), + negative_prompt_embeds=torch.full((1, 2, 2), 7.0), + positive_attention_mask=torch.tensor([[True, False]]), + negative_attention_mask=torch.tensor([[False, True]]), + ) + return SimpleNamespace(output=torch.zeros(1, 3, 5, 2, 2), prompt_conditions=conditions) + + def release_gpu_resources(self) -> None: + self.released = True + + class FakeRefiner: + def __init__(self) -> None: + self.called = False + self.closed = False + self.negative: torch.Tensor | None = None + self.negative_mask: torch.Tensor | None = None + + def refine(self, lowres_video, positive, negative, positive_mask, negative_mask, **kwargs) -> torch.Tensor: + del positive, positive_mask, kwargs + self.called = True + self.negative = negative + self.negative_mask = negative_mask + return lowres_video + + def close(self) -> None: + self.closed = True + + pipeline = FakePipeline() + refiner = FakeRefiner() + monkeypatch.setattr(module, "build_refiner", lambda *args, **kwargs: refiner) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_height", 2) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_width", 2) + + frames = module.run( + pipeline, + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2v", + refine=True, + ) + + assert frames.shape == (1, 3, 5, 2, 2) + assert pipeline.released + assert refiner.called + assert refiner.closed + assert torch.equal(refiner.negative, torch.zeros(1, 2, 2)) + assert torch.equal(refiner.negative_mask, torch.tensor([[True, False]])) + + +def test_moe_run_keeps_base_resident_when_requested(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example("moe") + + class FakePipeline: + device = "cpu" + text_stage = object() + refiner_co_resident = True + + def __init__(self) -> None: + self.released = False + + def generate(self, request: object, **_: object) -> SimpleNamespace: + del request + conditions = SimpleNamespace( + has_visual_condition=False, + positive_prompt_embeds=torch.zeros(1, 2, 2), + negative_prompt_embeds=torch.zeros(1, 2, 2), + positive_attention_mask=torch.ones(1, 2, dtype=torch.bool), + negative_attention_mask=torch.ones(1, 2, dtype=torch.bool), + ) + return SimpleNamespace(output=torch.zeros(1, 3, 5, 2, 2), prompt_conditions=conditions) + + def release_gpu_resources(self) -> None: + self.released = True + + class FakeRefiner: + def refine(self, lowres_video: torch.Tensor, *_: object, **__: object) -> torch.Tensor: + return lowres_video + + def close(self) -> None: + return None + + pipeline = FakePipeline() + monkeypatch.setattr(module, "build_refiner", lambda *args, **kwargs: FakeRefiner()) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_height", 2) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_width", 2) + + module.run( + pipeline, + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2v", + refine=True, + ) + + assert not pipeline.released + + +def test_moe_t2i_rejects_explicit_refiner_request() -> None: + module = _load_example("moe") + + with pytest.raises(ValueError, match="does not support t2i"): + module.run( + SimpleNamespace(), + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2i", + refine=True, + ) + + +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_cli_invokes_model_specific_entrypoints(variant: str, monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example(variant) + pipeline = SimpleNamespace(stop=lambda: None) + calls: list[tuple[object, dict[str, object]]] = [] + + monkeypatch.setattr(module, "get_pipeline", lambda *args, **kwargs: pipeline) + + def capture_run_with_file(pipe: object, *args: object, **kwargs: object) -> dict[str, str]: + calls.append((pipe, {"args": args, **kwargs})) + return {"output_path": "result.mp4"} + + monkeypatch.setattr(module, "run_with_file", capture_run_with_file) + result = CliRunner().invoke( + module.main, + ["--output_path", "result.mp4", "--no-refine"] if variant == "moe" else ["--output_path", "result.mp4"], + ) + + assert result.exit_code == 0, result.output + assert calls and calls[0][0] is pipeline + assert "Output saved to result.mp4" in result.output diff --git a/tests/unit/pipelines/lingbot_video/test_text_encoding.py b/tests/unit/pipelines/lingbot_video/test_text_encoding.py new file mode 100644 index 0000000..0289573 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_text_encoding.py @@ -0,0 +1,26 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video.text_encoding import LingBotVideoTextEncodingStage + + +class _VisionEncoder: + config = SimpleNamespace(vision_config=SimpleNamespace(patch_size=14)) + + +def test_prepare_ti2v_vlm_image_matches_source_smart_resize() -> None: + module_manager = ModuleManager(device="cpu", torch_dtype=torch.float32) + module_manager.add_module(_VisionEncoder(), name="text_encoder") + module_manager.add_module(object(), name="processor") + stage = LingBotVideoTextEncodingStage( + "text", module_manager, ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32) + ) + + image = stage.prepare_ti2v_vlm_image(torch.zeros(1, 3, 1, 192, 320)) + + assert image.size == (308, 196) diff --git a/tests/unit/pipelines/lingbot_video/test_vae.py b/tests/unit/pipelines/lingbot_video/test_vae.py new file mode 100644 index 0000000..8e3001c --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_vae.py @@ -0,0 +1,58 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from torch import nn + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video.vae import ( + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + denormalize_latent, +) + + +class _Posterior: + def __init__(self, sample: torch.Tensor) -> None: + self._sample = sample + + def sample(self, generator: torch.Generator | None = None) -> torch.Tensor: + return self._sample + + +class _VAE(nn.Module): + config = SimpleNamespace(latents_mean=[1.0, 2.0], latents_std=[2.0, 4.0]) + + def encode(self, values: torch.Tensor) -> SimpleNamespace: + return SimpleNamespace(latent_dist=_Posterior(torch.ones(values.shape[0], 2, 1, 1, 1))) + + def decode(self, latents: torch.Tensor) -> SimpleNamespace: + assert torch.allclose(latents, torch.ones_like(latents)) + assert latents.is_contiguous(memory_format=torch.channels_last_3d) + return SimpleNamespace(sample=torch.zeros(latents.shape[0], 3, 1, 1, 1)) + + +def test_vae_stages_apply_checkpoint_normalization() -> None: + runtime = ModelRuntimeConfig(device_type="cpu") + vae = _VAE() + module_manager = ModuleManager(device="cpu", torch_dtype=torch.float32) + module_manager.add_module(vae, name="vae") + encoded = LingBotVideoVAEEncodeStage("encode", module_manager, runtime).encode(torch.ones(1, 3, 1, 2, 2)) + decoded = LingBotVideoVAEDecodeStage("decode", module_manager, runtime).decode(encoded) + + assert encoded[:, 0].item() == 0.0 + assert encoded[:, 1].item() == -0.25 + assert torch.allclose(decoded, torch.full_like(decoded, 0.5)) + + +def test_denormalize_preserves_upstream_division_by_inverse_order() -> None: + latent = torch.tensor([[[[[0.12345679]]]]]) + mean = torch.tensor([0.27182818]) + std = torch.tensor([0.31415927]) + + actual = denormalize_latent(latent, mean, std) + expected = latent / std.reciprocal().reshape(1, 1, 1, 1, 1) + mean.reshape(1, 1, 1, 1, 1) + + assert torch.equal(actual, expected) diff --git a/tools/validation/benchmark_lingbot_video.py b/tools/validation/benchmark_lingbot_video.py new file mode 100644 index 0000000..79dd6c1 --- /dev/null +++ b/tools/validation/benchmark_lingbot_video.py @@ -0,0 +1,288 @@ +"""Benchmark a native LingBot-Video base or base-plus-refiner run on one GPU.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from collections import defaultdict +from collections.abc import Callable +from pathlib import Path +from typing import Any + +import torch +from PIL import Image + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline as build_dense_pipeline +from examples.lingbot_video.lingbot_video_moe_30b import build_pipeline as build_moe_pipeline +from examples.lingbot_video.lingbot_video_moe_30b import build_refiner +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + default_negative_caption, + load_lingbot_video_prompt, + load_refiner_first_frame, + num_frames_from_duration, + prepare_refiner_video, +) + +BenchmarkSamples = dict[str, list[float]] + + +def _resolve_negative_caption(negative_caption: str | None, num_frames: int) -> str: + """Keep benchmark sampling aligned with the pipeline's source defaults.""" + return default_negative_caption(num_frames) if negative_caption is None else negative_caption + + +def _synchronize() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _measure(samples: BenchmarkSamples, name: str, operation: Callable[[], Any]) -> Any: + _synchronize() + start = time.perf_counter() + result = operation() + _synchronize() + samples[name].append(time.perf_counter() - start) + return result + + +def _instrument_method( + samples_provider: Callable[[], BenchmarkSamples], owner: object, attribute: str, metric_name: str +) -> None: + original = getattr(owner, attribute) + + def measured(*args: object, **kwargs: object) -> Any: + return _measure(samples_provider(), metric_name, lambda: original(*args, **kwargs)) + + setattr(owner, attribute, measured) + + +def _image_to_tensor(path: str) -> torch.Tensor: + image = Image.open(path).convert("RGB") + return torch.from_numpy(__import__("numpy").asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).float() + + +def _summary(samples: BenchmarkSamples) -> dict[str, dict[str, float | int]]: + return { + name: { + "count": len(values), + "total_seconds": sum(values), + "mean_seconds": sum(values) / len(values), + "max_seconds": max(values), + } + for name, values in sorted(samples.items()) + } + + +def _encode_output(frames: torch.Tensor, output_path: Path, fps: int) -> None: + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + if video.shape[0] == 1: + Image.fromarray((video[0] * 255).round().astype("uint8")).save(output_path) + return + from diffusers.utils import export_to_video + + export_to_video(list(video), str(output_path), fps=fps) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--caption-json", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--variant", choices=("dense", "moe"), default="dense") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int) + parser.add_argument("--steps", type=int, default=40) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--negative-caption", default=None) + parser.add_argument("--image") + parser.add_argument("--refine", action="store_true") + parser.add_argument("--refiner-height", type=int, default=1088) + parser.add_argument("--refiner-width", type=int, default=1920) + parser.add_argument("--refiner-steps", type=int, default=8) + parser.add_argument("--refiner-guidance-scale", type=float, default=3.0) + parser.add_argument("--refiner-shift", type=float, default=3.0) + parser.add_argument("--refiner-t-thresh", type=float, default=0.85) + parser.add_argument("--refiner-tail-steps", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--disable-cpu-offload", action="store_true") + parser.add_argument("--warmup", type=int, default=1, help="Number of unreported warmup generations.") + parser.add_argument("--runs", type=int, default=1, help="Number of measured generations.") + args = parser.parse_args() + if args.refine and args.variant != "moe": + raise ValueError("--refine requires --variant moe") + if args.warmup < 0 or args.runs < 1: + raise ValueError("--warmup must be non-negative and --runs must be positive") + + caption, duration = load_lingbot_video_prompt(args.caption_json) + num_frames = ( + args.num_frames + if args.num_frames is not None + else num_frames_from_duration(duration) + if duration is not None + else 121 + ) + negative_caption = _resolve_negative_caption(args.negative_caption, num_frames) + setup_samples: BenchmarkSamples = defaultdict(list) + warmup_samples: BenchmarkSamples = defaultdict(list) + measured_samples: BenchmarkSamples = defaultdict(list) + active_samples = measured_samples + + def get_active_samples() -> BenchmarkSamples: + return active_samples + + if torch.cuda.is_available(): + torch.cuda.reset_peak_memory_stats() + pipeline_builder = build_dense_pipeline if args.variant == "dense" else build_moe_pipeline + pipeline = _measure( + setup_samples, + "base_load", + lambda: pipeline_builder( + args.model_dir, + guidance_scale=args.guidance_scale, + num_inference_steps=args.steps, + cpu_offload=False if args.disable_cpu_offload else None, + ), + ) + if pipeline.text_stage is None or pipeline.denoising_stage is None or pipeline.vae_decode_stage is None: + raise RuntimeError("LingBot runtime did not create all required base stages") + _instrument_method(get_active_samples, pipeline.text_stage, "encode", "base_text_encode") + _instrument_method(get_active_samples, pipeline.denoising_stage, "predict_noise_with_cfg", "base_denoise_step") + if pipeline.vae_encode_stage is not None: + _instrument_method(get_active_samples, pipeline.vae_encode_stage, "encode", "base_vae_encode") + _instrument_method(get_active_samples, pipeline.vae_decode_stage, "decode", "base_vae_decode") + source_image = _image_to_tensor(args.image) if args.image else None + refiner_prompt_conditions_reused = False + refiner = None + + def generate_once() -> torch.Tensor: + nonlocal refiner_prompt_conditions_reused, refiner + generation = _measure( + active_samples, + "base_total", + lambda: pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=args.height, + width=args.width, + num_frames=num_frames, + image=source_image, + ), + negative_caption=negative_caption, + generator=torch.Generator("cuda").manual_seed(args.seed), + ), + ) + frames = generation.output + if not args.refine: + return frames + if pipeline.text_stage is None: + raise RuntimeError("LingBot refiner requires the base text stage") + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = _measure( + active_samples, "refiner_text_encode", lambda: pipeline.text_stage.encode(caption) + ) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + refiner_prompt_conditions_reused = True + negative = torch.zeros_like(positive) + negative_mask = positive_mask.clone() + _measure(active_samples, "base_release", pipeline.release_gpu_resources) + if refiner is None: + refiner = _measure( + setup_samples, + "refiner_load", + lambda: build_refiner(args.model_dir, cpu_offload=not args.disable_cpu_offload), + ) + _instrument_method( + get_active_samples, + refiner.denoising_stage, + "predict_noise_with_cfg", + "refiner_denoise_step", + ) + _instrument_method(get_active_samples, refiner.vae_encode_stage, "encode", "refiner_vae_encode") + _instrument_method(get_active_samples, refiner.vae_decode_stage, "decode", "refiner_vae_decode") + lowres_video, _ = _measure( + active_samples, + "refiner_prepare", + lambda: prepare_refiner_video( + frames, source_fps=24.0, height=args.refiner_height, width=args.refiner_width + ), + ) + clean_first_frame = ( + _measure( + active_samples, + "refiner_first_frame_prepare", + lambda: load_refiner_first_frame( + args.image, + target_height=args.refiner_height, + target_width=args.refiner_width, + geometry_height=args.height, + geometry_width=args.width, + ), + ) + if args.image + else None + ) + return _measure( + active_samples, + "refiner_total", + lambda: refiner.refine( + lowres_video, + positive, + negative, + positive_mask, + negative_mask, + num_inference_steps=args.refiner_steps, + guidance_scale=args.refiner_guidance_scale, + shift=args.refiner_shift, + t_thresh=args.refiner_t_thresh, + tail_steps=args.refiner_tail_steps, + clean_first_frame=clean_first_frame, + generator=torch.Generator("cuda").manual_seed(args.seed), + ), + ) + + for _ in range(args.warmup): + active_samples = warmup_samples + generate_once() + frames = None + for _ in range(args.runs): + active_samples = measured_samples + frames = generate_once() + if frames is None: + raise RuntimeError("benchmark did not produce a measured output") + args.output.parent.mkdir(parents=True, exist_ok=True) + _measure(measured_samples, "output_encode", lambda: _encode_output(frames, args.output, fps=24)) + report = { + "model_dir": str(args.model_dir), + "variant": args.variant, + "refine": args.refine, + "height": args.height, + "width": args.width, + "num_frames": num_frames, + "steps": args.steps, + "warmup_runs": args.warmup, + "measured_runs": args.runs, + "cpu_offload": not args.disable_cpu_offload if args.variant == "moe" else False, + "refiner_prompt_conditions_reused": refiner_prompt_conditions_reused, + "setup_metrics": _summary(setup_samples), + "warmup_metrics": _summary(warmup_samples), + "metrics": _summary(measured_samples), + "peak_gpu_memory_mib": torch.cuda.max_memory_allocated() / 1024**2 if torch.cuda.is_available() else None, + } + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text(json.dumps(report, indent=2, sort_keys=True) + "\n", encoding="utf-8") + print(json.dumps(report, indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/capture_lingbot_video_reference.py b/tools/validation/capture_lingbot_video_reference.py new file mode 100644 index 0000000..db35261 --- /dev/null +++ b/tools/validation/capture_lingbot_video_reference.py @@ -0,0 +1,589 @@ +"""Capture reproducible LingBot-Video upstream reference artifacts. + +This tool deliberately executes the checked-out upstream Diffusers implementation. +It does not copy or modify upstream source files. The captured artifacts are intended +to remain under ``work_dirs/`` and provide the numerical oracle for later phases. +""" + +from __future__ import annotations + +import argparse +import contextlib +import hashlib +import importlib +import json +import os +import platform +import sys +import time +from collections.abc import Callable, Iterator +from dataclasses import dataclass, field +from datetime import datetime, timezone +from pathlib import Path +from types import ModuleType +from typing import Any + +import numpy as np +import torch +from PIL import Image +from safetensors.torch import load_model + +DEFAULT_UPSTREAM_ROOT = Path("work_dirs/lingbot-video-master") +DEFAULT_CASE_MANIFEST = "assets/cases/manifest.json" +DEFAULT_DENSE_MODEL_DIR = Path("/hhb-data/aigc/model_zoo/lingbot/lingbot-video-dense-1.3b") + + +def _sha256_file(path: Path) -> str: + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1024 * 1024), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _tree_fingerprint(root: Path, *, include_contents: bool) -> str: + """Return a stable fingerprint without loading checkpoint tensors into memory.""" + digest = hashlib.sha256() + for path in sorted(item for item in root.rglob("*") if item.is_file()): + relative = path.relative_to(root).as_posix() + stat = path.stat() + digest.update(f"{relative}\0{stat.st_size}\0".encode()) + if include_contents: + digest.update(_sha256_file(path).encode()) + return digest.hexdigest() + + +def _json_value(value: Any) -> Any: + if isinstance(value, Path): + return str(value) + if isinstance(value, torch.dtype): + return str(value).removeprefix("torch.") + if isinstance(value, torch.device): + return str(value) + if isinstance(value, np.generic): + return value.item() + raise TypeError(f"{type(value).__name__} is not JSON serializable") + + +def _tensor_summary(tensor: torch.Tensor) -> dict[str, Any]: + cpu = tensor.detach().cpu().clone(memory_format=torch.contiguous_format) + finite = torch.isfinite(cpu) + summary: dict[str, Any] = { + "shape": list(cpu.shape), + "dtype": str(cpu.dtype).removeprefix("torch."), + "numel": cpu.numel(), + "nan_count": int(torch.isnan(cpu).sum().item()) if cpu.is_floating_point() else 0, + "inf_count": int(torch.isinf(cpu).sum().item()) if cpu.is_floating_point() else 0, + "sha256": hashlib.sha256(cpu.reshape(-1).view(torch.uint8).numpy().tobytes()).hexdigest(), + } + if cpu.is_floating_point() and finite.any(): + values = cpu[finite].float() + summary.update( + min=float(values.min().item()), + max=float(values.max().item()), + mean=float(values.mean().item()), + std=float(values.std(unbiased=False).item()), + ) + return summary + + +def _first_tensor(value: Any) -> torch.Tensor | None: + if isinstance(value, torch.Tensor): + return value + if isinstance(value, (tuple, list)): + for item in value: + tensor = _first_tensor(item) + if tensor is not None: + return tensor + sample = getattr(value, "sample", None) + return sample if isinstance(sample, torch.Tensor) else None + + +@contextlib.contextmanager +def _upstream_import_path(root: Path) -> Iterator[None]: + """Temporarily expose the upstream checkout without editing it.""" + root = root.resolve() + if not (root / "lingbot_video" / "__init__.py").is_file(): + raise FileNotFoundError(f"LingBot-Video package not found under {root}") + sys.path.insert(0, str(root)) + try: + yield + finally: + sys.path.remove(str(root)) + + +def _import_upstream(root: Path) -> ModuleType: + with _upstream_import_path(root): + return importlib.import_module("lingbot_video.runner") + + +def _sample_cases(upstream_root: Path, modes: list[str], case_name: str | None) -> list[dict[str, Any]]: + manifest_path = upstream_root / DEFAULT_CASE_MANIFEST + manifest = json.loads(manifest_path.read_text(encoding="utf-8")) + selected: list[dict[str, Any]] = [] + for mode in modes: + matches = [ + item + for item in manifest["examples"] + if item["mode"] == mode and (case_name is None or item["name"] == case_name) + ] + if not matches: + requested = "all cases" if case_name is None else f"case {case_name!r}" + raise ValueError(f"No {mode} {requested} in {manifest_path}") + selected.extend(matches) + return selected + + +def _selected_step_indices(step_count: int, trace: str) -> set[int]: + if trace == "full": + return set(range(step_count)) + if trace == "sampled" and step_count: + return {0, step_count // 2, step_count - 1} + return set() + + +@dataclass +class ArtifactRecorder: + root: Path + trace: str + tensor_metadata: dict[str, dict[str, Any]] = field(default_factory=dict) + events: list[dict[str, Any]] = field(default_factory=list) + step_count: int = 0 + + def __post_init__(self) -> None: + (self.root / "tensors").mkdir(parents=True, exist_ok=True) + + @property + def selected_steps(self) -> set[int]: + return _selected_step_indices(self.step_count, self.trace) + + def record_tensor(self, name: str, value: Any) -> None: + tensor = _first_tensor(value) + if tensor is None: + return + safe_name = name.replace("/", "_").replace(" ", "_") + path = self.root / "tensors" / f"{safe_name}.pt" + cpu = tensor.detach().cpu().clone(memory_format=torch.contiguous_format) + torch.save(cpu, path) + metadata = _tensor_summary(cpu) + metadata["path"] = str(path.relative_to(self.root)) + self.tensor_metadata[safe_name] = metadata + + def event(self, name: str, **data: Any) -> None: + self.events.append({"name": name, **data}) + + +class PipelineTrace: + """Install reversible hooks around stable upstream pipeline boundaries.""" + + def __init__(self, pipe: Any, recorder: ArtifactRecorder): + self.pipe = pipe + self.recorder = recorder + self._restore: list[Callable[[], None]] = [] + self._prompt_call = 0 + self._transformer_call = 0 + self._step = 0 + + def _replace_method(self, obj: Any, name: str, wrapper: Callable[..., Any]) -> None: + original = getattr(obj, name) + setattr(obj, name, wrapper(original)) + self._restore.append(lambda: setattr(obj, name, original)) + + def __enter__(self) -> "PipelineTrace": + self._replace_method(self.pipe, "_build_prompt_inputs", self._wrap_prompt_inputs) + self._replace_method(self.pipe, "encode_prompt", self._wrap_encode_prompt) + self._replace_method(self.pipe, "prepare_latents", self._wrap_prepare_latents) + self._replace_method(self.pipe.scheduler, "set_timesteps", self._wrap_set_timesteps) + self._replace_method(self.pipe.scheduler, "step", self._wrap_scheduler_step) + self._replace_method(self.pipe.transformer, "forward", self._wrap_transformer_forward) + if hasattr(self.pipe, "encode_image_latent"): + self._replace_method(self.pipe, "encode_image_latent", self._wrap_encode_image_latent) + if hasattr(self.pipe, "preprocess_image"): + self._replace_method(self.pipe, "preprocess_image", self._wrap_preprocess_image) + if getattr(self.pipe, "vae", None) is not None: + self._replace_method(self.pipe.vae, "decode", self._wrap_vae_decode) + return self + + def __exit__(self, *_: Any) -> None: + for restore in reversed(self._restore): + restore() + + def _wrap_prompt_inputs(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + inputs = original(*args, **kwargs) + for name in ("input_ids", "attention_mask", "pixel_values", "image_grid_thw"): + value = inputs.get(name) if hasattr(inputs, "get") else None + if value is not None: + self.recorder.record_tensor(f"prompt_inputs_{self._prompt_call}_{name}", value) + return inputs + + return wrapped + + def _wrap_encode_prompt(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + call = self._prompt_call + self._prompt_call += 1 + output = original(*args, **kwargs) + if isinstance(output, tuple): + self.recorder.record_tensor(f"prompt_{call}_embeds", output[0]) + self.recorder.record_tensor(f"prompt_{call}_mask", output[1]) + self.recorder.event("encode_prompt", call=call) + return output + + return wrapped + + def _wrap_prepare_latents(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + output = original(*args, **kwargs) + self.recorder.record_tensor("initial_latents", output) + return output + + return wrapped + + def _wrap_set_timesteps(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + output = original(*args, **kwargs) + timesteps = getattr(self.pipe.scheduler, "timesteps", None) + sigmas = getattr(self.pipe.scheduler, "sigmas", None) + if isinstance(timesteps, torch.Tensor): + self.recorder.step_count = int(timesteps.numel()) + self.recorder.record_tensor("scheduler_timesteps", timesteps) + if isinstance(sigmas, torch.Tensor): + self.recorder.record_tensor("scheduler_sigmas", sigmas) + return output + + return wrapped + + def _wrap_transformer_forward(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + call = self._transformer_call + self._transformer_call += 1 + if self._step in self.recorder.selected_steps: + self.recorder.record_tensor(f"step_{self._step}_transformer_{call}_latent", args[0]) + self.recorder.record_tensor(f"step_{self._step}_transformer_{call}_timestep", args[1]) + self.recorder.record_tensor(f"step_{self._step}_transformer_{call}_prompt", args[2]) + output = original(*args, **kwargs) + if self._step in self.recorder.selected_steps: + self.recorder.record_tensor(f"step_{self._step}_transformer_{call}_noise", output) + return output + + return wrapped + + def _wrap_scheduler_step(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped( + model_output: torch.Tensor, timestep: torch.Tensor, sample: torch.Tensor, *args: Any, **kwargs: Any + ) -> Any: + selected = self._step in self.recorder.selected_steps + if selected: + self.recorder.record_tensor(f"step_{self._step}_noise_prediction", model_output) + self.recorder.record_tensor(f"step_{self._step}_latent_before", sample) + self.recorder.record_tensor(f"step_{self._step}_timestep", timestep) + output = original(model_output, timestep, sample, *args, **kwargs) + if selected: + self.recorder.record_tensor(f"step_{self._step}_latent_after", output) + self._step += 1 + return output + + return wrapped + + def _wrap_encode_image_latent(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + output = original(*args, **kwargs) + self.recorder.record_tensor("ti2v_condition_latent", output) + return output + + return wrapped + + def _wrap_preprocess_image(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + output = original(*args, **kwargs) + self.recorder.record_tensor("ti2v_preprocessed_image", output) + return output + + return wrapped + + def _wrap_vae_decode(self, original: Callable[..., Any]) -> Callable[..., Any]: + def wrapped(*args: Any, **kwargs: Any) -> Any: + self.recorder.record_tensor("vae_decode_input", args[0]) + output = original(*args, **kwargs) + self.recorder.record_tensor("vae_decode_output", output) + return output + + return wrapped + + +def _generator_state(generator: torch.Generator) -> dict[str, Any]: + return { + "device": str(generator.device), + "state_sha256": hashlib.sha256(generator.get_state().cpu().numpy().tobytes()).hexdigest(), + } + + +def _versions() -> dict[str, str]: + result = {"python": platform.python_version(), "torch": torch.__version__} + for name in ("diffusers", "transformers"): + try: + module = importlib.import_module(name) + result[name] = str(getattr(module, "__version__", "unknown")) + except ImportError: + result[name] = "not-installed" + return result + + +def _write_json(path: Path, value: Any) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text(json.dumps(value, indent=2, sort_keys=True, default=_json_value) + "\n", encoding="utf-8") + + +def _load_reference_pipeline( + runner: ModuleType, model_dir: Path, dtype_map: dict[str, torch.dtype], mode: str, transformer_subfolder: str +) -> Any: + """Load upstream components with a safetensors fallback when accelerate is absent.""" + try: + return runner._load_diffusers_pipe(model_dir, dtype_map, mode=mode, transformer_subfolder=transformer_subfolder) + except ValueError as exc: + if "low_cpu_mem_usage" not in str(exc) or "keep_in_fp32_modules" not in str(exc): + raise + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel + + config = json.loads((model_dir / transformer_subfolder / "config.json").read_text(encoding="utf-8")) + fields = ( + "patch_size", + "in_channels", + "out_channels", + "hidden_size", + "num_attention_heads", + "depth", + "intermediate_size", + "text_dim", + "freq_dim", + "norm_eps", + "rope_theta", + "axes_dims", + "qkv_bias", + "out_bias", + "patch_embed_bias", + "timestep_mlp_bias", + ) + device = runner._default_device() + transformer = LingBotVideoTransformer3DModel(**{field: config[field] for field in fields}).to( + device, dtype_map.get("transformer", dtype_map["default"]) + ) + load_model( + transformer, + model_dir / transformer_subfolder / "diffusion_pytorch_model.safetensors", + strict=True, + device=str(device), + ) + pipeline_class = runner._pipeline_class_for_mode(mode) + with runner._patch_qwen3vl_from_pretrained(): + pipeline = pipeline_class.from_pretrained( + str(model_dir), transformer=transformer, trust_remote_code=True, torch_dtype=dtype_map + ) + return pipeline.to(device) + + +def _capture_once( + *, + runner: ModuleType, + upstream_root: Path, + model_dir: Path, + case: dict[str, Any], + destination: Path, + args: argparse.Namespace, + repeat: int, +) -> dict[str, Any]: + prompt_path = upstream_root / case["prompt_json"] + sample = runner._load_prompt_sample(prompt_path) + prompt = runner._caption_from_sample(sample) + mode = case["mode"] + num_frames = 1 if mode == "t2i" else args.num_frames + image_path = upstream_root / case["image"] if "image" in case else None + dtype_map = { + "default": runner._parse_dtype(args.default_dtype), + "transformer": runner._parse_dtype(args.transformer_dtype), + "text_encoder": runner._parse_dtype(args.text_encoder_dtype), + "vae": runner._parse_dtype(args.vae_dtype), + } + pipe = _load_reference_pipeline( + runner, model_dir.resolve(), dtype_map, mode=mode, transformer_subfolder=args.transformer_subfolder + ) + device = runner._default_device() + generator = torch.Generator(device=device).manual_seed(args.seed) + recorder = ArtifactRecorder(destination, args.trace) + rng_before = _generator_state(generator) + started = time.perf_counter() + call_args: dict[str, Any] = { + "prompt": prompt, + "negative_prompt": runner.DEFAULT_NEGATIVE_PROMPT_IMAGE if mode == "t2i" else runner.DEFAULT_NEGATIVE_PROMPT, + "height": args.height, + "width": args.width, + "num_frames": num_frames, + "num_inference_steps": args.steps, + "guidance_scale": args.guidance_scale, + "shift": args.shift, + "generator": generator, + "output_type": "np", + "batch_cfg": False, + "null_cond_clone_zero": False, + } + if image_path is not None: + call_args["image"] = Image.open(image_path).convert("RGB") + with PipelineTrace(pipe, recorder): + result = pipe(**call_args) + frames = result.frames if hasattr(result, "frames") else result[0] + frames_array = np.asarray(frames[0] if isinstance(frames, list) else frames) + np.save(destination / "frames.npy", frames_array) + metadata = { + "capture_schema_version": 1, + "captured_at": datetime.now(timezone.utc).isoformat(), + "kind": "lingbot_video_numerical_oracle", + "mode": mode, + "case": case["name"], + "repeat": repeat, + "upstream_root": str(upstream_root.resolve()), + "upstream_source_fingerprint": _tree_fingerprint(upstream_root, include_contents=True), + "checkpoint_manifest_fingerprint": _tree_fingerprint(model_dir, include_contents=False), + "prompt_json": str(prompt_path), + "prompt_json_sha256": _sha256_file(prompt_path), + "image": str(image_path) if image_path is not None else None, + "image_sha256": _sha256_file(image_path) if image_path is not None else None, + "sampling": { + key: value + for key, value in call_args.items() + if key not in {"generator", "image", "prompt", "negative_prompt"} + }, + "seed": args.seed, + "prompt": prompt, + "negative_prompt": call_args["negative_prompt"], + "rng_before": rng_before, + "rng_after": _generator_state(generator), + "trace": args.trace, + "events": recorder.events, + "tensors": recorder.tensor_metadata, + "frames": { + "path": "frames.npy", + "shape": list(frames_array.shape), + "dtype": str(frames_array.dtype), + "sha256": hashlib.sha256(frames_array.tobytes()).hexdigest(), + }, + "elapsed_seconds": time.perf_counter() - started, + "environment": { + "versions": _versions(), + "cuda_available": torch.cuda.is_available(), + "cuda_device": torch.cuda.get_device_name(device) if device.type == "cuda" else None, + "device": str(device), + "qwen_attn_implementation": os.environ.get("LINGBOT_QWEN_ATTN_IMPLEMENTATION"), + }, + } + _write_json(destination / "metadata.json", metadata) + return metadata + + +def _repeatability_report(root: Path, captures: list[dict[str, Any]]) -> None: + grouped: dict[tuple[str, str], list[dict[str, Any]]] = {} + for capture in captures: + grouped.setdefault((capture["mode"], capture["case"]), []).append(capture) + report: dict[str, Any] = {"schema_version": 2, "groups": {}} + for (mode, case), values in grouped.items(): + frame_hashes = [item["frames"]["sha256"] for item in values] + shared_names = set(values[0]["tensors"]) + for value in values[1:]: + shared_names.intersection_update(value["tensors"]) + tensor_hashes = {name: [value["tensors"][name]["sha256"] for value in values] for name in sorted(shared_names)} + mismatched_tensors = [name for name, hashes in tensor_hashes.items() if len(set(hashes)) != 1] + report["groups"][f"{mode}/{case}"] = { + "runs": len(values), + "exact_frame_hash_match": len(set(frame_hashes)) == 1, + "frame_hashes": frame_hashes, + "exact_tensor_hash_match": not mismatched_tensors, + "shared_tensor_count": len(tensor_hashes), + "mismatched_tensor_hashes": mismatched_tensors, + "elapsed_seconds": [item["elapsed_seconds"] for item in values], + } + _write_json(root / "repeatability.json", report) + + +def _parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--upstream-root", type=Path, default=DEFAULT_UPSTREAM_ROOT) + parser.add_argument("--model-dir", type=Path, default=DEFAULT_DENSE_MODEL_DIR) + parser.add_argument("--output-dir", type=Path, default=Path("work_dirs/lingbot_video_reference")) + parser.add_argument("--all-cases", action="store_true", help="Capture every bundled case for the selected modes.") + parser.add_argument("--mode", choices=["t2i", "t2v", "ti2v"], action="append") + parser.add_argument("--case", default="example_1") + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--height", type=int, default=192) + parser.add_argument("--width", type=int, default=320) + parser.add_argument("--num-frames", type=int, default=9) + parser.add_argument("--steps", type=int, default=4) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--shift", type=float, default=3.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--trace", choices=["none", "sampled", "full"], default="full") + parser.add_argument("--default-dtype", default="bf16") + parser.add_argument("--transformer-dtype", default="bf16") + parser.add_argument("--text-encoder-dtype", default="bf16") + parser.add_argument("--vae-dtype", default="fp32") + parser.add_argument("--transformer-subfolder", default="transformer") + parser.add_argument("--qwen-attn-implementation", default="sdpa") + parser.add_argument("--dry-run", action="store_true") + args = parser.parse_args() + if args.repeats < 1: + parser.error("--repeats must be positive") + if args.num_frames < 1 or (args.num_frames - 1) % 4: + parser.error("--num-frames must be 1 or 4n+1") + if args.height % 16 or args.width % 16: + parser.error("--height and --width must be multiples of 16") + return args + + +def main() -> None: + args = _parse_args() + upstream_root = args.upstream_root.resolve() + model_dir = args.model_dir.resolve() + modes = args.mode or ["t2i", "t2v", "ti2v"] + cases = _sample_cases(upstream_root, modes, None if args.all_cases else args.case) + run_manifest = { + "schema_version": 1, + "upstream_root": str(upstream_root), + "all_cases": args.all_cases, + "selected_cases": [{"mode": item["mode"], "name": item["name"]} for item in cases], + "model_dir": str(model_dir), + "modes": modes, + "case": args.case, + "dry_run": args.dry_run, + "arguments": vars(args), + } + if args.dry_run: + _write_json(args.output_dir / "capture_manifest.json", run_manifest) + print(json.dumps(run_manifest, indent=2, default=_json_value)) + return + if not model_dir.is_dir(): + raise FileNotFoundError(f"Dense checkpoint is not available: {model_dir}") + os.environ["LINGBOT_QWEN_ATTN_IMPLEMENTATION"] = args.qwen_attn_implementation + os.environ["LINGBOT_QUIET_PROGRESS"] = "1" + runner = _import_upstream(upstream_root) + captures: list[dict[str, Any]] = [] + for case in cases: + for repeat in range(args.repeats): + destination = args.output_dir / case["mode"] / case["name"] / f"run-{repeat:02d}" + destination.mkdir(parents=True, exist_ok=False) + print(f"capturing {case['mode']}/{case['name']} run={repeat}", flush=True) + captures.append( + _capture_once( + runner=runner, + upstream_root=upstream_root, + model_dir=model_dir, + case=case, + destination=destination, + args=args, + repeat=repeat, + ) + ) + _write_json(args.output_dir / "capture_manifest.json", run_manifest) + _repeatability_report(args.output_dir, captures) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/compare_lingbot_video_parity.py b/tools/validation/compare_lingbot_video_parity.py new file mode 100644 index 0000000..cd5f8e4 --- /dev/null +++ b/tools/validation/compare_lingbot_video_parity.py @@ -0,0 +1,95 @@ +"""Compare a TeleFuser replay against a captured LingBot-Video reference. + +The comparator is intentionally framework-agnostic: captures contain tensor +metadata and optional ``.pt`` tensors, so this tool can be used before the +TeleFuser pipeline is wired into the serving stack. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import torch + + +def _load_metadata(root: Path) -> dict[str, Any]: + path = root / "metadata.json" + if not path.is_file(): + raise FileNotFoundError(f"Capture metadata not found: {path}") + return json.loads(path.read_text(encoding="utf-8")) + + +def _tensor_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, Any]: + if reference.shape != candidate.shape: + return { + "shape_match": False, + "reference_shape": list(reference.shape), + "candidate_shape": list(candidate.shape), + } + ref_raw = reference.detach().cpu() + got_raw = candidate.detach().cpu() + ref = ref_raw.float() + got = got_raw.float() + delta = (got - ref).abs() + ref_norm = ref.reshape(-1).norm() + cosine = ( + torch.nn.functional.cosine_similarity(ref.reshape(1, -1), got.reshape(1, -1)).item() if ref.numel() else 1.0 + ) + return { + "shape_match": True, + "dtype_match": reference.dtype == candidate.dtype, + "reference_dtype": str(reference.dtype).removeprefix("torch."), + "candidate_dtype": str(candidate.dtype).removeprefix("torch."), + "exact_mismatch_count": int(torch.count_nonzero(ref_raw != got_raw).item()), + "is_discrete": not reference.is_floating_point() and not candidate.is_floating_point(), + "max_abs_error": float(delta.max().item()) if delta.numel() else 0.0, + "mean_abs_error": float(delta.mean().item()) if delta.numel() else 0.0, + "relative_l2": float(delta.reshape(-1).norm().div(ref_norm.clamp_min(1e-12)).item()), + "cosine_similarity": float(max(-1.0, min(1.0, cosine))), + "nan_count": int(torch.isnan(got).sum().item()), + "inf_count": int(torch.isinf(got).sum().item()), + } + + +def compare_captures(reference_root: Path, candidate_root: Path) -> dict[str, Any]: + """Compare tensors with matching relative paths under two capture roots.""" + reference = _load_metadata(reference_root) + candidate = _load_metadata(candidate_root) + names = sorted(set(reference.get("tensors", {})) & set(candidate.get("tensors", {}))) + tensors: dict[str, Any] = {} + for name in names: + ref_path = reference_root / reference["tensors"][name]["path"] + got_path = candidate_root / candidate["tensors"][name]["path"] + if ref_path.is_file() and got_path.is_file(): + ref_tensor = torch.load(ref_path, map_location="cpu", weights_only=True) + candidate_tensor = torch.load(got_path, map_location="cpu", weights_only=True) + tensors[name] = _tensor_metrics(ref_tensor, candidate_tensor) + return { + "schema_version": 2, + "reference": str(reference_root), + "candidate": str(candidate_root), + "matched_tensors": len(tensors), + "missing_from_candidate": sorted(set(reference.get("tensors", {})) - set(candidate.get("tensors", {}))), + "unexpected_in_candidate": sorted(set(candidate.get("tensors", {})) - set(reference.get("tensors", {}))), + "tensors": tensors, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("reference", type=Path) + parser.add_argument("candidate", type=Path) + parser.add_argument("--output", type=Path) + args = parser.parse_args() + report = compare_captures(args.reference, args.candidate) + payload = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/inspect_lingbot_video_checkpoint.py b/tools/validation/inspect_lingbot_video_checkpoint.py new file mode 100644 index 0000000..79a856a --- /dev/null +++ b/tools/validation/inspect_lingbot_video_checkpoint.py @@ -0,0 +1,177 @@ +"""Report LingBot-Video transformer checkpoint loading acceptance evidence.""" + +from __future__ import annotations + +import argparse +import json +import time +from collections import defaultdict +from pathlib import Path +from typing import Any + +import torch +from safetensors import safe_open + +from telefuser.core.module_manager import ModuleManager + +_DENSE_CONFIG_KEYS = ( + "patch_size", + "in_channels", + "out_channels", + "hidden_size", + "num_attention_heads", + "depth", + "intermediate_size", + "text_dim", + "freq_dim", + "norm_eps", + "rope_theta", + "axes_dims", + "qkv_bias", + "out_bias", + "patch_embed_bias", + "timestep_mlp_bias", +) +_MOE_CONFIG_KEYS = _DENSE_CONFIG_KEYS + ( + "num_experts", + "num_experts_per_tok", + "moe_intermediate_size", + "decoder_sparse_step", + "mlp_only_layers", + "n_group", + "topk_group", + "routed_scaling_factor", + "n_shared_experts", +) + + +def checkpoint_keys(checkpoint_dir: Path) -> set[str]: + """Read checkpoint keys without materializing all checkpoint tensors.""" + index_path = checkpoint_dir / "diffusion_pytorch_model.safetensors.index.json" + if index_path.is_file(): + payload = json.loads(index_path.read_text(encoding="utf-8")) + weight_map = payload.get("weight_map") + if not isinstance(weight_map, dict): + raise ValueError(f"invalid safetensors index: {index_path}") + return set(weight_map) + checkpoint_path = checkpoint_dir / "diffusion_pytorch_model.safetensors" + with safe_open(checkpoint_path, framework="pt", device="cpu") as handle: + return set(handle.keys()) + + +def checkpoint_key_coverage(model: torch.nn.Module, available_keys: set[str]) -> dict[str, Any]: + """Report exact checkpoint key coverage without retaining checkpoint tensors.""" + expected = set(model.state_dict()) + matched = expected & available_keys + return { + "expected_key_count": len(expected), + "checkpoint_key_count": len(available_keys), + "matched_key_count": len(matched), + "coverage": len(matched) / len(expected) if expected else 1.0, + "missing_keys": sorted(expected - available_keys), + "unexpected_keys": sorted(available_keys - expected), + "matched_numel": sum(model.state_dict()[name].numel() for name in matched), + } + + +def _parameter_report(model: torch.nn.Module) -> dict[str, Any]: + """Summarize parameter count, placement, precision, and model structure.""" + by_dtype: dict[str, int] = defaultdict(int) + by_device: dict[str, int] = defaultdict(int) + by_component: dict[str, int] = defaultdict(int) + by_block: dict[str, int] = defaultdict(int) + total = 0 + fp32 = 0 + byte_count = 0 + for name, parameter in model.named_parameters(): + numel = parameter.numel() + total += numel + byte_count += numel * parameter.element_size() + by_dtype[str(parameter.dtype).removeprefix("torch.")] += numel + by_device[str(parameter.device)] += numel + by_component[name.split(".", 1)[0]] += numel + if name.startswith("blocks."): + by_block[name.split(".", 2)[1]] += numel + if parameter.dtype == torch.float32: + fp32 += numel + return { + "total_parameter_count": total, + "estimated_parameter_bytes": byte_count, + "retained_fp32_parameter_count": fp32, + "parameter_count_by_dtype": dict(sorted(by_dtype.items())), + "parameter_count_by_device": dict(sorted(by_device.items())), + "parameter_count_by_component": dict(sorted(by_component.items())), + "parameter_count_by_block": dict(sorted(by_block.items(), key=lambda item: int(item[0]))), + } + + +def build_load_report( + model: torch.nn.Module, + *, + checkpoint_dir: Path, + variant: str, + config: dict[str, Any], + available_keys: set[str], +) -> dict[str, Any]: + """Build the plan-required acceptance report after a strict model load.""" + consumed_keys = _DENSE_CONFIG_KEYS if variant == "dense" else _MOE_CONFIG_KEYS + return { + "variant": variant, + "checkpoint_dir": str(checkpoint_dir), + "consumed_config": {key: config[key] for key in consumed_keys}, + "checkpoint_key_coverage": checkpoint_key_coverage(model, available_keys), + "parameters": _parameter_report(model), + } + + +def _dtype(value: str) -> torch.dtype: + values = {"bf16": torch.bfloat16, "fp16": torch.float16, "fp32": torch.float32} + try: + return values[value] + except KeyError as exc: + raise ValueError(f"unsupported dtype: {value}") from exc + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument( + "--model-dir", type=Path, required=True, help="Checkpoint root containing transformer/ and refiner/." + ) + parser.add_argument("--variant", choices=("dense", "moe", "refiner"), required=True) + parser.add_argument("--device", default="cuda") + parser.add_argument("--dtype", choices=("bf16", "fp16", "fp32"), default="bf16") + parser.add_argument("--output", type=Path) + args = parser.parse_args() + + checkpoint_dir = args.model_dir / ("transformer" if args.variant in {"dense", "moe"} else "refiner") + config = json.loads((checkpoint_dir / "config.json").read_text(encoding="utf-8")) + available_keys = checkpoint_keys(checkpoint_dir) + cuda_before = ( + torch.cuda.memory_allocated(args.device) if torch.cuda.is_available() and "cuda" in args.device else None + ) + started = time.perf_counter() + module_manager = ModuleManager(device=args.device, torch_dtype=_dtype(args.dtype)) + module_manager.load_model(str(checkpoint_dir), name="transformer") + model = module_manager.fetch_module("transformer") + if model is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {checkpoint_dir}") + model.promote_stability_layers_to_fp32() + report = build_load_report( + model, + checkpoint_dir=checkpoint_dir, + variant=args.variant, + config=config, + available_keys=available_keys, + ) + report["load_seconds"] = time.perf_counter() - started + if cuda_before is not None: + report["measured_cuda_allocated_bytes"] = torch.cuda.memory_allocated(args.device) - cuda_before + payload = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/replay_lingbot_video_dense_reference.py b/tools/validation/replay_lingbot_video_dense_reference.py new file mode 100644 index 0000000..5f7c033 --- /dev/null +++ b/tools/validation/replay_lingbot_video_dense_reference.py @@ -0,0 +1,394 @@ +"""Replay a captured Dense LingBot-Video oracle through native TeleFuser stages. + +The upstream capture supplies prompt embeddings, initial latents, and optional +TI2V condition latents. This isolates native DiT, scheduler, condition, and +VAE-decode parity from text-encoder and RNG differences. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from PIL import Image + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video.data import preprocess_ti2v_image +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from telefuser.pipelines.lingbot_video.text_encoding import LingBotVideoTextEncodingStage +from telefuser.pipelines.lingbot_video.vae import ( + LingBotVideoVAEDecodeStage, + LingBotVideoVAEEncodeStage, + denormalize_latent, + first_frame_condition_mask, +) +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler + + +def tensor_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float | int | bool]: + """Return L0/L1 metrics for one captured tensor pair.""" + if reference.shape != candidate.shape: + return { + "shape_match": False, + "reference_numel": reference.numel(), + "candidate_numel": candidate.numel(), + } + ref = reference.float().cpu() + got = candidate.float().cpu() + delta = (got - ref).abs() + return { + "shape_match": True, + "max_abs": float(delta.max().item()) if delta.numel() else 0.0, + "mean_abs": float(delta.mean().item()) if delta.numel() else 0.0, + "relative_l2": float(delta.norm().div(ref.norm().clamp_min(1e-12)).item()), + "cosine": float(torch.nn.functional.cosine_similarity(ref.flatten(), got.flatten(), dim=0).item()), + "exact_mismatch_count": int(torch.count_nonzero(reference.cpu() != candidate.cpu()).item()), + } + + +def _load_tensor(reference_dir: Path, metadata: dict[str, Any], name: str) -> torch.Tensor: + entry = metadata["tensors"].get(name) + if entry is None: + raise KeyError(f"capture does not contain required tensor: {name}") + return torch.load(reference_dir / entry["path"], map_location="cpu", weights_only=True) + + +def _record( + results: dict[str, dict[str, float | int | bool]], name: str, reference: torch.Tensor, candidate: torch.Tensor +) -> None: + results[name] = tensor_metrics(reference, candidate) + + +def discover_reference_dirs(reference_root: Path) -> list[Path]: + """Return captured-run directories under a reference artifact root.""" + references = sorted(path.parent for path in reference_root.rglob("metadata.json")) + if not references: + raise FileNotFoundError(f"No capture metadata found under {reference_root}") + return references + + +def exact_replay_failures(report: dict[str, Any]) -> list[str]: + """Return metric paths that do not meet the exact numerical-oracle gate.""" + reports = report.get("reports", [report]) + failures: list[str] = [] + for item_index, item in enumerate(reports): + metrics = item.get("metrics", {}) + for name, value in metrics.items(): + if not value.get("shape_match", False) or value.get("exact_mismatch_count") != 0: + failures.append(f"report[{item_index}].{name}") + return failures + + +class _RecordingProcessor: + """Proxy a Hugging Face processor while retaining its pre-device input tensors.""" + + def __init__(self, processor: Any) -> None: + self._processor = processor + self.calls: list[dict[str, torch.Tensor]] = [] + + def __call__(self, *args: Any, **kwargs: Any) -> Any: + inputs = self._processor(*args, **kwargs) + self.calls.append( + { + name: value.detach().cpu().clone(memory_format=torch.contiguous_format) + for name in ("input_ids", "attention_mask", "pixel_values", "image_grid_thw") + if (value := inputs.get(name)) is not None + } + ) + return inputs + + def __getattr__(self, name: str) -> Any: + return getattr(self._processor, name) + + +def replay_text( + reference_dir: Path, + model_dir: Path, + *, + validate_ti2v_vae: bool = False, + seed: int = 42, +) -> dict[str, Any]: + """Replay source prompt construction and Qwen3-VL encoding for one capture.""" + try: + from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + except ImportError as exc: + raise RuntimeError("Text replay requires transformers with Qwen3-VL support") from exc + + metadata = json.loads((reference_dir / "metadata.json").read_text(encoding="utf-8")) + device = torch.device("cuda") + runtime_config = ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) + processor = AutoProcessor.from_pretrained(model_dir / "processor") + text_encoder = ( + Qwen3VLForConditionalGeneration.from_pretrained( + model_dir / "text_encoder", dtype=torch.bfloat16, attn_implementation="sdpa" + ) + .to(device) + .eval() + ) + module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) + module_manager.add_module(text_encoder, name="text_encoder") + module_manager.add_module(processor, name="processor") + stage = LingBotVideoTextEncodingStage("text_encoder", module_manager, runtime_config) + stage._crop_prefix_length() + recording_processor = _RecordingProcessor(processor) + stage.processor = recording_processor + + images: list[Image.Image] | None = None + results: dict[str, dict[str, float | int | bool]] = {} + if metadata["mode"] == "ti2v": + image_path = Path(metadata["image"]) + raw_image = ( + torch.from_numpy(np.array(Image.open(image_path).convert("RGB"), copy=True)).permute(2, 0, 1).unsqueeze(0) + ) + sampling = metadata["sampling"] + condition_pixels = preprocess_ti2v_image(raw_image, height=sampling["height"], width=sampling["width"]) + _record( + results, + "ti2v_preprocessed_image", + _load_tensor(reference_dir, metadata, "ti2v_preprocessed_image"), + condition_pixels, + ) + images = [stage.prepare_ti2v_vlm_image(condition_pixels)] + if validate_ti2v_vae: + try: + from diffusers import AutoencoderKLWan + except ImportError as exc: + raise RuntimeError("TI2V VAE replay requires diffusers AutoencoderKLWan") from exc + vae = AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).to(device).eval() + module_manager.add_module(vae, name="vae") + vae_encode = LingBotVideoVAEEncodeStage("vae_encode", module_manager, runtime_config) + generator = torch.Generator(device=device).manual_seed(metadata.get("seed", seed)) + _record( + results, + "ti2v_condition_latent", + _load_tensor(reference_dir, metadata, "ti2v_condition_latent"), + vae_encode.encode(condition_pixels, generator=generator), + ) + + for call, prompt in enumerate((metadata["prompt"], metadata["negative_prompt"])): + embeddings, mask = stage.encode(prompt, images=images) + _record( + results, f"prompt_{call}_embeds", _load_tensor(reference_dir, metadata, f"prompt_{call}_embeds"), embeddings + ) + _record(results, f"prompt_{call}_mask", _load_tensor(reference_dir, metadata, f"prompt_{call}_mask"), mask) + inputs = recording_processor.calls[call] + for name, candidate in inputs.items(): + reference_name = f"prompt_inputs_{call + 1}_{name}" + if reference_name in metadata["tensors"]: + _record(results, reference_name, _load_tensor(reference_dir, metadata, reference_name), candidate) + return { + "reference_dir": str(reference_dir), + "mode": metadata["mode"], + "case": metadata["case"], + "metrics": results, + } + + +def _load_replay_runtime( + model_dir: Path, +) -> tuple[ModelRuntimeConfig, LingBotVideoDenoisingStage, Any]: + """Load the Dense DiT and VAE once for one or more capture replays.""" + try: + from diffusers import AutoencoderKLWan + except ImportError as exc: + raise RuntimeError("Dense replay requires diffusers AutoencoderKLWan") from exc + device = torch.device("cuda") + runtime_config = ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) + module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) + module_manager.load_model(str(model_dir / "transformer"), name="transformer") + transformer = module_manager.fetch_module("transformer") + if transformer is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {model_dir / 'transformer'}") + transformer.promote_stability_layers_to_fp32() + denoising = LingBotVideoDenoisingStage("transformer", module_manager, runtime_config) + vae = AutoencoderKLWan.from_pretrained(model_dir / "vae", torch_dtype=torch.float32).to(device).eval() + return runtime_config, denoising, vae + + +def _replay_with_runtime( + reference_dir: Path, + model_dir: Path, + *, + runtime_config: ModelRuntimeConfig, + denoising: LingBotVideoDenoisingStage, + vae: Any, +) -> dict[str, Any]: + """Replay one capture using already-loaded Dense and VAE components.""" + metadata = json.loads((reference_dir / "metadata.json").read_text(encoding="utf-8")) + sampling = metadata["sampling"] + device = torch.device("cuda") + scheduler = FlowUniPCMultistepScheduler.from_pretrained(model_dir / "scheduler") + scheduler.set_timesteps(sampling["num_inference_steps"], device=device, shift=sampling["shift"]) + results: dict[str, dict[str, float | int | bool]] = {} + _record( + results, + "scheduler_timesteps", + _load_tensor(reference_dir, metadata, "scheduler_timesteps"), + scheduler.timesteps, + ) + _record(results, "scheduler_sigmas", _load_tensor(reference_dir, metadata, "scheduler_sigmas"), scheduler.sigmas) + + current = _load_tensor(reference_dir, metadata, "initial_latents").to(device) + positive = _load_tensor(reference_dir, metadata, "prompt_0_embeds").to(device) + positive_mask = _load_tensor(reference_dir, metadata, "prompt_0_mask").to(device) + negative = _load_tensor(reference_dir, metadata, "prompt_1_embeds").to(device) + negative_mask = _load_tensor(reference_dir, metadata, "prompt_1_mask").to(device) + condition = None + condition_mask = None + if metadata["mode"] == "ti2v": + condition = _load_tensor(reference_dir, metadata, "ti2v_condition_latent").to(device) + condition_mask = first_frame_condition_mask(current.shape[2], device=device) + + for index, timestep in enumerate(scheduler.timesteps): + if condition is not None and condition_mask is not None: + current = reinject_ti2v_condition(current, condition, condition_mask) + timestep_batch = timestep.expand(current.shape[0]) + noise = denoising.predict_noise_with_cfg( + current, + timestep_batch, + positive, + negative, + positive_mask, + negative_mask, + float(sampling["guidance_scale"]), + ) + prefix = f"step_{index}" + if f"{prefix}_noise_prediction" in metadata["tensors"]: + _record( + results, + f"{prefix}_noise_prediction", + _load_tensor(reference_dir, metadata, f"{prefix}_noise_prediction"), + noise, + ) + current = scheduler.step(noise, timestep, current) + if f"{prefix}_latent_after" in metadata["tensors"]: + _record( + results, + f"{prefix}_latent_after", + _load_tensor(reference_dir, metadata, f"{prefix}_latent_after"), + current, + ) + if condition is not None and condition_mask is not None: + current = reinject_ti2v_condition(current, condition, condition_mask) + + raw_latent = denormalize_latent( + current, + torch.tensor(vae.config.latents_mean, device=device), + torch.tensor(vae.config.latents_std, device=device), + ) + _record(results, "vae_decode_input", _load_tensor(reference_dir, metadata, "vae_decode_input"), raw_latent) + vae_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) + vae_manager.add_module(vae, name="vae") + decoder = LingBotVideoVAEDecodeStage("vae_decode", vae_manager, runtime_config) + frames = decoder.decode(current) + captured_frames = torch.from_numpy(np.load(reference_dir / metadata["frames"]["path"])) + candidate_frames = frames.permute(0, 2, 3, 4, 1)[0].cpu() + _record(results, "decoded_frames", captured_frames, candidate_frames) + return { + "reference_dir": str(reference_dir), + "mode": metadata["mode"], + "case": metadata["case"], + "metrics": results, + } + + +def replay(reference_dir: Path, model_dir: Path) -> dict[str, Any]: + """Replay a one-GPU Dense capture and return its native stage parity report.""" + runtime_config, denoising, vae = _load_replay_runtime(model_dir) + return _replay_with_runtime( + reference_dir, + model_dir, + runtime_config=runtime_config, + denoising=denoising, + vae=vae, + ) + + +def replay_many(reference_dirs: list[Path], model_dir: Path) -> dict[str, Any]: + """Replay multiple Dense captures while retaining one DiT/VAE runtime.""" + if not reference_dirs: + raise ValueError("reference_dirs must not be empty") + runtime_config, denoising, vae = _load_replay_runtime(model_dir) + reports = [ + _replay_with_runtime( + reference_dir, + model_dir, + runtime_config=runtime_config, + denoising=denoising, + vae=vae, + ) + for reference_dir in reference_dirs + ] + return { + "reference_count": len(reports), + "reference_dirs": [str(reference_dir) for reference_dir in reference_dirs], + "reports": reports, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + reference_group = parser.add_mutually_exclusive_group(required=True) + reference_group.add_argument("--reference-dir", type=Path) + reference_group.add_argument( + "--reference-root", + type=Path, + help="Replay every captured run below this artifact root while retaining one Dense/VAE runtime.", + ) + parser.add_argument( + "--model-dir", type=Path, default=Path("/hhb-data/aigc/model_zoo/lingbot/lingbot-video-dense-1.3b") + ) + parser.add_argument("--output", type=Path) + parser.add_argument( + "--validate-text", + action="store_true", + help="Validate source prompt construction and Qwen3-VL embeddings instead of DiT/VAE replay.", + ) + parser.add_argument( + "--validate-ti2v-vae", + action="store_true", + help="Also compare the sampled TI2V VAE condition latent; requires the capture seed.", + ) + parser.add_argument( + "--assert-exact", + action="store_true", + help="Exit nonzero unless every recorded tensor has an exact shape and value match.", + ) + parser.add_argument("--seed", type=int, default=42, help="Fallback seed for captures created before seed metadata.") + args = parser.parse_args() + reference_dirs = ( + [args.reference_dir] if args.reference_dir is not None else discover_reference_dirs(args.reference_root) + ) + if args.validate_text: + reports = [ + replay_text( + reference_dir, + args.model_dir, + validate_ti2v_vae=args.validate_ti2v_vae, + seed=args.seed, + ) + for reference_dir in reference_dirs + ] + report: dict[str, Any] = reports[0] if len(reports) == 1 else {"reports": reports} + elif len(reference_dirs) == 1: + report = replay(reference_dirs[0], args.model_dir) + else: + report = replay_many(reference_dirs, args.model_dir) + if args.assert_exact: + failures = exact_replay_failures(report) + if failures: + raise SystemExit(f"Exact replay gate failed: {', '.join(failures)}") + payload = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_dense_parity.py b/tools/validation/run_lingbot_video_dense_parity.py new file mode 100644 index 0000000..5778575 --- /dev/null +++ b/tools/validation/run_lingbot_video_dense_parity.py @@ -0,0 +1,85 @@ +"""Compare an upstream Dense DiT forward with TeleFuser on one GPU. + +Set ``PYTHONPATH=work_dirs/lingbot-video-master`` so the upstream package is +available without modifying its checkout. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from safetensors.torch import load_model + +from telefuser.core.module_manager import ModuleManager + + +def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: + delta = reference.float() - candidate.float() + return { + "max_abs": float(delta.abs().max().item()), + "mean_abs": float(delta.abs().mean().item()), + "relative_l2": float((delta.norm() / reference.float().norm().clamp_min(1e-12)).item()), + "cosine": float( + torch.nn.functional.cosine_similarity( + reference.float().flatten(), candidate.float().flatten(), dim=0 + ).item() + ), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--transformer-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--height", type=int, default=8) + parser.add_argument("--width", type=int, default=8) + parser.add_argument("--text-length", type=int, default=3) + args = parser.parse_args() + try: + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel as UpstreamTransformer + except ImportError as exc: + raise RuntimeError("Set PYTHONPATH to the upstream LingBot-Video checkout") from exc + + config = json.loads((args.transformer_dir / "config.json").read_text(encoding="utf-8")) + fields = ( + "patch_size", + "in_channels", + "out_channels", + "hidden_size", + "num_attention_heads", + "depth", + "intermediate_size", + "text_dim", + "freq_dim", + "norm_eps", + "rope_theta", + "axes_dims", + "qkv_bias", + "out_bias", + "patch_embed_bias", + "timestep_mlp_bias", + ) + device = torch.device("cuda") + upstream = UpstreamTransformer(**{field: config[field] for field in fields}).to(device, torch.bfloat16).eval() + load_model(upstream, args.transformer_dir / "diffusion_pytorch_model.safetensors", strict=True, device=str(device)) + module_manager = ModuleManager(device=str(device), torch_dtype=torch.bfloat16) + module_manager.load_model(str(args.transformer_dir), name="transformer") + telefuser = module_manager.fetch_module("transformer") + if telefuser is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {args.transformer_dir}") + telefuser.promote_stability_layers_to_fp32() + torch.manual_seed(args.seed) + latent = torch.randn(1, 16, 1, args.height, args.width, device=device, dtype=torch.bfloat16) + text = torch.randn(1, args.text_length, config["text_dim"], device=device, dtype=torch.bfloat16) + timestep = torch.tensor([500], device=device) + with torch.no_grad(): + reference = upstream(latent, timestep, text, return_dict=False)[0] + candidate = telefuser(latent, timestep, text) + print(json.dumps(_metrics(reference, candidate), indent=2, sort_keys=True)) + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_distributed.py b/tools/validation/run_lingbot_video_distributed.py new file mode 100644 index 0000000..0a5b2bf --- /dev/null +++ b/tools/validation/run_lingbot_video_distributed.py @@ -0,0 +1,130 @@ +"""Run one source-equivalent LingBot-Video request with FSDP and Ulysses SP.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +import time +from pathlib import Path + +import torch +import torch.distributed as dist + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline as build_dense_pipeline +from examples.lingbot_video.lingbot_video_moe_30b import build_pipeline as build_moe_pipeline +from telefuser.core.config import ParallelConfig +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + load_lingbot_video_prompt, +) + + +def _write_video(frames: torch.Tensor, output: Path, fps: int) -> None: + """Encode normalized TeleFuser RGB frames without a second uint8 conversion.""" + from diffusers.utils import export_to_video + + output.parent.mkdir(parents=True, exist_ok=True) + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + export_to_video(list(video), str(output), fps=fps) + + +def _max_across_ranks(value: float, device: torch.device) -> float: + """Return a scalar maximum across the default NCCL group.""" + tensor = torch.tensor(value, device=device) + dist.all_reduce(tensor, op=dist.ReduceOp.MAX) + return float(tensor.item()) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--caption-json", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--variant", choices=("dense", "moe"), default="dense") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--steps", type=int, default=40) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--fps", type=int, default=24) + args = parser.parse_args() + + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + try: + world_size = dist.get_world_size() + if world_size < 2: + raise ValueError("distributed LingBot validation requires at least two ranks") + parallel_config = ParallelConfig( + device_ids=list(range(world_size)), + sp_ulysses_degree=world_size, + enable_fsdp=True, + ) + caption, _ = load_lingbot_video_prompt(args.caption_json) + torch.cuda.reset_peak_memory_stats(device) + pipeline_builder = build_dense_pipeline if args.variant == "dense" else build_moe_pipeline + pipeline = pipeline_builder( + args.model_dir, + cpu_offload=False, + guidance_scale=args.guidance_scale, + num_inference_steps=args.steps, + parallel_config=parallel_config, + ) + dist.barrier() + started = time.perf_counter() + output = pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=args.height, + width=args.width, + num_frames=args.num_frames, + ), + generator=torch.Generator(device=device).manual_seed(args.seed), + ).output + elapsed_seconds = _max_across_ranks(time.perf_counter() - started, device) + reference = output.clone() if dist.get_rank() == 0 else torch.empty_like(output) + dist.broadcast(reference, src=0) + max_abs = _max_across_ranks(float((output.float() - reference.float()).abs().max().item()), device) + peak_memory_mib = _max_across_ranks(torch.cuda.max_memory_allocated(device) / 1024**2, device) + if dist.get_rank() == 0: + _write_video(output, args.output, args.fps) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text( + json.dumps( + { + "variant": args.variant, + "world_size": world_size, + "sp_ulysses_degree": world_size, + "fsdp": True, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "steps": args.steps, + "rank_output_max_abs": max_abs, + "elapsed_seconds": elapsed_seconds, + "peak_gpu_memory_mib": peak_memory_mib, + "output": str(args.output), + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + dist.barrier() + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_moe_parity.py b/tools/validation/run_lingbot_video_moe_parity.py new file mode 100644 index 0000000..f4bd0fd --- /dev/null +++ b/tools/validation/run_lingbot_video_moe_parity.py @@ -0,0 +1,106 @@ +"""Compare upstream and TeleFuser MoE/refiner forwards without dual model residency. + +Set ``PYTHONPATH=work_dirs/lingbot-video-master`` so the upstream package is +available without modifying its checkout. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from safetensors.torch import load_file + +from telefuser.core.module_manager import ModuleManager + + +def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: + delta = reference.float() - candidate.float() + return { + "max_abs": float(delta.abs().max().item()), + "mean_abs": float(delta.abs().mean().item()), + "relative_l2": float((delta.norm() / reference.float().norm().clamp_min(1e-12)).item()), + "cosine": float( + torch.nn.functional.cosine_similarity( + reference.float().flatten(), candidate.float().flatten(), dim=0 + ).item() + ), + } + + +def exact_parity_failures(metrics: dict[str, float]) -> list[str]: + """Return numerical fields that do not meet the zero-drift oracle gate.""" + return [name for name in ("max_abs", "mean_abs", "relative_l2") if metrics[name] != 0.0] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--transformer-dir", type=Path, required=True) + parser.add_argument("--seed", type=int, default=7) + parser.add_argument("--height", type=int, default=8) + parser.add_argument("--width", type=int, default=8) + parser.add_argument("--text-length", type=int, default=3) + parser.add_argument("--output", type=Path, help="Optional JSON metrics destination.") + parser.add_argument("--assert-exact", action="store_true", help="Exit nonzero unless the oracle replay is exact.") + args = parser.parse_args() + try: + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel as UpstreamTransformer + except ImportError as exc: + raise RuntimeError("Set PYTHONPATH to the upstream LingBot-Video checkout") from exc + + directory = args.transformer_dir + config = json.loads((directory / "config.json").read_text(encoding="utf-8")) + config_fields = {name: value for name, value in config.items() if not name.startswith("_")} + device = torch.device("cuda") + torch.manual_seed(args.seed) + latent = torch.randn(1, 16, 1, args.height, args.width, dtype=torch.bfloat16) + text = torch.randn(1, args.text_length, config["text_dim"], dtype=torch.bfloat16) + timestep = torch.tensor([500]) + + upstream = UpstreamTransformer(**config_fields).to(torch.bfloat16).to(device).eval() + index = json.loads((directory / "diffusion_pytorch_model.safetensors.index.json").read_text(encoding="utf-8"))[ + "weight_map" + ] + expected = set(upstream.state_dict()) + found: set[str] = set() + unexpected: set[str] = set() + for shard in sorted(set(index.values())): + shard_state = load_file(directory / shard, device=str(device)) + unexpected.update(upstream.load_state_dict(shard_state, strict=False).unexpected_keys) + found.update(shard_state) + del shard_state + missing = expected - found + if missing or unexpected: + raise RuntimeError( + f"Upstream MoE checkpoint mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}" + ) + with torch.no_grad(): + reference = ( + upstream(latent.to(device), timestep.to(device), text.to(device), return_dict=False)[0].float().cpu() + ) + del upstream + torch.cuda.empty_cache() + + module_manager = ModuleManager(device=str(device), torch_dtype=torch.bfloat16) + module_manager.load_model(str(directory), name="transformer") + telefuser = module_manager.fetch_module("transformer") + if telefuser is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {directory}") + telefuser.promote_stability_layers_to_fp32() + with torch.no_grad(): + candidate = telefuser(latent.to(device), timestep.to(device), text.to(device)).float().cpu() + metrics = _metrics(reference, candidate) + if args.assert_exact: + failures = exact_parity_failures(metrics) + if failures: + raise SystemExit(f"Exact MoE parity gate failed: {', '.join(failures)}") + payload = json.dumps(metrics, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_native_parallel.py b/tools/validation/run_lingbot_video_native_parallel.py new file mode 100644 index 0000000..31eb6c4 --- /dev/null +++ b/tools/validation/run_lingbot_video_native_parallel.py @@ -0,0 +1,167 @@ +"""Run LingBot-Video through the native TeleFuser multi-process worker topology.""" + +from __future__ import annotations + +import argparse +import json +import sys +import time +from pathlib import Path + +import torch + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline as build_dense_pipeline +from examples.lingbot_video.lingbot_video_moe_30b import build_pipeline as build_moe_pipeline +from examples.lingbot_video.lingbot_video_moe_30b import build_refiner +from telefuser.core.config import ParallelConfig +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + load_lingbot_video_prompt, + prepare_refiner_video, +) + + +def _write_video(frames: torch.Tensor, output: Path, fps: int) -> None: + """Encode normalized RGB frames from the rank-zero worker result.""" + from diffusers.utils import export_to_video + + output.parent.mkdir(parents=True, exist_ok=True) + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + export_to_video(list(video), str(output), fps=fps) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--caption-json", type=Path, required=True) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument("--report", type=Path, required=True) + parser.add_argument("--variant", choices=("dense", "moe"), default="dense") + parser.add_argument("--refine", action="store_true") + parser.add_argument("--refiner-height", type=int, default=1088) + parser.add_argument("--refiner-width", type=int, default=1920) + parser.add_argument("--refiner-steps", type=int, default=8) + parser.add_argument("--refiner-guidance-scale", type=float, default=3.0) + parser.add_argument("--refiner-shift", type=float, default=3.0) + parser.add_argument("--refiner-t-thresh", type=float, default=0.85) + parser.add_argument("--refiner-tail-steps", type=int, default=2) + parser.add_argument( + "--refiner-batch-cfg", + action="store_true", + help="Use batched CFG for refiner; disabled by default so 1088p SP4 fits on four H100s.", + ) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=121) + parser.add_argument("--steps", type=int, default=40) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--fps", type=int, default=24) + args = parser.parse_args() + + if args.refine and args.variant != "moe": + raise ValueError("--refine requires --variant moe") + + if torch.cuda.device_count() < 4: + raise RuntimeError("native LingBot SP4 validation requires four visible CUDA devices") + caption, _ = load_lingbot_video_prompt(args.caption_json) + pipeline_builder = build_dense_pipeline if args.variant == "dense" else build_moe_pipeline + pipeline = pipeline_builder( + args.model_dir, + cpu_offload=False, + guidance_scale=args.guidance_scale, + num_inference_steps=args.steps, + parallel_config=ParallelConfig( + device_ids=[0, 1, 2, 3], + sp_ulysses_degree=4, + enable_fsdp=True, + ), + ) + try: + torch.cuda.reset_peak_memory_stats(0) + started = time.perf_counter() + generation = pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=args.height, + width=args.width, + num_frames=args.num_frames, + ), + generator=torch.Generator("cuda").manual_seed(args.seed), + ) + output = generation.output + base_elapsed_seconds = time.perf_counter() - started + refiner_elapsed_seconds = None + if args.refine: + pipeline.release_gpu_resources() + refiner = build_refiner( + args.model_dir, + cpu_offload=False, + parallel_config=ParallelConfig( + device_ids=[0, 1, 2, 3], + sp_ulysses_degree=4, + enable_fsdp=True, + ), + batch_cfg=args.refiner_batch_cfg, + ) + try: + lowres_video, _ = prepare_refiner_video( + output, + source_fps=args.fps, + height=args.refiner_height, + width=args.refiner_width, + ) + refiner_started = time.perf_counter() + output = refiner.refine( + lowres_video, + generation.prompt_conditions.positive_prompt_embeds, + torch.zeros_like(generation.prompt_conditions.positive_prompt_embeds), + generation.prompt_conditions.positive_attention_mask, + generation.prompt_conditions.positive_attention_mask.clone(), + num_inference_steps=args.refiner_steps, + guidance_scale=args.refiner_guidance_scale, + shift=args.refiner_shift, + t_thresh=args.refiner_t_thresh, + tail_steps=args.refiner_tail_steps, + generator=torch.Generator("cuda").manual_seed(args.seed), + ) + refiner_elapsed_seconds = time.perf_counter() - refiner_started + finally: + refiner.close() + _write_video(output, args.output, args.fps) + args.report.parent.mkdir(parents=True, exist_ok=True) + args.report.write_text( + json.dumps( + { + "world_size": 4, + "variant": args.variant, + "refine": args.refine, + "sp_ulysses_degree": 4, + "fsdp": True, + "batch_cfg": True, + "refiner_batch_cfg": args.refiner_batch_cfg if args.refine else None, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "steps": args.steps, + "base_elapsed_seconds": base_elapsed_seconds, + "refiner_elapsed_seconds": refiner_elapsed_seconds, + "peak_rank0_gpu_memory_mib": torch.cuda.max_memory_allocated(0) / 1024**2, + "output": str(args.output), + }, + indent=2, + sort_keys=True, + ) + + "\n", + encoding="utf-8", + ) + finally: + pipeline.stop() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_refiner_core_parity.py b/tools/validation/run_lingbot_video_refiner_core_parity.py new file mode 100644 index 0000000..d7b2269 --- /dev/null +++ b/tools/validation/run_lingbot_video_refiner_core_parity.py @@ -0,0 +1,177 @@ +"""Compare upstream and TeleFuser refiner low-noise sampling with injected tensors.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import torch +from safetensors.torch import load_file + +from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage +from telefuser.pipelines.lingbot_video.refiner import ( + LingBotVideoRefinerStage, + compute_refiner_sigmas, + prepare_refiner_latent, +) +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler + + +class _Encode: + def encode(self, values: torch.Tensor, generator: torch.Generator | None = None) -> torch.Tensor: + del generator + return values + + +class _Decode: + def decode(self, values: torch.Tensor) -> torch.Tensor: + return values + + +def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: + delta = reference.float() - candidate.float() + return { + "max_abs": float(delta.abs().max()), + "mean_abs": float(delta.abs().mean()), + "relative_l2": float(delta.norm() / reference.float().norm().clamp_min(1e-12)), + "cosine": float( + torch.nn.functional.cosine_similarity(reference.float().flatten(), candidate.float().flatten(), dim=0) + ), + } + + +def exact_parity_failures(metrics: dict[str, float]) -> list[str]: + """Return numerical fields that do not meet the zero-drift oracle gate.""" + return [name for name in ("max_abs", "mean_abs", "relative_l2") if metrics[name] != 0.0] + + +def _load_upstream(): + from lingbot_video.transformer_lingbot_video import LingBotVideoTransformer3DModel + + config = json.loads((DIRECTORY / "config.json").read_text()) + model = LingBotVideoTransformer3DModel(**{key: value for key, value in config.items() if not key.startswith("_")}) + model = model.to(torch.bfloat16).to(DEVICE).eval() + index = json.loads((DIRECTORY / "diffusion_pytorch_model.safetensors.index.json").read_text())["weight_map"] + expected, found, unexpected = set(model.state_dict()), set(), set() + for shard in sorted(set(index.values())): + values = load_file(DIRECTORY / shard, device=str(DEVICE)) + unexpected.update(model.load_state_dict(values, strict=False).unexpected_keys) + found.update(values) + del values + if expected - found or unexpected: + raise RuntimeError("upstream checkpoint key mismatch") + return model + + +def main() -> None: + global ROOT, DIRECTORY, DEVICE + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-root", type=Path, required=True) + parser.add_argument("--refiner-dir", type=Path) + parser.add_argument("--device", default="cuda") + parser.add_argument("--seed", type=int, default=13) + parser.add_argument("--output", type=Path) + parser.add_argument("--assert-exact", action="store_true", help="Exit nonzero unless the oracle replay is exact.") + args = parser.parse_args() + ROOT = args.model_root + DIRECTORY = args.refiner_dir or ROOT / "refiner" + DEVICE = torch.device(args.device) + from lingbot_video.pipeline_lingbot_video import LingBotVideoPipeline + from lingbot_video.scheduling_flow_unipc import FlowUniPCMultistepScheduler as UpstreamScheduler + + torch.manual_seed(args.seed) + x_up = torch.randn(1, 16, 1, 8, 8) + noise = torch.randn_like(x_up) + condition = torch.randn(1, 16, 1, 8, 8) + positive = torch.randn(1, 3, 2560, dtype=torch.bfloat16) + negative = torch.randn(1, 3, 2560, dtype=torch.bfloat16) + mask = torch.ones(1, 3, dtype=torch.long) + sigmas = compute_refiner_sigmas( + sigma_max=1.0, sigma_min=0.0, num_inference_steps=2, shift=3.0, t_thresh=0.25, tail_steps=1 + ) + initial = prepare_refiner_latent(x_up, noise, 0.25) + + upstream = _load_upstream() + upstream_scheduler = UpstreamScheduler.from_pretrained(ROOT / "scheduler") + source_pipeline = LingBotVideoPipeline(upstream, None, None, None, upstream_scheduler) + with torch.no_grad(): + reference = ( + source_pipeline( + "", + height=64, + width=64, + num_frames=1, + guidance_scale=3.0, + num_inference_steps=2, + shift=3.0, + latents=initial.to(DEVICE), + cond_latent=condition.to(DEVICE), + prompt_embeds=positive.to(DEVICE), + prompt_mask=mask.to(DEVICE), + negative_prompt_embeds=negative.to(DEVICE), + negative_prompt_mask=mask.to(DEVICE), + output_type="latent", + t_thresh=0.25, + refiner_sigma_tail_steps=1, + ) + .frames.float() + .cpu() + ) + upstream.to("cpu") + source_pipeline.transformer = None + del source_pipeline, upstream + import gc + + gc.collect() + torch.cuda.empty_cache() + + module_manager = ModuleManager(device=str(DEVICE), torch_dtype=torch.bfloat16) + module_manager.load_model(str(DIRECTORY), name="transformer") + local = module_manager.fetch_module("transformer") + if local is None: + raise RuntimeError(f"Unable to load LingBot-Video transformer from {DIRECTORY}") + local.promote_stability_layers_to_fp32() + local_scheduler = FlowUniPCMultistepScheduler.from_pretrained(ROOT / "scheduler") + stage = LingBotVideoRefinerStage( + denoising_stage=LingBotVideoDenoisingStage( + "refiner", module_manager, ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) + ), + vae_encode_stage=_Encode(), + vae_decode_stage=_Decode(), + scheduler=local_scheduler, + ) + with torch.no_grad(): + candidate = ( + stage.refine( + x_up.to(DEVICE), + positive.to(DEVICE), + negative.to(DEVICE), + mask.to(DEVICE), + mask.to(DEVICE), + num_inference_steps=2, + guidance_scale=3.0, + shift=3.0, + t_thresh=0.25, + tail_steps=1, + clean_first_frame=condition.to(DEVICE), + noise=noise.to(DEVICE), + ) + .float() + .cpu() + ) + metrics = _metrics(reference, candidate) + if args.assert_exact: + failures = exact_parity_failures(metrics) + if failures: + raise SystemExit(f"Exact refiner parity gate failed: {', '.join(failures)}") + payload = json.dumps({"sigmas": sigmas.tolist(), "metrics": metrics}, indent=2, sort_keys=True) + "\n" + if args.output: + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/run_lingbot_video_sp_parity.py b/tools/validation/run_lingbot_video_sp_parity.py new file mode 100644 index 0000000..efb3a9a --- /dev/null +++ b/tools/validation/run_lingbot_video_sp_parity.py @@ -0,0 +1,129 @@ +"""Validate LingBot-Video Ulysses sequence parallelism against an eager model.""" + +from __future__ import annotations + +import argparse +import json +import os + +import torch +import torch.distributed as dist + +from telefuser.core.config import ModelRuntimeConfig, OffloadConfig, ParallelConfig, WeightOffloadType +from telefuser.core.module_manager import ModuleManager +from telefuser.models.lingbot_video_dit import LingBotVideoTransformer3DModel +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, transformer_timestep + + +def _build_model(device: torch.device) -> LingBotVideoTransformer3DModel: + """Create a compact deterministic DiT whose joint stream needs SP padding.""" + torch.manual_seed(7) + return ( + LingBotVideoTransformer3DModel( + patch_size=(1, 2, 2), + in_channels=2, + out_channels=2, + hidden_size=16, + num_attention_heads=4, + depth=2, + intermediate_size=32, + text_dim=8, + freq_dim=8, + axes_dims=(2, 2, 0), + ) + .to(device=device, dtype=torch.bfloat16) + .eval() + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--output") + parser.add_argument("--fsdp", action="store_true", help="Also apply TeleFuser block FSDP.") + args = parser.parse_args() + local_rank = int(os.environ["LOCAL_RANK"]) + torch.cuda.set_device(local_rank) + device = torch.device("cuda", local_rank) + dist.init_process_group("nccl") + try: + world_size = dist.get_world_size() + if world_size < 2 or 4 % world_size: + raise ValueError("SP parity requires a world size that divides four") + model = _build_model(device) + # LingBot checkpoints intentionally preserve these modulation tables in FP32. + # Exercise the FSDP ignored-state path with the same mixed-dtype layout. + for module in model.modules(): + if hasattr(module, "scale_shift_table"): + module.scale_shift_table.data = module.scale_shift_table.data.float() + torch.manual_seed(11) + latent = torch.randn(1, 2, 1, 2, 8, device=device, dtype=torch.bfloat16) + prompt = torch.randn(1, 3, 8, device=device, dtype=torch.bfloat16) + negative_prompt = torch.randn(1, 3, 8, device=device, dtype=torch.bfloat16) + prompt_mask = torch.tensor([[True, True, False]], device=device) + timestep = torch.tensor([125.0], device=device) + with torch.inference_mode(): + model_timestep = transformer_timestep(timestep, torch.bfloat16) + with torch.autocast(device_type="cuda", dtype=torch.bfloat16): + eager_positive = model(latent, model_timestep, prompt, prompt_mask) + eager_negative = model(latent, model_timestep, negative_prompt, prompt_mask) + eager = eager_negative.float() + 3.0 * (eager_positive.float() - eager_negative.float()) + if args.fsdp: + runtime_config = ModelRuntimeConfig( + device_type="cuda", + device_id=local_rank, + torch_dtype=torch.bfloat16, + offload_config=OffloadConfig(offload_type=WeightOffloadType.NO_CPU_OFFLOAD), + parallel_config=ParallelConfig( + device_ids=list(range(world_size)), + sp_ulysses_degree=world_size, + enable_fsdp=True, + ), + ) + module_manager = ModuleManager(device="cpu", torch_dtype=torch.bfloat16) + module_manager.add_module(model, name="transformer") + stage = LingBotVideoDenoisingStage("sp4-test", module_manager, runtime_config, batch_cfg=True) + stage.parallel_models() + distributed = stage.predict_noise_with_cfg( + latent, + timestep, + prompt, + negative_prompt, + positive_attention_mask=prompt_mask, + negative_attention_mask=prompt_mask, + guidance_scale=3.0, + ) + else: + model.set_ulysses_group(dist.group.WORLD) + positive = model(latent, model_timestep, prompt, prompt_mask) + negative = model(latent, model_timestep, negative_prompt, prompt_mask) + distributed = negative.float() + 3.0 * (positive.float() - negative.float()) + delta = (distributed.float() - eager.float()).abs() + metrics = torch.tensor( + [delta.max(), delta.mean(), delta.norm() / eager.float().norm().clamp_min(1e-12)], device=device + ) + dist.all_reduce(metrics, op=dist.ReduceOp.MAX) + report = { + "fsdp": args.fsdp, + "batch_cfg": args.fsdp, + "world_size": world_size, + "joint_tokens": 7, + "ulysses_padding_tokens": 1 if world_size == 4 else 0, + "fp32_unsharded_parameters": sum(parameter.dtype == torch.float32 for parameter in model.parameters()), + "max_abs": float(metrics[0].item()), + "mean_abs": float(metrics[1].item()), + "relative_l2": float(metrics[2].item()), + } + if args.output and dist.get_rank() == 0: + with open(args.output, "w", encoding="utf-8") as handle: + json.dump(report, handle, indent=2, sort_keys=True) + handle.write("\n") + if dist.get_rank() == 0: + print(json.dumps(report, sort_keys=True)) + if report["max_abs"] != 0.0: + raise AssertionError(f"Ulysses SP parity failed: {report}") + finally: + dist.destroy_process_group() + + +if __name__ == "__main__": + main() diff --git a/tools/validation/validate_lingbot_video_refiner_handoff.py b/tools/validation/validate_lingbot_video_refiner_handoff.py new file mode 100644 index 0000000..e20d8cc --- /dev/null +++ b/tools/validation/validate_lingbot_video_refiner_handoff.py @@ -0,0 +1,255 @@ +"""Validate a LingBot-Video refiner MP4 handoff against the upstream loader. + +The upstream refiner reads the base result after it was written to MP4. The +native runtime uses an in-memory RGB tensor instead. This tool first verifies +that the compatibility MP4 path matches the upstream loader exactly, then can +optionally quantify the input change avoided by the native in-memory handoff. +""" + +from __future__ import annotations + +import argparse +import contextlib +import importlib +import json +import sys +import types +from collections.abc import Iterator +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import numpy as np +import torch + +from telefuser.pipelines.lingbot_video.refiner import load_refiner_video_file, prepare_refiner_video + +DEFAULT_UPSTREAM_ROOT = Path("work_dirs/lingbot-video-master") + + +@contextlib.contextmanager +def _upstream_import_path(root: Path) -> Iterator[None]: + """Temporarily expose the checked-out upstream package without modifying it.""" + root = root.resolve() + if not (root / "lingbot_video" / "__init__.py").is_file(): + raise FileNotFoundError(f"LingBot-Video package not found under {root}") + sys.path.insert(0, str(root)) + try: + yield + finally: + sys.path.remove(str(root)) + + +def tensor_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float | int | bool]: + """Return L0/L1 metrics for one refiner handoff tensor pair.""" + if reference.shape != candidate.shape: + return { + "shape_match": False, + "reference_shape": list(reference.shape), + "candidate_shape": list(candidate.shape), + } + ref = reference.float().cpu() + got = candidate.float().cpu() + delta = (got - ref).abs() + exact_mismatch_count = int(torch.count_nonzero(reference.cpu() != candidate.cpu()).item()) + cosine = ( + 1.0 + if exact_mismatch_count == 0 + else torch.nn.functional.cosine_similarity(ref.flatten(), got.flatten(), dim=0).item() + ) + return { + "shape_match": True, + "dtype_match": reference.dtype == candidate.dtype, + "max_abs": float(delta.max().item()) if delta.numel() else 0.0, + "mean_abs": float(delta.mean().item()) if delta.numel() else 0.0, + "relative_l2": float(delta.norm().div(ref.norm().clamp_min(1e-12)).item()), + "cosine": float(max(-1.0, min(1.0, cosine))), + "exact_mismatch_count": exact_mismatch_count, + } + + +def _pyav_decord_module() -> types.ModuleType: + """Return the small decord surface used by the source refiner loader.""" + try: + import av + except ImportError as exc: + raise RuntimeError("upstream MP4 handoff validation requires the optional PyAV dependency") from exc + + class Batch: + def __init__(self, values: np.ndarray) -> None: + self.values = values + + def asnumpy(self) -> np.ndarray: + return self.values + + class VideoReader: + def __init__(self, source: str, ctx: object) -> None: + del ctx + container = av.open(source) + try: + stream = next(iter(container.streams.video), None) + if stream is None: + raise ValueError(f"video has no video stream: {source}") + self._fps = float(stream.average_rate) if stream.average_rate is not None else 0.0 + self._frames = [frame.to_ndarray(format="rgb24") for frame in container.decode(stream)] + finally: + container.close() + + def __len__(self) -> int: + return len(self._frames) + + def get_avg_fps(self) -> float: + return self._fps + + def get_batch(self, indices: np.ndarray) -> Batch: + return Batch(np.stack([self._frames[int(index)] for index in indices])) + + module = types.ModuleType("decord") + module.VideoReader = VideoReader + module.cpu = lambda _: object() + return module + + +def load_upstream_refiner_video_file( + path: str | Path, + *, + upstream_root: Path, + height: int, + width: int, + sample_fps: int, + vae_tc: int, + max_frames: int | None, +) -> tuple[torch.Tensor, dict[str, Any]]: + """Run the source MP4 loader without requiring decord to be installed.""" + decord = _pyav_decord_module() + with _upstream_import_path(upstream_root): + upstream_utils = importlib.import_module("lingbot_video.utils") + with patch.dict(sys.modules, {"decord": decord}): + return upstream_utils.load_refiner_video_tensor( + path, + height, + width, + sample_fps=sample_fps, + vae_tc=vae_tc, + max_frames=max_frames, + ) + + +def validate_refiner_handoff( + path: str | Path, + *, + upstream_root: Path, + height: int, + width: int, + sample_fps: int = 24, + vae_tc: int = 4, + max_frames: int | None = None, + in_memory_video: torch.Tensor | None = None, + in_memory_fps: float | None = None, +) -> dict[str, Any]: + """Compare source and native MP4 loaders plus an optional memory handoff.""" + reference, reference_metadata = load_upstream_refiner_video_file( + path, + upstream_root=upstream_root, + height=height, + width=width, + sample_fps=sample_fps, + vae_tc=vae_tc, + max_frames=max_frames, + ) + candidate, candidate_metadata = load_refiner_video_file( + path, + height=height, + width=width, + sample_fps=sample_fps, + vae_tc=vae_tc, + max_frames=max_frames, + ) + report: dict[str, Any] = { + "input": str(path), + "source_mp4_to_telefuser_mp4": tensor_metrics(reference, candidate), + "metadata_match": reference_metadata == candidate_metadata, + "reference_metadata": reference_metadata, + "candidate_metadata": candidate_metadata, + } + if in_memory_video is not None: + if in_memory_fps is None: + raise ValueError("in_memory_fps is required when in_memory_video is provided") + prepared, memory_metadata = prepare_refiner_video( + in_memory_video, + source_fps=in_memory_fps, + height=height, + width=width, + sample_fps=sample_fps, + vae_tc=vae_tc, + max_frames=max_frames, + ) + report["in_memory_to_mp4"] = tensor_metrics(prepared, candidate) + report["in_memory_metadata"] = memory_metadata + report["in_memory_metadata_match"] = memory_metadata == candidate_metadata + return report + + +def exact_handoff_failures(report: dict[str, Any]) -> list[str]: + """Return source-MP4 compatibility failures without judging native handoff quality.""" + metrics = report["source_mp4_to_telefuser_mp4"] + failures: list[str] = [] + if not report["metadata_match"]: + failures.append("metadata_match") + for name in ("shape_match", "dtype_match"): + if not metrics.get(name, False): + failures.append(f"source_mp4_to_telefuser_mp4.{name}") + for name in ("max_abs", "mean_abs", "relative_l2", "exact_mismatch_count"): + if metrics.get(name) != 0: + failures.append(f"source_mp4_to_telefuser_mp4.{name}") + return failures + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--input", type=Path, required=True, help="Base MP4 passed to the refiner.") + parser.add_argument("--height", type=int, required=True, help="Refiner input height.") + parser.add_argument("--width", type=int, required=True, help="Refiner input width.") + parser.add_argument("--upstream-root", type=Path, default=DEFAULT_UPSTREAM_ROOT) + parser.add_argument("--sample-fps", type=int, default=24) + parser.add_argument("--vae-tc", type=int, default=4) + parser.add_argument("--max-frames", type=int) + parser.add_argument( + "--in-memory-video", + type=Path, + help="Optional [B,3,F,H,W] torch tensor saved before MP4 encoding.", + ) + parser.add_argument("--in-memory-fps", type=float, help="FPS for --in-memory-video.") + parser.add_argument("--output", type=Path, help="Optional JSON report path.") + parser.add_argument( + "--assert-exact", + action="store_true", + help="Exit nonzero unless the source MP4 compatibility path matches exactly.", + ) + args = parser.parse_args() + in_memory_video = None + if args.in_memory_video is not None: + in_memory_video = torch.load(args.in_memory_video, map_location="cpu", weights_only=True) + report = validate_refiner_handoff( + args.input, + upstream_root=args.upstream_root, + height=args.height, + width=args.width, + sample_fps=args.sample_fps, + vae_tc=args.vae_tc, + max_frames=args.max_frames, + in_memory_video=in_memory_video, + in_memory_fps=args.in_memory_fps, + ) + if args.assert_exact: + failures = exact_handoff_failures(report) + if failures: + raise SystemExit(f"Exact refiner MP4 handoff gate failed: {', '.join(failures)}") + payload = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main() diff --git a/tools/validation/validate_lingbot_video_refiner_output_handoff.py b/tools/validation/validate_lingbot_video_refiner_output_handoff.py new file mode 100644 index 0000000..9242eed --- /dev/null +++ b/tools/validation/validate_lingbot_video_refiner_output_handoff.py @@ -0,0 +1,314 @@ +"""Compare final Refiner outputs from in-memory and MP4 base-video handoffs.""" + +from __future__ import annotations + +import argparse +import json +import sys +import tempfile +import time +from pathlib import Path +from typing import Any + +import torch +from PIL import Image + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +if str(PROJECT_ROOT) not in sys.path: + sys.path.insert(0, str(PROJECT_ROOT)) + +from examples.lingbot_video.lingbot_video_moe_30b import build_pipeline, build_refiner +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + default_negative_caption, + load_lingbot_video_prompt, + load_refiner_first_frame, + load_refiner_video_file, + num_frames_from_duration, + prepare_refiner_video, +) + + +def _synchronize() -> None: + if torch.cuda.is_available(): + torch.cuda.synchronize() + + +def tensor_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float | int | bool]: + """Return L0/L1 metrics for one in-memory versus MP4 handoff tensor pair.""" + if reference.shape != candidate.shape: + return { + "shape_match": False, + "reference_shape": list(reference.shape), + "candidate_shape": list(candidate.shape), + } + ref = reference.float().cpu() + got = candidate.float().cpu() + delta = (got - ref).abs() + exact_mismatch_count = int(torch.count_nonzero(reference.cpu() != candidate.cpu()).item()) + cosine = ( + 1.0 + if exact_mismatch_count == 0 + else torch.nn.functional.cosine_similarity(ref.flatten(), got.flatten(), dim=0).item() + ) + return { + "shape_match": True, + "dtype_match": reference.dtype == candidate.dtype, + "max_abs": float(delta.max().item()) if delta.numel() else 0.0, + "mean_abs": float(delta.mean().item()) if delta.numel() else 0.0, + "relative_l2": float(delta.norm().div(ref.norm().clamp_min(1e-12)).item()), + "cosine": float(max(-1.0, min(1.0, cosine))), + "exact_mismatch_count": exact_mismatch_count, + } + + +def decoded_frame_quality_metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float | None | bool]: + """Return local-SSIM and PSNR for decoded ``[B,3,F,H,W]`` RGB tensors.""" + if reference.shape != candidate.shape: + return {"shape_match": False, "psnr_db": None, "ssim": None} + if reference.ndim != 5 or reference.shape[1] != 3: + raise ValueError("decoded frames must have shape [B,3,F,H,W]") + ref = reference.float().cpu().clamp(0.0, 1.0) + got = candidate.float().cpu().clamp(0.0, 1.0) + mse = torch.mean((got - ref).square()) + psnr = None if mse == 0 else float(-10.0 * torch.log10(mse).item()) + batch, channels, frames, height, width = ref.shape + kernel = max(1, min(11, height, width) | 1) + padding = kernel // 2 + ref_2d = ref.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width) + got_2d = got.permute(0, 2, 1, 3, 4).reshape(batch * frames, channels, height, width) + mean_ref = torch.nn.functional.avg_pool2d(ref_2d, kernel, stride=1, padding=padding) + mean_got = torch.nn.functional.avg_pool2d(got_2d, kernel, stride=1, padding=padding) + variance_ref = ( + torch.nn.functional.avg_pool2d(ref_2d.square(), kernel, stride=1, padding=padding) - mean_ref.square() + ) + variance_got = ( + torch.nn.functional.avg_pool2d(got_2d.square(), kernel, stride=1, padding=padding) - mean_got.square() + ) + covariance = ( + torch.nn.functional.avg_pool2d(ref_2d * got_2d, kernel, stride=1, padding=padding) - mean_ref * mean_got + ) + c1, c2 = 0.01**2, 0.03**2 + ssim = ((2 * mean_ref * mean_got + c1) * (2 * covariance + c2)) / ( + (mean_ref.square() + mean_got.square() + c1) * (variance_ref + variance_got + c2) + ) + return {"shape_match": True, "psnr_db": psnr, "ssim": float(ssim.mean().clamp(-1.0, 1.0).item())} + + +def side_by_side_frames(reference: torch.Tensor, candidate: torch.Tensor) -> torch.Tensor: + """Place matching decoded videos side by side for visual handoff review.""" + if reference.shape != candidate.shape: + raise ValueError("comparison videos must have matching shapes") + if reference.ndim != 5 or reference.shape[1] != 3: + raise ValueError("comparison videos must have shape [B,3,F,H,W]") + return torch.cat((reference, candidate), dim=-1) + + +def _write_video_tensor(video: torch.Tensor, path: Path, fps: int = 24) -> None: + """Write a normalized ``[B,3,F,H,W]`` RGB tensor without rescaling uint8 frames.""" + frames = video[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + path.parent.mkdir(parents=True, exist_ok=True) + if frames.shape[0] == 1: + Image.fromarray((frames[0] * 255).round().astype("uint8")).save(path) + return + from diffusers.utils import export_to_video + + export_to_video(list(frames), str(path), fps=fps) + + +def write_side_by_side_video(reference: torch.Tensor, candidate: torch.Tensor, path: Path, fps: int = 24) -> None: + """Write in-memory and MP4-refined decoded frames as a visual comparison artifact.""" + comparison = side_by_side_frames(reference, candidate) + _write_video_tensor(comparison, path, fps) + + +def _measure(operation: Any) -> tuple[Any, float]: + _synchronize() + started = time.perf_counter() + result = operation() + _synchronize() + return result, time.perf_counter() - started + + +def _image_to_tensor(path: str) -> torch.Tensor: + image = Image.open(path).convert("RGB") + return torch.from_numpy(__import__("numpy").asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).float() + + +def build_refiner_handoff_report( + *, + memory_input: torch.Tensor, + mp4_input: torch.Tensor, + memory_metadata: dict[str, Any], + mp4_metadata: dict[str, Any], + memory_output: torch.Tensor, + mp4_output: torch.Tensor, + memory_seconds: float, + mp4_seconds: float, +) -> dict[str, Any]: + """Report effects of replacing the source's lossy MP4 handoff.""" + return { + "comparison_baseline": "source_mp4_round_trip", + "mp4_round_trip_is_lossy": True, + "metadata_match": memory_metadata == mp4_metadata, + "in_memory_to_mp4_input": tensor_metrics(memory_input.cpu(), mp4_input.cpu()), + "final_output": tensor_metrics(memory_output.cpu(), mp4_output.cpu()), + "final_output_quality": decoded_frame_quality_metrics(memory_output, mp4_output), + "memory_refiner_seconds": memory_seconds, + "mp4_refiner_seconds": mp4_seconds, + } + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", type=Path, required=True) + parser.add_argument("--caption-json", type=Path, required=True) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int) + parser.add_argument("--steps", type=int, default=40) + parser.add_argument("--guidance-scale", type=float, default=3.0) + parser.add_argument("--negative-caption", default=None) + parser.add_argument("--image") + parser.add_argument("--refiner-height", type=int, default=1088) + parser.add_argument("--refiner-width", type=int, default=1920) + parser.add_argument("--refiner-steps", type=int, default=8) + parser.add_argument("--refiner-guidance-scale", type=float, default=3.0) + parser.add_argument("--refiner-shift", type=float, default=3.0) + parser.add_argument("--refiner-t-thresh", type=float, default=0.85) + parser.add_argument("--refiner-tail-steps", type=int, default=2) + parser.add_argument("--seed", type=int, default=42) + parser.add_argument("--disable-cpu-offload", action="store_true") + parser.add_argument("--output", type=Path) + parser.add_argument( + "--comparison-output", + type=Path, + help="Optional side-by-side memory|MP4 Refiner output video for human review.", + ) + args = parser.parse_args() + + caption, duration = load_lingbot_video_prompt(args.caption_json) + num_frames = ( + args.num_frames + if args.num_frames is not None + else num_frames_from_duration(duration) + if duration is not None + else 121 + ) + source_image = _image_to_tensor(args.image) if args.image else None + negative_caption = ( + args.negative_caption if args.negative_caption is not None else default_negative_caption(num_frames) + ) + pipeline, base_load_seconds = _measure( + lambda: build_pipeline( + args.model_dir, + guidance_scale=args.guidance_scale, + num_inference_steps=args.steps, + cpu_offload=False if args.disable_cpu_offload else None, + ) + ) + generation, base_generation_seconds = _measure( + lambda: pipeline.generate( + LingBotVideoRequest( + caption=caption, + height=args.height, + width=args.width, + num_frames=num_frames, + image=source_image, + ), + negative_caption=negative_caption, + generator=torch.Generator("cuda").manual_seed(args.seed), + ) + ) + frames = generation.output + if pipeline.text_stage is None: + raise RuntimeError("LingBot refiner requires the base text stage") + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = pipeline.text_stage.encode(caption) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative = torch.zeros_like(positive) + negative_mask = positive_mask.clone() + _, base_release_seconds = _measure(pipeline.release_gpu_resources) + refiner, refiner_load_seconds = _measure( + lambda: build_refiner(args.model_dir, cpu_offload=not args.disable_cpu_offload) + ) + memory_input, memory_metadata = prepare_refiner_video( + frames, source_fps=24.0, height=args.refiner_height, width=args.refiner_width + ) + clean_first_frame = ( + load_refiner_first_frame( + args.image, + target_height=args.refiner_height, + target_width=args.refiner_width, + geometry_height=args.height, + geometry_width=args.width, + ) + if args.image + else None + ) + + def refine(video: torch.Tensor) -> tuple[torch.Tensor, float]: + return _measure( + lambda: refiner.refine( + video, + positive, + negative, + positive_mask, + negative_mask, + num_inference_steps=args.refiner_steps, + guidance_scale=args.refiner_guidance_scale, + shift=args.refiner_shift, + t_thresh=args.refiner_t_thresh, + tail_steps=args.refiner_tail_steps, + clean_first_frame=clean_first_frame, + generator=torch.Generator("cuda").manual_seed(args.seed), + ) + ) + + memory_output, memory_seconds = refine(memory_input) + with tempfile.TemporaryDirectory() as temporary_directory: + mp4_path = Path(temporary_directory) / "base.mp4" + _write_video_tensor(frames, mp4_path) + mp4_input, mp4_metadata = load_refiner_video_file( + mp4_path, height=args.refiner_height, width=args.refiner_width + ) + mp4_output, mp4_seconds = refine(mp4_input) + report = build_refiner_handoff_report( + memory_input=memory_input, + mp4_input=mp4_input, + memory_metadata=memory_metadata, + mp4_metadata=mp4_metadata, + memory_output=memory_output, + mp4_output=mp4_output, + memory_seconds=memory_seconds, + mp4_seconds=mp4_seconds, + ) + report.update( + model_dir=str(args.model_dir), + height=args.height, + width=args.width, + num_frames=num_frames, + steps=args.steps, + refiner_height=args.refiner_height, + refiner_width=args.refiner_width, + refiner_steps=args.refiner_steps, + cpu_offload=not args.disable_cpu_offload, + base_load_seconds=base_load_seconds, + base_generation_seconds=base_generation_seconds, + base_release_seconds=base_release_seconds, + refiner_load_seconds=refiner_load_seconds, + ) + if args.comparison_output is not None: + write_side_by_side_video(memory_output, mp4_output, args.comparison_output) + report["comparison_output"] = str(args.comparison_output) + payload = json.dumps(report, indent=2, sort_keys=True) + "\n" + if args.output is not None: + args.output.parent.mkdir(parents=True, exist_ok=True) + args.output.write_text(payload, encoding="utf-8") + print(payload, end="") + + +if __name__ == "__main__": + main()