From 054a16640e467417cf6d9faad6abadf4b00ac124 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 14:04:59 +0000 Subject: [PATCH 1/9] feat(lingbot-video): add single-process inference runtime - add Dense and MoE model implementations with T2V/TI2V pipelines and refiner support - add service contracts, examples, documentation, and source attribution - add benchmark, parity, replay, and handoff validation tooling with unit coverage - preserve VAE, visual conditioning, RGB handoff, geometry, and prompt-envelope behavior Verification: - git diff --cached --check --- CLAUDE.md | 10 +- README.md | 1 + docs/en/index.md | 1 + docs/en/lingbot_video.md | 216 +++++++ docs/zh/index.md | 1 + docs/zh/lingbot_video.md | 162 +++++ examples/lingbot_video/README.md | 30 + .../lingbot_video/lingbot_video_generate.py | 142 +++++ .../lingbot_video/lingbot_video_service.py | 205 ++++++ telefuser/models/lingbot_video_dit.py | 338 ++++++++++ telefuser/models/lingbot_video_moe.py | 249 ++++++++ telefuser/ops/__init__.py | 2 + telefuser/ops/moe.py | 49 ++ telefuser/pipelines/lingbot_video/__init__.py | 86 +++ telefuser/pipelines/lingbot_video/data.py | 206 ++++++ .../pipelines/lingbot_video/denoising.py | 109 ++++ telefuser/pipelines/lingbot_video/loading.py | 172 +++++ telefuser/pipelines/lingbot_video/pipeline.py | 198 ++++++ telefuser/pipelines/lingbot_video/refiner.py | 278 +++++++++ telefuser/pipelines/lingbot_video/runtime.py | 129 ++++ .../pipelines/lingbot_video/text_encoding.py | 141 +++++ telefuser/pipelines/lingbot_video/vae.py | 102 +++ tests/unit/models/test_lingbot_video_dit.py | 68 ++ tests/unit/models/test_lingbot_video_moe.py | 96 +++ .../unit/pipelines/lingbot_video/test_data.py | 67 ++ .../pipelines/lingbot_video/test_pipeline.py | 183 ++++++ .../pipelines/lingbot_video/test_refiner.py | 159 +++++ .../lingbot_video/test_refiner_stage.py | 89 +++ .../pipelines/lingbot_video/test_runtime.py | 88 +++ .../lingbot_video/test_service_contract.py | 234 +++++++ .../lingbot_video/test_text_encoding.py | 25 + .../unit/pipelines/lingbot_video/test_vae.py | 55 ++ .../tools/test_benchmark_lingbot_video.py | 38 ++ .../test_capture_lingbot_video_reference.py | 59 ++ .../test_inspect_lingbot_video_checkpoint.py | 54 ++ ...st_replay_lingbot_video_dense_reference.py | 86 +++ .../test_run_lingbot_video_moe_parity.py | 34 + ...t_run_lingbot_video_refiner_core_parity.py | 34 + ..._validate_lingbot_video_refiner_handoff.py | 92 +++ ...te_lingbot_video_refiner_output_handoff.py | 85 +++ tools/validation/benchmark_lingbot_video.py | 285 +++++++++ .../capture_lingbot_video_reference.py | 589 ++++++++++++++++++ .../compare_lingbot_video_parity.py | 95 +++ .../inspect_lingbot_video_checkpoint.py | 162 +++++ .../replay_lingbot_video_dense_reference.py | 383 ++++++++++++ .../run_lingbot_video_dense_parity.py | 80 +++ .../run_lingbot_video_moe_parity.py | 101 +++ .../run_lingbot_video_refiner_core_parity.py | 172 +++++ .../validate_lingbot_video_refiner_handoff.py | 255 ++++++++ ...te_lingbot_video_refiner_output_handoff.py | 312 ++++++++++ 50 files changed, 6806 insertions(+), 1 deletion(-) create mode 100644 docs/en/lingbot_video.md create mode 100644 docs/zh/lingbot_video.md create mode 100644 examples/lingbot_video/README.md create mode 100644 examples/lingbot_video/lingbot_video_generate.py create mode 100644 examples/lingbot_video/lingbot_video_service.py create mode 100644 telefuser/models/lingbot_video_dit.py create mode 100644 telefuser/models/lingbot_video_moe.py create mode 100644 telefuser/ops/moe.py create mode 100644 telefuser/pipelines/lingbot_video/__init__.py create mode 100644 telefuser/pipelines/lingbot_video/data.py create mode 100644 telefuser/pipelines/lingbot_video/denoising.py create mode 100644 telefuser/pipelines/lingbot_video/loading.py create mode 100644 telefuser/pipelines/lingbot_video/pipeline.py create mode 100644 telefuser/pipelines/lingbot_video/refiner.py create mode 100644 telefuser/pipelines/lingbot_video/runtime.py create mode 100644 telefuser/pipelines/lingbot_video/text_encoding.py create mode 100644 telefuser/pipelines/lingbot_video/vae.py create mode 100644 tests/unit/models/test_lingbot_video_dit.py create mode 100644 tests/unit/models/test_lingbot_video_moe.py create mode 100644 tests/unit/pipelines/lingbot_video/test_data.py create mode 100644 tests/unit/pipelines/lingbot_video/test_pipeline.py create mode 100644 tests/unit/pipelines/lingbot_video/test_refiner.py create mode 100644 tests/unit/pipelines/lingbot_video/test_refiner_stage.py create mode 100644 tests/unit/pipelines/lingbot_video/test_runtime.py create mode 100644 tests/unit/pipelines/lingbot_video/test_service_contract.py create mode 100644 tests/unit/pipelines/lingbot_video/test_text_encoding.py create mode 100644 tests/unit/pipelines/lingbot_video/test_vae.py create mode 100644 tests/unit/tools/test_benchmark_lingbot_video.py create mode 100644 tests/unit/tools/test_capture_lingbot_video_reference.py create mode 100644 tests/unit/tools/test_inspect_lingbot_video_checkpoint.py create mode 100644 tests/unit/tools/test_replay_lingbot_video_dense_reference.py create mode 100644 tests/unit/tools/test_run_lingbot_video_moe_parity.py create mode 100644 tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py create mode 100644 tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py create mode 100644 tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py create mode 100644 tools/validation/benchmark_lingbot_video.py create mode 100644 tools/validation/capture_lingbot_video_reference.py create mode 100644 tools/validation/compare_lingbot_video_parity.py create mode 100644 tools/validation/inspect_lingbot_video_checkpoint.py create mode 100644 tools/validation/replay_lingbot_video_dense_reference.py create mode 100644 tools/validation/run_lingbot_video_dense_parity.py create mode 100644 tools/validation/run_lingbot_video_moe_parity.py create mode 100644 tools/validation/run_lingbot_video_refiner_core_parity.py create mode 100644 tools/validation/validate_lingbot_video_refiner_handoff.py create mode 100644 tools/validation/validate_lingbot_video_refiner_output_handoff.py diff --git a/CLAUDE.md b/CLAUDE.md index 1b98247..007bb78 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,13 @@ 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. Grouped-GEMM, FP8, and distributed execution require separate parity and benchmark evidence before being enabled. + ### Layer Architecture Principles For Models TeleFuser's model follows a strict layered architecture for operations: diff --git a/README.md b/README.md index cb9027a..4af7fa9 100644 --- a/README.md +++ b/README.md @@ -195,6 +195,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) | +| `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | Precision-first Dense/MoE generation with a single-process base-to-refiner path; see [docs/en/lingbot_video.md](docs/en/lingbot_video.md) | ### Image Generation and Other Multimodal Pipelines diff --git a/docs/en/index.md b/docs/en/index.md index c96457e..88b01ec 100644 --- a/docs/en/index.md +++ b/docs/en/index.md @@ -68,6 +68,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/lingbot_video.md b/docs/en/lingbot_video.md new file mode 100644 index 0000000..ef402ca --- /dev/null +++ b/docs/en/lingbot_video.md @@ -0,0 +1,216 @@ +# LingBot-Video + +LingBot-Video supports Dense and MoE base DiTs for T2I, T2V, and TI2V. The +MoE checkpoint also includes a separate low-noise refiner. The integration is +precision-first: use the upstream reference capture artifacts before enabling +backend or distributed optimizations. + +## Checkpoints + +The Dense loader consumes a Diffusers `transformer/` directory directly: + +```python +from telefuser.pipelines.lingbot_video import load_lingbot_video_dense_transformer + +transformer = load_lingbot_video_dense_transformer( + "/path/to/lingbot-video-dense-1.3b/transformer" +) +``` + +The MoE and refiner directories are sharded. Use +`load_lingbot_video_moe_transformer` for either `transformer/` or `refiner/`. +The default sorted eager expert path preserves upstream route ordering and keeps +a `where`-based diagnostic fallback. It is a validated single-GPU BF16 path, +but not a grouped-GEMM, FP8, or distributed production-throughput backend. For +`variant="moe"`, the runtime defaults to stage CPU offload so the base DiT, +text encoder, VAE, and separately loaded refiner do not need to reside on one +GPU. Set `cpu_offload=False` only when GPU capacity is known to be sufficient. + +## Prompt preparation + +The generation pipeline intentionally consumes the structured JSON caption, while +prompt rewriting remains an optional, separately deployable workflow. Preserve the +official two-stage contract: EXPAND uses the base VLM with its LoRA disabled, and +MAP uses the same VLM with the LingBot rewriter LoRA enabled. For TI2V, provide +the identical first frame to both the rewriter and TeleFuser generation. + +```bash +REWRITER_BASE_MODEL=/path/to/Qwen3.6-27B \ +REWRITER_ADAPTER=/path/to/lingbot-video-rewriter-lora \ +python work_dirs/lingbot-video-master/rewriter/inference.py \ + --mode t2v --prompt "" --duration 5 --output prompt.json +``` + +Pass the resulting `prompt.json` to `--caption-json`, or serialize its +`caption` object as the service `prompt`. The rewriter is deliberately not loaded +inside a DiT service process because the two models have separate deployment and +capacity requirements. + +Unless explicitly overridden, the pipeline, CLI, and service use the checkpoint's +structured negative CFG caption. T2I uses the source still-image variant, while +T2V and TI2V use the source video variant, including its temporal-stability +constraints. Do not replace an omitted negative prompt with an empty string when +reproducing an upstream sample: it changes the Qwen3-VL negative condition and +can materially alter color and image quality. + +## Runtime composition + +`LingBotVideoPipeline` composes independently loaded stages: + +- `LingBotVideoTextEncodingStage` encodes the structured JSON caption with + Qwen3-VL. +- `LingBotVideoDenoisingStage` runs source-order, two-forward CFG. +- `LingBotVideoVAEEncodeStage` and `LingBotVideoVAEDecodeStage` apply the + checkpoint VAE's latent mean/std normalization. +- `FlowUniPCMultistepScheduler` owns the sigma/timestep sequence. + +Attach stages with `pipeline.set_runtime(...)` after initializing the pipeline +config. Provide a structured JSON caption, not casual unstructured text. + +For the standard checkpoint layout, `build_lingbot_video_pipeline` loads those +components directly without importing the upstream runtime: + +```python +from telefuser.pipelines.lingbot_video import LingBotVideoRequest, build_lingbot_video_pipeline + +pipeline = build_lingbot_video_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) +frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) +``` + +Direct API and CLI heights and widths must be divisible by 16: the Wan VAE +downsamples by eight and the DiT spatially patchifies the resulting latents by two. + +The default `AttentionConfig` uses TeleFuser's SDPA dispatcher and remains the +source-equivalent numerical path. Alternative attention backends are opt-in +through `attention_config=` and require a separate L2 parity report; they are +not enabled by the service or CLI defaults. + +The VAE decode stage returns RGB video in `[0,1]`. Video callers must pass these +float frames directly to Diffusers `export_to_video`, which performs the uint8 +conversion itself. Converting to uint8 before that call applies a second 255 +scale and wraps channel values, producing a negative-like MP4. + +## TI2V + +TI2V has two independent first-frame condition paths: + +1. The image is supplied to Qwen3-VL as visual input with the caption. +2. The image is VAE-encoded to a clean temporal-prefix latent. It is written + before every denoising step and once after the final scheduler step. + +The current pipeline accepts a raw RGB condition image in the range [0,255] with shape +`[B, 3, H, W]` (or `[B, 3, F, H, W]`, using frame zero). It resizes and center-crops the image before separately passing it to Qwen3-VL and the VAE. + +## Service + +The service entrypoint exposes `t2i`, `t2v`, and `i2v` through TeleFuser service APIs and requires a structured JSON string as `prompt`: + +```bash +LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ + telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +``` + +Set `LINGBOT_VIDEO_VARIANT=moe` and `LINGBOT_VIDEO_ENABLE_REFINER=1` to use the MoE checkpoint with its separate refiner. The service releases base-stage weights before loading the refiner; `refine` is also an explicit per-request boolean parameter. +Requested service resolutions are rounded up to the LingBot VAE-and-DiT +sixteen-pixel spatial grid; for example, `480p` at `16:9` resolves to 864x480. +Pass `negative_prompt` only to override the source-compatible default; an +explicit empty string remains a supported override. + +## Refiner + +`LingBotVideoRefinerStage` takes the base RGB output in memory, VAE-encodes it, +mixes it with noise at `t_thresh`, and samples the low-noise sigma tail. It can +also preserve a clean TI2V frame zero. Base and refiner are separate runtime +stages. Call `base_pipeline.release_gpu_resources()` before loading the refiner +when they share a GPU. +The included CLI implements this lifecycle for MoE checkpoints: + +```bash +python examples/lingbot_video/lingbot_video_generate.py \ + --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ + --caption-json /path/to/caption.json --output result.mp4 +``` + +With `--image first_frame.png`, the CLI also applies the upstream TI2V frame-zero geometry to the refiner condition. + +For an in-memory base-to-refiner handoff, call `prepare_refiner_video(...)` before +`LingBotVideoRefinerStage.refine(...)`. It matches upstream training-aligned frame +selection and bicubic resize without an MP4 write/read round trip; pass the base +output FPS explicitly. Validate this path against the corresponding source MP4 baseline. +The MP4 compatibility test uses the upstream Diffusers writer and compares the upstream loader (through a PyAV-backed decord adapter when decord is unavailable) with this loader tensor-for-tensor. + + +## Validation + +Capture upstream artifacts before comparing a TeleFuser run: + +```bash +python tools/validation/capture_lingbot_video_reference.py --dry-run +python tools/validation/capture_lingbot_video_reference.py --all-cases --mode t2i --mode t2v --mode ti2v --trace sampled +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-dense-1.3b --variant dense --output dense-load-report.json +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --output moe-load-report.json +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant refiner --output refiner-load-report.json +python tools/validation/compare_lingbot_video_parity.py REFERENCE CANDIDATE +python tools/validation/replay_lingbot_video_dense_reference.py --reference-dir work_dirs/lingbot_video_reference/t2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --validate-ti2v-vae --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --reference-root work_dirs/lingbot_video_reference_all_cases --assert-exact --output dense-all-cases-replay.json +PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_moe_parity.py --transformer-dir /path/to/lingbot-video-moe-30b-a3b/transformer --assert-exact +PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_refiner_core_parity.py --model-root /path/to/lingbot-video-moe-30b-a3b --assert-exact +python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.mp4 --height 1088 --width 1920 --assert-exact +python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 +python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +``` + +`validate_lingbot_video_refiner_handoff.py` checks that the TeleFuser MP4 +compatibility loader is tensor-identical to the source loader. Add +`--in-memory-video` and `--in-memory-fps` to quantify the input difference +introduced by MP4 encoding; this comparison does not replace a final refiner +output-quality evaluation. +Use `--assert-exact` to make the source MP4 compatibility comparison fail on +metadata, tensor shape, dtype, or value drift. It intentionally does not judge +the native in-memory handoff, whose difference from lossy MP4 is expected. +`validate_lingbot_video_refiner_output_handoff.py` generates one MoE base +sample, drives the refiner with both the native RGB tensor and a temporary MP4 +round trip using identical prompt conditions and RNG state, then reports the +final-output L2 difference. This is intentionally a quality comparison, not an +equivalence test: the upstream refiner uses the lossy MP4 round trip, while the +native path removes that intermediate encoding. The report includes decoded-frame +PSNR and local SSIM; `--comparison-output` writes memory output on the left and +MP4-round-trip output on the right for human review. + +The capture tool records scheduler tensors, prompt tensors, selected denoising steps, latent inputs/outputs, the generation seed, RNG state hashes, and decoded frame hashes. Add `--validate-text` to the replay command to compare Qwen3-VL processor inputs and final embeddings; TI2V validation also compares the preprocessed first frame. Add `--validate-ti2v-vae` to compare the sampled clean condition latent; use `--seed` only with pre-existing captures that lack seed metadata. +Use `--reference-root` for an all-case Dense DiT/VAE replay. It retains one +loaded Dense transformer and VAE across captured runs, while instantiating a +fresh scheduler for each run so each capture preserves its own sampling setup. +Add `--assert-exact` to make the command fail when any recorded tensor differs +in shape or value, so it can serve as a CI parity gate rather than a report-only +diagnostic. +The checkpoint-inspection tool performs the normal strict load and records the +consumed config fields, checkpoint-key coverage, component/block parameter +counts, dtype/device distribution, retained FP32 parameter count, and model +memory allocation evidence. +The benchmark tool reports one-time setup separately from warmup and measured metrics for checkpoint load, text encoding, each denoising step, VAE, refiner, output encoding, and peak GPU memory. Use it to establish a baseline before enabling an optimization; record full-resolution measurements separately from smoke runs. +When `--negative-caption` is omitted, the benchmark uses the same source-compatible +T2I or video negative caption as the pipeline, CLI, and service. Pass an explicit +empty string only when intentionally benchmarking that semantic override. +For a base-plus-refiner run, it also records the serial base release and refiner +load phases. The default sorted eager MoE path is source-equivalent and has an +explicit diagnostic fallback, but it is not a grouped-GEMM or FP8 +production-throughput backend. +For T2I/T2V, the refiner reuses the exact CFG text conditions from the base +generation and reports this as `refiner_prompt_conditions_reused`; TI2V keeps +the source-compatible text-only refiner encoding path. +The refiner core CLI injects identical latent, noise, prompt, and frame-zero condition tensors into the upstream and TeleFuser low-noise paths. It offloads the upstream DiT before loading the TeleFuser DiT, so both 30B models do not overlap in GPU memory. +Add `--assert-exact` to the MoE or refiner core validator to enforce the +zero-drift numerical-oracle gate instead of only writing metrics. + +## Requirements and limitations + +The numerical-oracle path requires CUDA, PyTorch, Diffusers, Transformers, and the +checkpoint components `transformer/`, `text_encoder/`, `processor/`, `vae/`, and +`scheduler/`. Dense runs source-equivalently on one GPU. The current MoE expert +implementation is an eager correctness path; do not use it as a 30B production +backend. FSDP, Ulysses/CFG parallelism, FlashAttention, and FP8 experts are intentionally not +enabled. The service and CLI both support a serial single-process base-plus-refiner lifecycle; distributed alignment is outside the current support scope. diff --git a/docs/zh/index.md b/docs/zh/index.md index 2690f9d..5a051e4 100644 --- a/docs/zh/index.md +++ b/docs/zh/index.md @@ -67,6 +67,7 @@ AdaTaylorCache 和运行时缓存控制，面向重复生成工作负载。 | FlashVSR | VSR | 视频超分辨率 | | LiveAct | S2V | 语音转视频 | | LongCat-Video | T2V, I2V | 长视频生成 | +| LingBot-Video | T2I, T2V, TI2V, MoE refiner | 精度优先的 Dense/MoE 视频生成 | ### 图像生成 diff --git a/docs/zh/lingbot_video.md b/docs/zh/lingbot_video.md new file mode 100644 index 0000000..6b1a440 --- /dev/null +++ b/docs/zh/lingbot_video.md @@ -0,0 +1,162 @@ +# LingBot-Video + +LingBot-Video 支持 Dense 与 MoE base DiT 的 T2I、T2V、TI2V；MoE checkpoint +还包含独立的低噪 refiner。本集成按精度优先实现，在启用优化或分布式后端前应先完成上游参考对比。 + +## Checkpoint + +Dense checkpoint 可直接加载 Diffusers `transformer/` 目录： + +```python +from telefuser.pipelines.lingbot_video import load_lingbot_video_dense_transformer + +transformer = load_lingbot_video_dense_transformer("/path/to/transformer") +``` + +MoE 与 refiner 使用分片 safetensors，调用 +`load_lingbot_video_moe_transformer` 加载 `transformer/` 或 `refiner/`。 +默认的 sorted eager expert 路径保持上游 route 顺序，并保留 `where` 诊断 fallback。 +它是经过验证的单卡 BF16 路径，但不是 grouped-GEMM、FP8 或分布式生产吞吐后端。 +`variant="moe"` 默认启用 stage CPU offload，避免 base DiT、text encoder、VAE 与 +独立 refiner 同时驻留单张 GPU。仅在显存明确充足时才设置 `cpu_offload=False`。 + +## Prompt 准备 + +生成管线有意只消费结构化 JSON caption；prompt rewriter 是可独立部署的可选流程。 +必须保留上游两阶段语义：EXPAND 使用不挂 LoRA 的 base VLM，MAP 使用同一 base VLM +并启用 LingBot rewriter LoRA。TI2V 必须将同一张首帧同时传给 rewriter 与 TeleFuser。 + +```bash +REWRITER_BASE_MODEL=/path/to/Qwen3.6-27B \ +REWRITER_ADAPTER=/path/to/lingbot-video-rewriter-lora \ +python work_dirs/lingbot-video-master/rewriter/inference.py \ + --mode t2v --prompt "" --duration 5 --output prompt.json +``` + +将输出的 `prompt.json` 传给 `--caption-json`，或将其中的 `caption` 对象序列化为 +服务请求的 `prompt`。由于 rewriter 与 DiT 的部署和显存需求独立，TeleFuser 不会在 +同一 DiT 服务进程内加载 rewriter。 + +未显式覆盖时，pipeline、CLI 与 service 都会使用 checkpoint 的结构化负向 CFG +caption。T2I 使用上游的静态图像版本；T2V 与 TI2V 使用包含时序稳定性约束的上游视频 +版本。复现上游样本时，不能以空字符串替代未传入的负向 prompt：这会改变 Qwen3-VL 的 +负向 condition，并可能显著影响颜色与画质。 + +## 运行时组成 + +`LingBotVideoPipeline` 通过 `set_runtime(...)` 组合独立加载的 stage： + +- `LingBotVideoTextEncodingStage`：Qwen3-VL structured JSON caption 编码。 +- `LingBotVideoDenoisingStage`：source-order 的 two-forward CFG。 +- `LingBotVideoVAEEncodeStage` / `LingBotVideoVAEDecodeStage`：使用 checkpoint 的 latent mean/std。 +- `FlowUniPCMultistepScheduler`：sigma/timestep 调度。 + +核心 pipeline 应输入结构化 JSON caption，不应将普通自然语言 prompt 直接替代 rewriter 输出。 + +标准 checkpoint 目录可直接使用 `build_lingbot_video_pipeline` 装配，不依赖上游 runtime： + +```python +from telefuser.pipelines.lingbot_video import LingBotVideoRequest, build_lingbot_video_pipeline + +pipeline = build_lingbot_video_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) +frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) +``` + +直接使用 API 或 CLI 时，高和宽必须能被 16 整除：Wan VAE 会以 8 倍下采样，DiT 会再以 2 倍对 latent 作空间 patchify。 + +默认 `AttentionConfig` 使用 TeleFuser 的 SDPA dispatcher，仍属于 source-equivalent +数值路径。其他 attention backend 必须通过 `attention_config=` 显式启用并单独提供 L2 +parity 报告；service 与 CLI 默认不会启用它们。 + +VAE decode stage 的 RGB video 范围为 `[0,1]`。视频调用方必须将这些浮点帧直接传给 +Diffusers 的 `export_to_video`，由它自行转换为 uint8。若在调用前先转 uint8，就会再次 +乘以 255 并使通道值溢出，生成近似负片的 MP4。 + +## 服务 + +服务入口通过 TeleFuser API 暴露 `t2i`、`t2v`、`i2v`，`prompt` 必须是 structured JSON 字符串： + +```bash +LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ + telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +``` + +设置 `LINGBOT_VIDEO_VARIANT=moe` 与 `LINGBOT_VIDEO_ENABLE_REFINER=1` 可启用带独立 refiner 的 MoE checkpoint。服务会在加载 refiner 前释放 base stage 的权重；`refine` 也可作为显式的单请求布尔参数。 +服务请求的分辨率会向上对齐到 LingBot VAE 与 DiT 所需的 16 像素网格；例如 `480p`、`16:9` 会解析为 864x480。 +仅在确实需要覆盖上游默认值时才传入 `negative_prompt`；显式传入空字符串仍是受支持的覆盖方式。 + +## TI2V 与 Refiner + +TI2V 接收范围为 [0,255] 的 RGB tensor（`[B, 3, H, W]` 或 `[B, 3, F, H, W]`，后者取第 0 帧）。管线会先 resize 和 center-crop，再将同一视觉帧传给正向与负向 Qwen3-VL CFG 分支，并编码为 VAE clean temporal-prefix latent；该 latent 会在每次 denoising step 前及最终 step 后重写。 + +`LingBotVideoRefinerStage` 直接接收 base RGB tensor，避免 MP4/decord 往返；它按 `t_thresh` 与噪声混合，并使用 low-noise sigma tail 采样。base 与 refiner 是独立运行时 stage；共享 GPU 时，请在加载 refiner 前调用 `base_pipeline.release_gpu_resources()`。 + +示例 CLI 已为 MoE checkpoint 自动完成这个生命周期： + +```bash +python examples/lingbot_video/lingbot_video_generate.py \ + --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ + --caption-json /path/to/caption.json --output result.mp4 +``` + +附加 `--image first_frame.png` 时，CLI 也会把上游 TI2V frame-zero 几何规则应用到 refiner condition。 + +如需在内存中交接基座输出与 refiner，可先调用 `prepare_refiner_video(...)`，再调用 +`LingBotVideoRefinerStage.refine(...)`。该函数复刻上游训练对齐抽帧与双三次缩放， +无需 MP4 写入/读取；必须显式传入基座输出 FPS，并应与对应的上游 MP4 基线进行验证。 +MP4 兼容性测试使用上游 Diffusers writer；当 decord 不可用时，通过 PyAV-backed decord adapter 调用上游 loader，并逐 tensor 对比本 loader。 + + +## 验证 + +```bash +python tools/validation/capture_lingbot_video_reference.py --dry-run +python tools/validation/capture_lingbot_video_reference.py --all-cases --mode t2i --mode t2v --mode ti2v --trace sampled +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-dense-1.3b --variant dense --output dense-load-report.json +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --output moe-load-report.json +python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant refiner --output refiner-load-report.json +python tools/validation/compare_lingbot_video_parity.py REFERENCE CANDIDATE +python tools/validation/replay_lingbot_video_dense_reference.py --reference-dir work_dirs/lingbot_video_reference/t2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --validate-ti2v-vae --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 +python tools/validation/replay_lingbot_video_dense_reference.py --reference-root work_dirs/lingbot_video_reference_all_cases --assert-exact --output dense-all-cases-replay.json +PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_moe_parity.py --transformer-dir /path/to/lingbot-video-moe-30b-a3b/transformer --assert-exact +PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_refiner_core_parity.py --model-root /path/to/lingbot-video-moe-30b-a3b --assert-exact +python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.mp4 --height 1088 --width 1920 --assert-exact +python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 +python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +``` + +`validate_lingbot_video_refiner_handoff.py` 验证 TeleFuser 的 MP4 兼容 loader 与上游 +loader 的输入 tensor 完全一致。可通过 `--in-memory-video` 和 `--in-memory-fps` 量化 MP4 +编码带来的输入差异；该比较不替代最终 Refiner 输出质量评估。 +使用 `--assert-exact` 可在 source MP4 兼容路径的 metadata、tensor shape、dtype 或值发生 +漂移时使命令失败。它刻意不评判原生内存 handoff，因为它与有损 MP4 的差异是预期行为。 +`validate_lingbot_video_refiner_output_handoff.py` 会先生成一个 MoE base sample，再以 +完全相同的 prompt condition 与 RNG state 分别通过原生 RGB tensor 和临时 MP4 回环驱动 +refiner，并报告最终输出的 L2 差异。这是质量对比而非等价性测试：上游 refiner 使用有损 +MP4 回环，而原生路径去除了这一中间编码。报告还会给出 decoded-frame PSNR 与局部 SSIM； +`--comparison-output` 会写出左侧为内存输出、右侧为 MP4 回环输出的并排视频，供人工审核。 + +捕获工具记录 prompt/scheduler/selected denoising tensors、生成 seed、RNG state hash 和 decoded frame hash，用于 L0/L1 parity。replay 命令附加 `--validate-text` 会逐项对比 Qwen3-VL processor 输入和最终 embedding；TI2V 还会对比首帧预处理结果。附加 `--validate-ti2v-vae` 会对比采样后的 clean condition latent；仅对不含 seed metadata 的旧 capture 使用 `--seed`。 +使用 `--reference-root` 可批量复放所有 Dense DiT/VAE capture。它会在所有样本间复用同一已加载的 Dense transformer 和 VAE；每个样本仍会重新实例化 scheduler，以保持各自的采样配置。 +添加 `--assert-exact` 后，只要任何记录 tensor 的 shape 或值不一致，命令就会以失败状态退出，可作为 CI parity gate，而非仅输出诊断报告。 +checkpoint inspection 工具会执行正常的严格加载，并记录已消费的 config 字段、checkpoint key 覆盖、component/block 参数量、dtype/device 分布、保留的 FP32 参数量及模型显存分配证据。 +基准工具会将一次性 setup 与 warmup、measured 指标分开记录，覆盖 checkpoint load、text encoding、每个 denoising step、VAE、refiner、输出编码与峰值 GPU 显存。启用优化前应先用它建立基线，并将全分辨率测量与 smoke run 分开记录。 +未传入 `--negative-caption` 时，benchmark 会与 pipeline、CLI、service 一样使用上游兼容的 +T2I 或视频负向 caption。只有需要特意评测该语义覆盖时，才显式传入空字符串。 +对于 base+refiner，它还会记录串行的 base release 与 refiner load 阶段。默认 sorted eager +MoE 路径与 source 等价，并有显式诊断 fallback，但不是 grouped-GEMM 或 FP8 生产吞吐后端。 +T2I/T2V 的 Refiner 会复用 base generation 生成的完全相同 CFG text condition，并在报告中 +记录 `refiner_prompt_conditions_reused`；TI2V 保持上游兼容的 text-only Refiner 编码路径。 +Refiner core CLI 会向上游与 TeleFuser 的 low-noise 路径注入相同的 latent、noise、prompt 与 frame-zero condition；它会先 offload 上游 DiT，再加载 TeleFuser DiT，避免两个 30B 模型同时占用 GPU。 +对 MoE 或 Refiner core validator 添加 `--assert-exact`，可强制执行零漂移的 numerical-oracle gate，而非只写出指标。 + +## 依赖与已知限制 + +数值 oracle 路径需要 CUDA、PyTorch、Diffusers、Transformers 以及 checkpoint 的 +`transformer/`、`text_encoder/`、`processor/`、`vae/`、`scheduler/` 组件。Dense +可在单 GPU 上按 source-equivalent 路径运行。当前 MoE expert 是 eager correctness +实现，不应作为 30B 的生产后端。FSDP、Ulysses/CFG 并行、FlashAttention 与 FP8 expert +尚未启用。服务与 CLI 都支持串行单进程的 base+refiner 生命周期；当前不支持多卡对齐。 diff --git a/examples/lingbot_video/README.md b/examples/lingbot_video/README.md new file mode 100644 index 0000000..0ee74ed --- /dev/null +++ b/examples/lingbot_video/README.md @@ -0,0 +1,30 @@ +# LingBot-Video + +Use a structured JSON caption produced by the LingBot rewriter. Dense T2V: + +```bash +python examples/lingbot_video/lingbot_video_generate.py \ + --model-dir /path/to/lingbot-video-dense-1.3b \ + --caption-json /path/to/caption.json --output result.mp4 +``` + +Pass `--image first_frame.png` for TI2V, or `--variant moe` for the MoE base. + +For the MoE checkpoint refiner, use the in-memory base-to-refiner path: + +```bash +python examples/lingbot_video/lingbot_video_generate.py \ + --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ + --caption-json /path/to/caption.json --output result.mp4 +``` + +The command releases base-stage GPU weights before loading the separate refiner. Add `--image first_frame.png` to preserve the TI2V clean frame-zero condition. The default sorted eager MoE route path is source-equivalent with a diagnostic fallback, but it is not a grouped-GEMM or FP8 production-performance backend. + +Serve structured-caption T2I/T2V/TI2V requests with: + +```bash +LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ + telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +``` + +Set `LINGBOT_VIDEO_VARIANT=moe LINGBOT_VIDEO_ENABLE_REFINER=1` for the MoE base-plus-refiner service lifecycle. diff --git a/examples/lingbot_video/lingbot_video_generate.py b/examples/lingbot_video/lingbot_video_generate.py new file mode 100644 index 0000000..030b0b1 --- /dev/null +++ b/examples/lingbot_video/lingbot_video_generate.py @@ -0,0 +1,142 @@ +"""Generate a LingBot-Video T2I, T2V, or TI2V sample with native TeleFuser stages.""" + +from __future__ import annotations + +import argparse + +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + build_lingbot_video_pipeline, + build_lingbot_video_refiner_stage, + default_negative_caption, + load_lingbot_video_prompt, + load_refiner_first_frame, + num_frames_from_duration, + prepare_refiner_video, +) + + +def _image_to_tensor(path: str) -> torch.Tensor: + """Load an RGB image as the raw [0,255] tensor required by TI2V.""" + image = Image.open(path).convert("RGB") + return torch.from_numpy(__import__("numpy").asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).float() + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--model-dir", required=True) + parser.add_argument("--caption-json", required=True, help="Path to the structured JSON caption.") + parser.add_argument("--output", required=True, help="Output MP4 path, or PNG when --num-frames=1.") + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument( + "--num-frames", type=int, help="Override prompt-file duration with an explicit 4n+1 frame count." + ) + 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("--image", help="Optional first-frame image for TI2V.") + parser.add_argument("--variant", choices=("dense", "moe"), default="dense") + parser.add_argument("--negative-caption", default=None, help="Override the official structured negative caption.") + 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) + args = parser.parse_args() + if args.refine and args.variant != "moe": + raise ValueError("--refine requires --variant moe because the refiner is shipped with the MoE checkpoint") + caption_text, 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 = build_lingbot_video_pipeline( + args.model_dir, + variant=args.variant, + guidance_scale=args.guidance_scale, + num_inference_steps=args.steps, + ) + generator = torch.Generator("cuda").manual_seed(args.seed) + generation = pipeline.generate( + LingBotVideoRequest( + caption=caption_text, + height=args.height, + width=args.width, + num_frames=num_frames, + image=source_image, + ), + negative_caption=negative_caption, + generator=generator, + ) + frames = generation.output + if args.refine: + if pipeline.text_stage is None: + raise RuntimeError("LingBot-Video text stage is required for refiner prompt encoding") + # The upstream TI2V refiner consumes text-only prompt conditioning. Its + # frame-zero visual condition is injected through the VAE latent instead. + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = pipeline.text_stage.encode(caption_text) + negative, negative_mask = pipeline.text_stage.encode(negative_caption) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + negative = generation.prompt_conditions.negative_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative_mask = generation.prompt_conditions.negative_attention_mask + pipeline.release_gpu_resources() + refiner = build_lingbot_video_refiner_stage(args.model_dir) + lowres_video, _ = 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 + ) + frames = 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), + ) + 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(args.output) + return + from diffusers.utils import export_to_video + + export_to_video(list(video), args.output, fps=24) + + +if __name__ == "__main__": + main() diff --git a/examples/lingbot_video/lingbot_video_service.py b/examples/lingbot_video/lingbot_video_service.py new file mode 100644 index 0000000..28a53f4 --- /dev/null +++ b/examples/lingbot_video/lingbot_video_service.py @@ -0,0 +1,205 @@ +"""Service entrypoint for LingBot-Video Dense generation.""" + +from __future__ import annotations + +import json +import os + +import numpy as np +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_video import ( + DEFAULT_NEGATIVE_PROMPT, + DEFAULT_NEGATIVE_PROMPT_IMAGE, + LingBotVideoRequest, + build_lingbot_video_pipeline, + build_lingbot_video_refiner_stage, + default_negative_caption, + load_refiner_first_frame, + parse_lingbot_video_prompt, + prepare_refiner_video, +) +from telefuser.utils.video import get_target_video_size_from_ratio + +PPL_CONFIG = { + "model_root": os.environ.get( + "LINGBOT_VIDEO_MODEL_ROOT", "/hhb-data/aigc/model_zoo/lingbot/lingbot-video-dense-1.3b" + ), + "num_inference_steps": 40, + "guidance_scale": 3.0, + "fps": 24, + "variant": os.environ.get("LINGBOT_VIDEO_VARIANT", "dense"), + "enable_refiner": os.environ.get("LINGBOT_VIDEO_ENABLE_REFINER", "0").lower() in {"1", "true", "yes"}, + "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, +} + + +def get_pipeline_contract() -> dict: + """Declare the task contract consumed by TeleFuser service APIs.""" + parameters = { + "seed": {"type": "integer", "default": 42}, + "target_video_length": {"type": "integer", "default": 4}, + "resolution": {"type": "string", "default": "480p"}, + "aspect_ratio": {"type": "string", "default": "16:9"}, + "refine": {"type": "boolean", "default": PPL_CONFIG["enable_refiner"]}, + } + image_parameters = { + **parameters, + "negative_prompt": {"type": "string", "default": DEFAULT_NEGATIVE_PROMPT_IMAGE}, + } + video_parameters = { + **parameters, + "negative_prompt": {"type": "string", "default": DEFAULT_NEGATIVE_PROMPT}, + } + return { + "pipeline_name": "lingbot_video", + "supported_tasks": ["t2i", "t2v", "i2v"], + "supported_media_types": ["image", "video"], + "execution_mode": "serial_single_pipeline", + "effective_max_concurrent_tasks": 1, + "task_contracts": { + "t2i": {"media_type": "image", "parameters": image_parameters}, + "t2v": {"media_type": "video", "parameters": video_parameters}, + "i2v": { + "media_type": "video", + "required_inputs": ["first_image_path"], + "parameters": video_parameters, + }, + }, + } + + +def get_pipeline(parallelism: int = 1, model_root: str = PPL_CONFIG["model_root"]): + """Load a single-GPU source-equivalent LingBot base pipeline.""" + if parallelism != 1: + raise ValueError("LingBot-Video service currently supports parallelism=1") + return build_lingbot_video_pipeline( + model_root, + variant=PPL_CONFIG["variant"], + guidance_scale=PPL_CONFIG["guidance_scale"], + num_inference_steps=PPL_CONFIG["num_inference_steps"], + ) + + +def _load_condition_image(path: str) -> torch.Tensor: + """Load a user-uploaded RGB image as a raw [0,255] TI2V tensor.""" + array = np.asarray(Image.open(path).convert("RGB")).copy() + return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).float() + + +def _num_frames(seconds: int, fps: int) -> int: + """Convert service duration to the nearest valid ``4n+1`` LingBot count.""" + return max(1, 4 * round(seconds * fps / 4) + 1) + + +def _lingbot_video_size(aspect_ratio: str, resolution: str) -> tuple[int, int]: + """Resolve a service resolution for the combined VAE and DiT spatial contract.""" + 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 run_with_file( + pipeline, + prompt: str, + *, + first_image_path: str = "", + negative_prompt: str | None = None, + seed: int = 42, + output_path: str = "output.mp4", + target_video_length: int = 4, + resolution: str = "480p", + aspect_ratio: str = "16:9", + task: str = "t2v", + refine: bool | None = None, + **_: object, +) -> dict[str, str]: + """Generate and encode a T2I, T2V, or TI2V service result.""" + width, height = _lingbot_video_size(aspect_ratio, resolution) + caption, _ = parse_lingbot_video_prompt(json.loads(prompt)) + if task not in {"t2i", "t2v", "i2v"}: + raise ValueError(f"unsupported LingBot-Video service task: {task}") + if task == "i2v" and not first_image_path: + raise ValueError("LingBot-Video i2v requires first_image_path") + refine_enabled = PPL_CONFIG["enable_refiner"] if refine is None else refine + if refine_enabled and pipeline.variant != "moe": + raise ValueError("LingBot-Video refiner requires a pipeline loaded with variant=moe") + num_frames = 1 if task == "t2i" else _num_frames(target_video_length, 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=torch.Generator("cuda").manual_seed(seed), + ) + frames = generation.output + if refine_enabled: + if pipeline.text_stage is None or pipeline.model_dir is None: + raise RuntimeError("LingBot-Video refiner requires text-stage and checkpoint metadata") + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = pipeline.text_stage.encode(caption) + negative, negative_mask = pipeline.text_stage.encode(resolved_negative_prompt) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + negative = generation.prompt_conditions.negative_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative_mask = generation.prompt_conditions.negative_attention_mask + pipeline.release_gpu_resources() + refiner = build_lingbot_video_refiner_stage(pipeline.model_dir) + lowres_video, _ = prepare_refiner_video( + frames, + source_fps=PPL_CONFIG["fps"], + height=PPL_CONFIG["refiner_height"], + width=PPL_CONFIG["refiner_width"], + ) + clean_first_frame = ( + load_refiner_first_frame( + first_image_path, + target_height=PPL_CONFIG["refiner_height"], + target_width=PPL_CONFIG["refiner_width"], + geometry_height=height, + geometry_width=width, + ) + if first_image_path + else None + ) + frames = refiner.refine( + lowres_video, + positive, + negative, + positive_mask, + negative_mask, + num_inference_steps=PPL_CONFIG["refiner_steps"], + guidance_scale=PPL_CONFIG["refiner_guidance_scale"], + shift=PPL_CONFIG["refiner_shift"], + t_thresh=PPL_CONFIG["refiner_t_thresh"], + tail_steps=PPL_CONFIG["refiner_tail_steps"], + clean_first_frame=clean_first_frame, + generator=torch.Generator("cuda").manual_seed(seed), + ) + video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() + if task == "t2i": + Image.fromarray((video[0] * 255).round().astype("uint8")).save(output_path) + return {"output_path": output_path} + from diffusers.utils import export_to_video + + export_to_video(list(video), output_path, fps=PPL_CONFIG["fps"]) + return {"output_path": output_path} diff --git a/telefuser/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py new file mode 100644 index 0000000..7f6c5d6 --- /dev/null +++ b/telefuser/models/lingbot_video_dit.py @@ -0,0 +1,338 @@ +"""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.nn.functional as F +from torch import nn + +from telefuser.ops.attention import attention + + +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 + + def set_attention_config(self, attention_config: object) -> None: + """Attach a runtime-selected TeleFuser attention implementation.""" + self.attention_config = attention_config + + def forward( + self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor, attention_mask: torch.Tensor | 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) + output = attention( + query.transpose(1, 2), + key.transpose(1, 2), + value.transpose(1, 2), + attention_config=self.attention_config, + attn_mask=attention_mask, + input_layout="BNSD", + output_layout="BNSD", + ) + if not isinstance(output, torch.Tensor): + raise RuntimeError("LingBot attention does not support log-sum-exp outputs") + return self.to_out(output.transpose(1, 2).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, + ) -> 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) + 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) + + 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) + + 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) + text = self.text_embedder(encoder_hidden_states) + joint = torch.cat((video, text), dim=1) + text_length = text.shape[1] + positions = make_lingbot_video_joint_position_ids(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) + .expand(batch, -1, -1) + ) + time_embedding = self.time_embedder(timestep) + temb_input = time_embedding.unsqueeze(1).expand(-1, joint.shape[1], -1) + temb6 = self.time_modulation(temb_input) + attention_mask = None + if encoder_attention_mask is not None and not bool(encoder_attention_mask.bool().all()): + video_mask = torch.ones(batch, video_tokens, dtype=torch.bool, device=hidden_states.device) + attention_mask = torch.cat((video_mask, encoder_attention_mask.bool()), dim=1)[:, None, None, :] + for block in self.blocks: + joint = block(joint, temb6, rotary, attention_mask) + 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)) + 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) diff --git a/telefuser/models/lingbot_video_moe.py b/telefuser/models/lingbot_video_moe.py new file mode 100644 index 0000000..22a07c6 --- /dev/null +++ b/telefuser/models/lingbot_video_moe.py @@ -0,0 +1,249 @@ +"""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.ops.moe import route_topk + +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.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 {"sorted", "where"}: + raise ValueError("execution backend must be 'sorted' or 'where'") + self.execution_backend = backend + + def forward(self, tokens: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + 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) diff --git a/telefuser/ops/__init__.py b/telefuser/ops/__init__.py index 10d060f..d99594b 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 route_topk from .normalization import AdaLayerNormContinuous, LayerNorm, RMSNorm, fused_scale_shift, modulate from .rotary import apply_rotary_emb @@ -28,6 +29,7 @@ "AdaLayerNormContinuous", "fused_scale_shift", "modulate", + "route_topk", # Rotary "apply_rotary_emb", ] diff --git a/telefuser/ops/moe.py b/telefuser/ops/moe.py new file mode 100644 index 0000000..d5069de --- /dev/null +++ b/telefuser/ops/moe.py @@ -0,0 +1,49 @@ +"""Correctness-first Mixture-of-Experts routing primitives.""" + +from __future__ import annotations + +import torch + + +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).values.sum(dim=-1) + selected_groups = group_values.topk(top_k_groups, dim=-1).indices + group_mask = torch.zeros_like(grouped, dtype=torch.bool) + group_mask.scatter_(1, selected_groups.unsqueeze(-1).expand(-1, -1, experts_per_group), True) + masked = selection_scores.masked_fill(~group_mask.reshape(tokens, experts), float("-inf")) + indices = masked.topk(top_k, dim=-1).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 diff --git a/telefuser/pipelines/lingbot_video/__init__.py b/telefuser/pipelines/lingbot_video/__init__.py new file mode 100644 index 0000000..48892c8 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/__init__.py @@ -0,0 +1,86 @@ +"""LingBot-Video pipeline contracts and runtime 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 .loading import checkpoint_key_coverage, load_lingbot_video_dense_transformer, load_lingbot_video_moe_transformer +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 .runtime import build_lingbot_video_pipeline, build_lingbot_video_refiner_stage +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", + "build_lingbot_video_pipeline", + "build_lingbot_video_refiner_stage", + "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", + "checkpoint_key_coverage", + "load_lingbot_video_dense_transformer", + "load_lingbot_video_moe_transformer", + "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..05958ca --- /dev/null +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -0,0 +1,109 @@ +"""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 + +import torch + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig +from telefuser.models.lingbot_video_dit import LingBotVideoTransformer3DModel + + +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 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, transformer: LingBotVideoTransformer3DModel, model_runtime_config: ModelRuntimeConfig + ) -> None: + super().__init__(name, model_runtime_config) + self.transformer = transformer + 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"] + + @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.""" + 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): + positive = self.transformer( + latents, model_timestep, positive_prompt_embeds, positive_attention_mask + ).float() + if guidance_scale == 1.0: + return positive + if negative_prompt_embeds is None: + raise ValueError("negative_prompt_embeds is required when guidance_scale is not 1") + negative = self.transformer( + latents, model_timestep, negative_prompt_embeds, negative_attention_mask + ).float() + return negative + guidance_scale * (positive - negative) diff --git a/telefuser/pipelines/lingbot_video/loading.py b/telefuser/pipelines/lingbot_video/loading.py new file mode 100644 index 0000000..f5e1985 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/loading.py @@ -0,0 +1,172 @@ +"""Checkpoint loading diagnostics for LingBot-Video components.""" + +from __future__ import annotations + +import json +from collections.abc import Collection, Mapping +from pathlib import Path +from typing import Any + +import torch +from torch import nn + +from telefuser.models.lingbot_video_dit import LingBotVideoTransformer3DModel +from telefuser.models.lingbot_video_moe import LingBotVideoMoeTransformer3DModel + + +def checkpoint_key_coverage( + module: nn.Module, state_dict: Mapping[str, torch.Tensor] | Collection[str] +) -> dict[str, Any]: + """Return exact checkpoint key coverage without mutating the module.""" + expected = set(module.state_dict()) + available = set(state_dict) + matched = expected & available + return { + "expected_key_count": len(expected), + "checkpoint_key_count": len(available), + "matched_key_count": len(matched), + "coverage": len(matched) / len(expected) if expected else 1.0, + "missing_keys": sorted(expected - available), + "unexpected_keys": sorted(available - expected), + "matched_numel": sum(module.state_dict()[name].numel() for name in matched), + } + + +def load_lingbot_video_dense_transformer( + checkpoint_dir: str | Path, *, device: torch.device | str = "cuda", torch_dtype: torch.dtype = torch.bfloat16 +) -> "LingBotVideoTransformer3DModel": + """Strictly load the official Diffusers Dense transformer checkpoint.""" + + from safetensors.torch import load_model + + directory = Path(checkpoint_dir) + config = json.loads((directory / "config.json").read_text(encoding="utf-8")) + 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", + ) + transformer = LingBotVideoTransformer3DModel(**{key: config[key] for key in keys}).to( + device=device, dtype=torch_dtype + ) + 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 transformer.named_modules(): + if any(part in fp32_names for part in name.split(".")): + module.float() + for name, parameter in transformer.named_parameters(): + if any(part in fp32_names for part in name.split(".")): + parameter.data = parameter.data.float() + + missing, unexpected = load_model( + transformer, directory / "diffusion_pytorch_model.safetensors", strict=True, device=str(device) + ) + if missing or unexpected: + raise RuntimeError(f"LingBot checkpoint mismatch: missing={missing}, unexpected={unexpected}") + return transformer.eval() + + +def load_lingbot_video_moe_transformer( + checkpoint_dir: str | Path, *, device: torch.device | str = "cuda", torch_dtype: torch.dtype = torch.bfloat16 +) -> LingBotVideoMoeTransformer3DModel: + """Strictly load the official sharded MoE/refiner transformer checkpoint.""" + from safetensors.torch import load_file + + directory = Path(checkpoint_dir) + config = json.loads((directory / "config.json").read_text(encoding="utf-8")) + 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", + "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", + ) + transformer = LingBotVideoMoeTransformer3DModel(**{key: config[key] for key in keys}).to( + device=device, dtype=torch_dtype + ) + 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 transformer.named_modules(): + if any(part in fp32_names for part in name.split(".")): + module.float() + for name, parameter in transformer.named_parameters(): + if any(part in fp32_names for part in name.split(".")): + parameter.data = parameter.data.float() + + index = json.loads((directory / "diffusion_pytorch_model.safetensors.index.json").read_text(encoding="utf-8"))[ + "weight_map" + ] + expected = set(transformer.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)) + shard_unexpected = transformer.load_state_dict(shard_state, strict=False).unexpected_keys + unexpected.update(shard_unexpected) + found.update(shard_state) + del shard_state + missing = expected - found + if missing or unexpected: + raise RuntimeError( + f"LingBot MoE checkpoint mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}" + ) + return transformer.eval() diff --git a/telefuser/pipelines/lingbot_video/pipeline.py b/telefuser/pipelines/lingbot_video/pipeline.py new file mode 100644 index 0000000..e8d6c34 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/pipeline.py @@ -0,0 +1,198 @@ +"""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 + +import torch + +from telefuser.core.base_pipeline import BasePipeline +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler + +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(frozen=True) +class LingBotVideoPipelineConfig: + """Runtime configuration shared by Dense and MoE pipelines.""" + + model: LingBotVideoModelConfig = LingBotVideoModelConfig() + guidance_scale: float = 3.0 + num_inference_steps: int = 50 + shift: float = 3.0 + + 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 + positive_attention_mask: torch.Tensor + negative_attention_mask: torch.Tensor + 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): + """Lifecycle-aware facade for independently loaded LingBot-Video stages.""" + + 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.model_dir: str | None = None + self.variant: str = "dense" + + def init(self, module_manager: object, config: LingBotVideoPipelineConfig) -> None: + del module_manager + self.config = config + + @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 set_runtime( + self, + *, + text_stage: LingBotVideoTextEncodingStage, + denoising_stage: LingBotVideoDenoisingStage, + scheduler: FlowUniPCMultistepScheduler, + vae_decode_stage: LingBotVideoVAEDecodeStage | None = None, + vae_encode_stage: LingBotVideoVAEEncodeStage | None = None, + ) -> None: + """Attach independently loaded stages to this pipeline instance.""" + self.text_stage = text_stage + self.denoising_stage = denoising_stage + self.scheduler = scheduler + self.vae_decode_stage = vae_decode_stage + self.vae_encode_stage = vae_encode_stage + + 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): + 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 __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) + 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) + negative, negative_mask = self.text_stage.encode(negative_caption, images=vision_images) + 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, + ) + 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.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") + return LingBotVideoGeneration(self.vae_decode_stage.decode(latent), prompt_conditions) diff --git a/telefuser/pipelines/lingbot_video/refiner.py b/telefuser/pipelines/lingbot_video/refiner.py new file mode 100644 index 0000000..9216922 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/refiner.py @@ -0,0 +1,278 @@ +"""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 pathlib import Path + +import numpy as np +import torch +import torch.nn.functional as F + +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from telefuser.pipelines.lingbot_video.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 + + 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, + ) -> 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) + 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.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) + return self.vae_decode_stage.decode(latent) diff --git a/telefuser/pipelines/lingbot_video/runtime.py b/telefuser/pipelines/lingbot_video/runtime.py new file mode 100644 index 0000000..31c35bd --- /dev/null +++ b/telefuser/pipelines/lingbot_video/runtime.py @@ -0,0 +1,129 @@ +"""Native checkpoint assembly for the LingBot-Video pipeline.""" + +from __future__ import annotations + +from pathlib import Path + +import torch + +from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, OffloadConfig, WeightOffloadType +from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler + +from .data import load_lingbot_video_model_config +from .denoising import LingBotVideoDenoisingStage +from .loading import load_lingbot_video_dense_transformer, load_lingbot_video_moe_transformer +from .pipeline import LingBotVideoPipeline, LingBotVideoPipelineConfig +from .refiner import LingBotVideoRefinerStage +from .text_encoding import LingBotVideoTextEncodingStage +from .vae import LingBotVideoVAEDecodeStage, LingBotVideoVAEEncodeStage + + +def build_lingbot_video_pipeline( + model_dir: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + variant: str = "dense", + cpu_offload: bool | None = None, + guidance_scale: float = 3.0, + num_inference_steps: int = 50, + shift: float = 3.0, + attention_config: AttentionConfig | None = None, +) -> LingBotVideoPipeline: + """Load official checkpoint components and assemble a native pipeline runtime. + + ``variant`` selects a Dense transformer or the MoE base transformer. The + refiner remains a separately loaded stage so its weights do not need to + coexist with the base model. + """ + if cpu_offload is None: + cpu_offload = variant == "moe" + if variant not in {"dense", "moe"}: + raise ValueError("variant must be 'dense' or 'moe'") + try: + from diffusers import AutoencoderKLWan + from transformers import AutoProcessor, Qwen3VLForConditionalGeneration + except ImportError as exc: + raise RuntimeError("LingBot-Video runtime requires diffusers and transformers") from exc + + root = Path(model_dir) + transformer_dir = root / "transformer" + runtime_config = ModelRuntimeConfig( + device_type=torch.device(device).type, + torch_dtype=torch_dtype, + attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + offload_config=OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD + ), + ) + load_device = "cpu" if cpu_offload else device + transformer = ( + load_lingbot_video_dense_transformer(transformer_dir, device=load_device, torch_dtype=torch_dtype) + if variant == "dense" + else load_lingbot_video_moe_transformer(transformer_dir, device=load_device, torch_dtype=torch_dtype) + ) + processor = AutoProcessor.from_pretrained(root / "processor") + text_encoder = ( + Qwen3VLForConditionalGeneration.from_pretrained( + root / "text_encoder", dtype=torch_dtype, attn_implementation="sdpa" + ) + .to(load_device) + .eval() + ) + vae = AutoencoderKLWan.from_pretrained(root / "vae", torch_dtype=torch.float32).to(load_device).eval() + scheduler = FlowUniPCMultistepScheduler.from_pretrained(root / "scheduler") + pipeline = LingBotVideoPipeline(device=device, torch_dtype=torch_dtype) + pipeline.model_dir = str(root) + pipeline.variant = variant + pipeline.init( + None, + LingBotVideoPipelineConfig( + model=load_lingbot_video_model_config(transformer_dir, variant=variant), + guidance_scale=guidance_scale, + num_inference_steps=num_inference_steps, + shift=shift, + ), + ) + pipeline.set_runtime( + text_stage=LingBotVideoTextEncodingStage("text_encoder", text_encoder, processor, runtime_config), + denoising_stage=LingBotVideoDenoisingStage("transformer", transformer, runtime_config), + vae_encode_stage=LingBotVideoVAEEncodeStage("vae_encode", vae, runtime_config), + vae_decode_stage=LingBotVideoVAEDecodeStage("vae_decode", vae, runtime_config), + scheduler=scheduler, + ) + return pipeline + + +def build_lingbot_video_refiner_stage( + model_dir: str | Path, + *, + device: str = "cuda", + torch_dtype: torch.dtype = torch.bfloat16, + cpu_offload: bool = True, + attention_config: AttentionConfig | None = None, +) -> LingBotVideoRefinerStage: + """Load the separately checkpointed MoE refiner and its VAE/scheduler stages.""" + try: + from diffusers import AutoencoderKLWan + except ImportError as exc: + raise RuntimeError("LingBot-Video refiner runtime requires diffusers") from exc + + root = Path(model_dir) + runtime_config = ModelRuntimeConfig( + device_type=torch.device(device).type, + torch_dtype=torch_dtype, + attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), + offload_config=OffloadConfig( + offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD + ), + ) + load_device = "cpu" if cpu_offload else device + transformer = load_lingbot_video_moe_transformer(root / "refiner", device=load_device, torch_dtype=torch_dtype) + vae = AutoencoderKLWan.from_pretrained(root / "vae", torch_dtype=torch.float32).to(load_device).eval() + scheduler = FlowUniPCMultistepScheduler.from_pretrained(root / "scheduler") + return LingBotVideoRefinerStage( + denoising_stage=LingBotVideoDenoisingStage("refiner", transformer, runtime_config), + vae_encode_stage=LingBotVideoVAEEncodeStage("refiner_vae_encode", vae, runtime_config), + vae_decode_stage=LingBotVideoVAEDecodeStage("refiner_vae_decode", vae, runtime_config), + scheduler=scheduler, + ) diff --git a/telefuser/pipelines/lingbot_video/text_encoding.py b/telefuser/pipelines/lingbot_video/text_encoding.py new file mode 100644 index 0000000..fa00284 --- /dev/null +++ b/telefuser/pipelines/lingbot_video/text_encoding.py @@ -0,0 +1,141 @@ +"""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 torch import nn + +from telefuser.core.base_stage import BaseStage, with_model_offload +from telefuser.core.config import ModelRuntimeConfig + +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, + text_encoder: nn.Module, + processor: Any, + 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 = text_encoder + self.processor = 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..528ac9e --- /dev/null +++ b/telefuser/pipelines/lingbot_video/vae.py @@ -0,0 +1,102 @@ +"""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 + + +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, vae: torch.nn.Module, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.vae = 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, vae: torch.nn.Module, model_runtime_config: ModelRuntimeConfig) -> None: + super().__init__(name, model_runtime_config) + self.vae = 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/tests/unit/models/test_lingbot_video_dit.py b/tests/unit/models/test_lingbot_video_dit.py new file mode 100644 index 0000000..845d3b2 --- /dev/null +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -0,0 +1,68 @@ +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)) 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..b8ff12b --- /dev/null +++ b/tests/unit/models/test_lingbot_video_moe.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +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) 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..46e3b80 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -0,0 +1,183 @@ +from __future__ import annotations + +import torch +import torch.nn as nn + +from telefuser.core.config import ModelRuntimeConfig +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 _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_pipeline_runs_t2v_sampling_with_injected_runtime() -> None: + pipeline = LingBotVideoPipeline(device="cpu") + pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) + pipeline.set_runtime(text_stage=_TextStage(), denoising_stage=_DenoisingStage(), 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, + ) + + +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", _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)) + + +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 = LingBotVideoPipeline(device="cpu") + pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) + pipeline.set_runtime(text_stage=text_stage, denoising_stage=_DenoisingStage(), scheduler=_Scheduler()) + + pipeline.generate(LingBotVideoRequest(caption="{}", height=16, width=16, num_frames=5), decode=False) + + assert text_stage.captions == ["{}", DEFAULT_NEGATIVE_PROMPT] + + +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 = LingBotVideoPipeline(device="cpu") + pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) + pipeline.set_runtime( + text_stage=text_stage, + denoising_stage=_DenoisingStage(), + scheduler=_Scheduler(), + 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 = LingBotVideoPipeline(device="cpu") + pipeline.set_runtime( + text_stage=text_stage, + denoising_stage=denoising_stage, + scheduler=_Scheduler(), + vae_encode_stage=vae_encode_stage, + 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 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..88ef8ab --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_refiner.py @@ -0,0 +1,159 @@ +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 + + +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") + path = tmp_path / "base.mp4" + from diffusers.utils import export_to_video + + 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) + upstream_root = Path("work_dirs/lingbot-video-master") + 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 + + image = np.arange(9 * 15 * 3, dtype=np.uint8).reshape(9, 15, 3) + path = tmp_path / "first.png" + Image.fromarray(image).save(path) + upstream_root = Path("work_dirs/lingbot-video-master") + 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..a5f56a5 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py @@ -0,0 +1,89 @@ +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 _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)) 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..7218eb2 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -0,0 +1,88 @@ +"""Tests for LingBot checkpoint assembly and memory lifecycle defaults.""" + +from __future__ import annotations + +import sys +import types + +from telefuser.core.config import WeightOffloadType +from telefuser.pipelines.lingbot_video import runtime + + +class _FakeModule: + def __init__(self) -> None: + self.device: 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 test_moe_runtime_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: + loaded_transformers: list[tuple[object, str]] = [] + + def load_dense(path, *, device: str, torch_dtype): + del torch_dtype + loaded_transformers.append((path, device)) + return _FakeModule() + + def load_moe(path, *, device: str, torch_dtype): + del torch_dtype + loaded_transformers.append((path, device)) + return _FakeModule() + + 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(runtime, "load_lingbot_video_dense_transformer", load_dense) + monkeypatch.setattr(runtime, "load_lingbot_video_moe_transformer", load_moe) + monkeypatch.setattr(runtime, "load_lingbot_video_model_config", lambda *args, **kwargs: None) + monkeypatch.setattr(runtime.FlowUniPCMultistepScheduler, "from_pretrained", _Scheduler.from_pretrained) + + base = runtime.build_lingbot_video_pipeline(tmp_path, device="cuda", variant="moe") + refiner = runtime.build_lingbot_video_refiner_stage(tmp_path, device="cuda") + + assert loaded_transformers == [(tmp_path / "transformer", "cpu"), (tmp_path / "refiner", "cpu")] + assert base.model_dir == str(tmp_path) + assert base.variant == "moe" + 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..bd1d306 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -0,0 +1,234 @@ +"""Tests for the LingBot-Video service entrypoint contract.""" + +from __future__ import annotations + +import importlib.util +import json +import sys +from pathlib import Path +from types import SimpleNamespace + +import numpy as np +import pytest +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_video import DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT_IMAGE +from telefuser.service.core.pipeline_contract import load_pipeline_contract + + +def _load_service_module(): + path = Path("examples/lingbot_video/lingbot_video_service.py") + spec = importlib.util.spec_from_file_location("lingbot_video_service_test", path) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_service_contract_declares_t2i_t2v_and_i2v() -> None: + module = _load_service_module() + + contract, is_explicit = load_pipeline_contract(module, ppl_file="lingbot_video_service.py", default_task="t2v") + + assert is_explicit + assert contract.supported_tasks == ("t2i", "t2v", "i2v") + assert contract.get_task_contract("t2i").media_type == "image" + assert contract.get_task_contract("t2i").parameters["negative_prompt"].default == DEFAULT_NEGATIVE_PROMPT_IMAGE + assert contract.get_task_contract("t2v").parameters["refine"].type == "boolean" + 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_service_writes_single_image_for_t2i(tmp_path: Path) -> None: + module = _load_service_module() + + class FakePipeline: + variant = "dense" + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.num_frames == 1 + assert request.caption == '{"scene":"test"}' + assert (request.height, request.width) == (480, 864) + return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2)) + + output_path = tmp_path / "result.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() + + +def test_service_writes_video_for_t2v(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_service_module() + exported: list[object] = [] + + import diffusers.utils + + def capture_export(frames, output_path: str, fps: int) -> str: + del output_path, fps + exported.extend(frames) + return "unused.mp4" + + monkeypatch.setattr(diffusers.utils, "export_to_video", capture_export) + + class FakePipeline: + variant = "dense" + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.num_frames == 97 + assert (request.height, request.width) == (480, 864) + return SimpleNamespace(output=torch.full((1, 3, 5, 2, 2), 0.5)) + + output_path = tmp_path / "result.mp4" + result = module.run_with_file( + FakePipeline(), + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2v", + output_path=str(output_path), + ) + + assert result == {"output_path": str(output_path)} + assert len(exported) == 5 + assert exported[0].dtype == np.float32 + assert float(exported[0][0, 0, 0]) == 0.5 + + +def test_service_resolution_is_aligned_to_lingbot_vae_and_dit() -> None: + module = _load_service_module() + + for resolution in ("480p", "720p", "1080p", "2k", "4k"): + width, height = module._lingbot_video_size("16:9", resolution) + assert width % 16 == 0 + assert height % 16 == 0 + + +def test_service_rejects_i2v_without_first_image_path() -> None: + module = _load_service_module() + + with pytest.raises(ValueError, match="requires first_image_path"): + module.run_with_file( + object(), + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="i2v", + ) + + +def test_service_runs_base_then_refiner_on_one_pipeline(tmp_path: Path, monkeypatch) -> None: + module = _load_service_module() + + class FakeTextStage: + def __init__(self) -> None: + self.calls = 0 + + def encode(self, caption: str) -> tuple[torch.Tensor, torch.Tensor]: + self.calls += 1 + assert caption in {'{"scene":"test"}', DEFAULT_NEGATIVE_PROMPT_IMAGE} + return torch.zeros(1, 1, 2), torch.ones(1, 1, dtype=torch.bool) + + class FakePipeline: + variant = "moe" + model_dir = "fake-model" + text_stage = FakeTextStage() + + def __init__(self) -> None: + self.released = False + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.num_frames == 1 + conditions = SimpleNamespace( + has_visual_condition=False, + positive_prompt_embeds=torch.zeros(1, 1, 2), + negative_prompt_embeds=torch.zeros(1, 1, 2), + positive_attention_mask=torch.ones(1, 1, dtype=torch.bool), + negative_attention_mask=torch.ones(1, 1, dtype=torch.bool), + ) + return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2), prompt_conditions=conditions) + + def release_gpu_resources(self) -> None: + self.released = True + + class FakeRefiner: + def __init__(self) -> None: + self.called = False + + def refine(self, lowres_video, *args, **kwargs) -> torch.Tensor: + del args, kwargs + self.called = True + assert lowres_video.shape == (1, 3, 1, 2, 2) + return lowres_video + + pipeline = FakePipeline() + refiner = FakeRefiner() + monkeypatch.setattr(module, "build_lingbot_video_refiner_stage", lambda _: refiner) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_height", 2) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_width", 2) + + output_path = tmp_path / "refined.png" + result = module.run_with_file( + pipeline, + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="t2i", + refine=True, + output_path=str(output_path), + ) + + assert result == {"output_path": str(output_path)} + assert pipeline.released + assert refiner.called + assert pipeline.text_stage.calls == 0 + assert output_path.is_file() + + +def test_service_reencodes_text_only_conditions_for_ti2v_refiner(tmp_path: Path, monkeypatch) -> None: + module = _load_service_module() + + class FakeTextStage: + def __init__(self) -> None: + self.calls = 0 + + def encode(self, caption: str) -> tuple[torch.Tensor, torch.Tensor]: + self.calls += 1 + assert caption in {'{"scene":"test"}', DEFAULT_NEGATIVE_PROMPT} + return torch.zeros(1, 1, 2), torch.ones(1, 1, dtype=torch.bool) + + class FakePipeline: + variant = "moe" + model_dir = "fake-model" + text_stage = FakeTextStage() + + def generate(self, request, **_: object) -> SimpleNamespace: + assert request.image is not None + conditions = SimpleNamespace(has_visual_condition=True) + return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2), prompt_conditions=conditions) + + def release_gpu_resources(self) -> None: + return None + + class FakeRefiner: + def refine(self, lowres_video, *args, **kwargs) -> torch.Tensor: + del args, kwargs + return lowres_video + + monkeypatch.setattr(module, "build_lingbot_video_refiner_stage", lambda _: FakeRefiner()) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_height", 2) + monkeypatch.setitem(module.PPL_CONFIG, "refiner_width", 2) + image_path = tmp_path / "first.png" + Image.new("RGB", (2, 2)).save(image_path) + + module.run_with_file( + FakePipeline(), + json.dumps({"caption": {"scene": "test"}, "duration": 5}), + task="i2v", + first_image_path=str(image_path), + refine=True, + output_path=str(tmp_path / "refined.mp4"), + ) + + assert FakePipeline.text_stage.calls == 2 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..1deed0c --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_text_encoding.py @@ -0,0 +1,25 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch + +from telefuser.core.config import ModelRuntimeConfig +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: + stage = LingBotVideoTextEncodingStage( + "text", + _VisionEncoder(), + processor=object(), + model_runtime_config=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..0ed6af5 --- /dev/null +++ b/tests/unit/pipelines/lingbot_video/test_vae.py @@ -0,0 +1,55 @@ +from __future__ import annotations + +from types import SimpleNamespace + +import torch +from torch import nn + +from telefuser.core.config import ModelRuntimeConfig +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() + encoded = LingBotVideoVAEEncodeStage("encode", vae, runtime).encode(torch.ones(1, 3, 1, 2, 2)) + decoded = LingBotVideoVAEDecodeStage("decode", vae, 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/tests/unit/tools/test_benchmark_lingbot_video.py b/tests/unit/tools/test_benchmark_lingbot_video.py new file mode 100644 index 0000000..936d43a --- /dev/null +++ b/tests/unit/tools/test_benchmark_lingbot_video.py @@ -0,0 +1,38 @@ +from __future__ import annotations + +from collections import defaultdict + +from telefuser.pipelines.lingbot_video import DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT_IMAGE +from tools.validation.benchmark_lingbot_video import _instrument_method, _resolve_negative_caption, _summary + + +def test_summary_reports_each_metric_count_and_aggregate_timings() -> None: + summary = _summary({"base_total": [1.0, 3.0], "vae_decode": [0.5]}) + + assert summary["base_total"] == {"count": 2, "total_seconds": 4.0, "mean_seconds": 2.0, "max_seconds": 3.0} + assert summary["vae_decode"] == {"count": 1, "total_seconds": 0.5, "mean_seconds": 0.5, "max_seconds": 0.5} + + +def test_instrumented_method_records_into_the_active_benchmark_phase() -> None: + class Owner: + def run(self) -> str: + return "ok" + + first_phase: dict[str, list[float]] = defaultdict(list) + second_phase: dict[str, list[float]] = defaultdict(list) + active_phase = first_phase + owner = Owner() + _instrument_method(lambda: active_phase, owner, "run", "operation") + + assert owner.run() == "ok" + active_phase = second_phase + assert owner.run() == "ok" + + assert len(first_phase["operation"]) == 1 + assert len(second_phase["operation"]) == 1 + + +def test_benchmark_uses_source_negative_caption_when_not_overridden() -> None: + assert _resolve_negative_caption(None, 1) == DEFAULT_NEGATIVE_PROMPT_IMAGE + assert _resolve_negative_caption(None, 5) == DEFAULT_NEGATIVE_PROMPT + assert _resolve_negative_caption("", 5) == "" diff --git a/tests/unit/tools/test_capture_lingbot_video_reference.py b/tests/unit/tools/test_capture_lingbot_video_reference.py new file mode 100644 index 0000000..944fc93 --- /dev/null +++ b/tests/unit/tools/test_capture_lingbot_video_reference.py @@ -0,0 +1,59 @@ +from __future__ import annotations + +import json +import subprocess +from pathlib import Path + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +TOOL = REPOSITORY_ROOT / "tools" / "validation" / "capture_lingbot_video_reference.py" + + +def test_lingbot_video_reference_capture_dry_run_writes_manifest(tmp_path: Path) -> None: + output_dir = tmp_path / "reference" + completed = subprocess.run( + [ + "python", + str(TOOL), + "--dry-run", + "--output-dir", + str(output_dir), + "--mode", + "t2i", + "--case", + "example_1", + ], + check=True, + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + ) + + manifest = json.loads((output_dir / "capture_manifest.json").read_text(encoding="utf-8")) + assert manifest["dry_run"] is True + assert manifest["modes"] == ["t2i"] + assert manifest["case"] == "example_1" + assert '"dry_run": true' in completed.stdout + + +def test_lingbot_video_reference_capture_dry_run_selects_all_mode_cases(tmp_path: Path) -> None: + output_dir = tmp_path / "all-cases" + subprocess.run( + [ + "python", + str(TOOL), + "--dry-run", + "--output-dir", + str(output_dir), + "--mode", + "t2i", + "--all-cases", + ], + check=True, + cwd=REPOSITORY_ROOT, + capture_output=True, + text=True, + ) + + manifest = json.loads((output_dir / "capture_manifest.json").read_text(encoding="utf-8")) + assert manifest["all_cases"] is True + assert [item["name"] for item in manifest["selected_cases"]] == [f"example_{index}" for index in range(1, 7)] diff --git a/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py b/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py new file mode 100644 index 0000000..ee1cd25 --- /dev/null +++ b/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py @@ -0,0 +1,54 @@ +from __future__ import annotations + +from pathlib import Path + +import torch + +from tools.validation.inspect_lingbot_video_checkpoint import build_load_report + + +class _TinyModel(torch.nn.Module): + def __init__(self) -> None: + super().__init__() + self.blocks = torch.nn.ModuleList((torch.nn.Linear(2, 3), torch.nn.Linear(3, 2))) + self.head = torch.nn.Linear(2, 1, bias=False) + + +def test_load_report_covers_keys_and_parameter_layout() -> None: + model = _TinyModel() + keys = set(model.state_dict()) + config = { + "patch_size": [1, 2, 2], + "in_channels": 16, + "out_channels": 16, + "hidden_size": 8, + "num_attention_heads": 2, + "depth": 2, + "intermediate_size": 12, + "text_dim": 10, + "freq_dim": 4, + "norm_eps": 1e-6, + "rope_theta": 10000.0, + "axes_dims": [2, 2, 4], + "qkv_bias": True, + "out_bias": True, + "patch_embed_bias": True, + "timestep_mlp_bias": True, + } + + report = build_load_report( + model, + checkpoint_dir=Path("checkpoint"), + variant="dense", + config=config, + available_keys=keys, + ) + + coverage = report["checkpoint_key_coverage"] + parameters = report["parameters"] + assert coverage["coverage"] == 1.0 + assert coverage["missing_keys"] == [] + assert coverage["unexpected_keys"] == [] + assert parameters["total_parameter_count"] == sum(parameter.numel() for parameter in model.parameters()) + assert parameters["parameter_count_by_block"] == {"0": 9, "1": 8} + assert parameters["parameter_count_by_component"]["head"] == 2 diff --git a/tests/unit/tools/test_replay_lingbot_video_dense_reference.py b/tests/unit/tools/test_replay_lingbot_video_dense_reference.py new file mode 100644 index 0000000..576c609 --- /dev/null +++ b/tests/unit/tools/test_replay_lingbot_video_dense_reference.py @@ -0,0 +1,86 @@ +"""Unit tests for the capture-driven Dense replay validator.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +import pytest +import torch + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +TOOL = REPOSITORY_ROOT / "tools" / "validation" / "replay_lingbot_video_dense_reference.py" + + +def _load_tool_module(): + spec = importlib.util.spec_from_file_location("replay_lingbot_video_dense_reference_test", TOOL) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_tensor_metrics_reports_exact_replay() -> None: + module = _load_tool_module() + tensor = torch.tensor([1.0, -2.0]) + + metrics = module.tensor_metrics(tensor, tensor.clone()) + + assert metrics["shape_match"] + assert metrics["max_abs"] == 0.0 + assert metrics["mean_abs"] == 0.0 + assert metrics["relative_l2"] == 0.0 + assert metrics["cosine"] == pytest.approx(1.0) + assert metrics["exact_mismatch_count"] == 0 + + +def test_tensor_metrics_reports_shape_mismatch() -> None: + module = _load_tool_module() + + metrics = module.tensor_metrics(torch.ones(2), torch.ones(3)) + + assert metrics == {"shape_match": False, "reference_numel": 2, "candidate_numel": 3} + + +def test_discover_reference_dirs_returns_sorted_capture_runs(tmp_path: Path) -> None: + module = _load_tool_module() + late = tmp_path / "t2v" / "example_2" / "run-00" + early = tmp_path / "t2i" / "example_1" / "run-00" + late.mkdir(parents=True) + early.mkdir(parents=True) + (late / "metadata.json").write_text("{}", encoding="utf-8") + (early / "metadata.json").write_text("{}", encoding="utf-8") + + references = module.discover_reference_dirs(tmp_path) + + assert references == [early, late] + + +def test_discover_reference_dirs_rejects_empty_root(tmp_path: Path) -> None: + module = _load_tool_module() + + with pytest.raises(FileNotFoundError, match="No capture metadata"): + module.discover_reference_dirs(tmp_path) + + +def test_exact_replay_failures_accepts_exact_metrics() -> None: + module = _load_tool_module() + + assert ( + module.exact_replay_failures({"metrics": {"decoded_frames": {"shape_match": True, "exact_mismatch_count": 0}}}) + == [] + ) + + +def test_exact_replay_failures_reports_shape_and_value_mismatches() -> None: + module = _load_tool_module() + report = { + "reports": [ + {"metrics": {"shape": {"shape_match": False}}}, + {"metrics": {"values": {"shape_match": True, "exact_mismatch_count": 1}}}, + ] + } + + assert module.exact_replay_failures(report) == ["report[0].shape", "report[1].values"] diff --git a/tests/unit/tools/test_run_lingbot_video_moe_parity.py b/tests/unit/tools/test_run_lingbot_video_moe_parity.py new file mode 100644 index 0000000..bac7e75 --- /dev/null +++ b/tests/unit/tools/test_run_lingbot_video_moe_parity.py @@ -0,0 +1,34 @@ +"""Unit tests for strict MoE oracle-parity verdicts.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +TOOL = REPOSITORY_ROOT / "tools" / "validation" / "run_lingbot_video_moe_parity.py" + + +def _load_tool_module(): + spec = importlib.util.spec_from_file_location("run_lingbot_video_moe_parity_test", TOOL) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_exact_parity_failures_accepts_zero_drift() -> None: + module = _load_tool_module() + + assert module.exact_parity_failures({"max_abs": 0.0, "mean_abs": 0.0, "relative_l2": 0.0, "cosine": 1.0}) == [] + + +def test_exact_parity_failures_reports_drift() -> None: + module = _load_tool_module() + + assert module.exact_parity_failures({"max_abs": 0.25, "mean_abs": 0.0, "relative_l2": 0.125, "cosine": 0.99}) == [ + "max_abs", + "relative_l2", + ] diff --git a/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py b/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py new file mode 100644 index 0000000..9e77de0 --- /dev/null +++ b/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py @@ -0,0 +1,34 @@ +"""Unit tests for strict refiner oracle-parity verdicts.""" + +from __future__ import annotations + +import importlib.util +import sys +from pathlib import Path + +REPOSITORY_ROOT = Path(__file__).resolve().parents[3] +TOOL = REPOSITORY_ROOT / "tools" / "validation" / "run_lingbot_video_refiner_core_parity.py" + + +def _load_tool_module(): + spec = importlib.util.spec_from_file_location("run_lingbot_video_refiner_core_parity_test", TOOL) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_exact_parity_failures_accepts_zero_drift() -> None: + module = _load_tool_module() + + assert module.exact_parity_failures({"max_abs": 0.0, "mean_abs": 0.0, "relative_l2": 0.0, "cosine": 1.0}) == [] + + +def test_exact_parity_failures_reports_drift() -> None: + module = _load_tool_module() + + assert module.exact_parity_failures({"max_abs": 0.25, "mean_abs": 0.0, "relative_l2": 0.125, "cosine": 0.99}) == [ + "max_abs", + "relative_l2", + ] diff --git a/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py b/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py new file mode 100644 index 0000000..adc7eb5 --- /dev/null +++ b/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tools.validation.validate_lingbot_video_refiner_handoff import exact_handoff_failures, validate_refiner_handoff + + +def test_refiner_mp4_handoff_tool_matches_source_loader(tmp_path: Path) -> None: + pytest.importorskip("av") + from diffusers.utils import export_to_video + + 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) + + report = validate_refiner_handoff( + path, + upstream_root=Path("work_dirs/lingbot-video-master"), + height=16, + width=16, + sample_fps=24, + vae_tc=4, + ) + + assert report["metadata_match"] + assert report["source_mp4_to_telefuser_mp4"] == { + "shape_match": True, + "dtype_match": True, + "max_abs": 0.0, + "mean_abs": 0.0, + "relative_l2": 0.0, + "cosine": 1.0, + "exact_mismatch_count": 0, + } + + in_memory_video = torch.from_numpy(np.stack(frames)).permute(3, 0, 1, 2).unsqueeze(0).float().div(255.0) + report_with_memory = validate_refiner_handoff( + path, + upstream_root=Path("work_dirs/lingbot-video-master"), + height=16, + width=16, + sample_fps=24, + vae_tc=4, + in_memory_video=in_memory_video, + in_memory_fps=48.0, + ) + + assert report_with_memory["in_memory_metadata_match"] + assert report_with_memory["in_memory_to_mp4"]["shape_match"] + + +def test_exact_handoff_failures_accepts_source_compatibility() -> None: + report = { + "metadata_match": True, + "source_mp4_to_telefuser_mp4": { + "shape_match": True, + "dtype_match": True, + "max_abs": 0.0, + "mean_abs": 0.0, + "relative_l2": 0.0, + "exact_mismatch_count": 0, + }, + } + + assert exact_handoff_failures(report) == [] + + +def test_exact_handoff_failures_ignores_native_handoff_quality_but_reports_source_drift() -> None: + report = { + "metadata_match": False, + "source_mp4_to_telefuser_mp4": { + "shape_match": True, + "dtype_match": False, + "max_abs": 0.25, + "mean_abs": 0.0, + "relative_l2": 0.5, + "exact_mismatch_count": 1, + }, + "in_memory_to_mp4": {"relative_l2": 0.75}, + } + + assert exact_handoff_failures(report) == [ + "metadata_match", + "source_mp4_to_telefuser_mp4.dtype_match", + "source_mp4_to_telefuser_mp4.max_abs", + "source_mp4_to_telefuser_mp4.relative_l2", + "source_mp4_to_telefuser_mp4.exact_mismatch_count", + ] diff --git a/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py b/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py new file mode 100644 index 0000000..71688b5 --- /dev/null +++ b/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py @@ -0,0 +1,85 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from tools.validation.validate_lingbot_video_refiner_output_handoff import ( + _write_video_tensor, + build_refiner_handoff_report, + decoded_frame_quality_metrics, + side_by_side_frames, + write_side_by_side_video, +) + + +def test_final_handoff_report_includes_input_and_final_output_metrics() -> None: + memory_input = torch.tensor([[[[[0.0]]]]]) + mp4_input = torch.tensor([[[[[0.25]]]]]) + memory_output = torch.full((1, 3, 1, 1, 1), 0.5) + mp4_output = torch.full((1, 3, 1, 1, 1), 0.75) + + report = build_refiner_handoff_report( + memory_input=memory_input, + mp4_input=mp4_input, + memory_metadata={"frames": 5}, + mp4_metadata={"frames": 5}, + memory_output=memory_output, + mp4_output=mp4_output, + memory_seconds=1.25, + mp4_seconds=1.5, + ) + + assert report["comparison_baseline"] == "source_mp4_round_trip" + assert report["mp4_round_trip_is_lossy"] + assert report["metadata_match"] + assert report["in_memory_to_mp4_input"]["max_abs"] == 0.25 + assert report["final_output"]["max_abs"] == 0.25 + assert report["final_output_quality"]["shape_match"] + assert report["final_output_quality"]["psnr_db"] is not None + assert report["memory_refiner_seconds"] == 1.25 + assert report["mp4_refiner_seconds"] == 1.5 + + +def test_decoded_frame_quality_metrics_are_exact_for_identical_video() -> None: + frames = torch.linspace(0.0, 1.0, 3 * 2 * 4 * 4).reshape(1, 3, 2, 4, 4) + + metrics = decoded_frame_quality_metrics(frames, frames.clone()) + + assert metrics == {"shape_match": True, "psnr_db": None, "ssim": 1.0} + + +def test_side_by_side_frames_concatenates_video_width() -> None: + reference = torch.zeros(1, 3, 2, 4, 5) + candidate = torch.ones_like(reference) + + comparison = side_by_side_frames(reference, candidate) + + assert comparison.shape == (1, 3, 2, 4, 10) + assert torch.equal(comparison[..., :5], reference) + assert torch.equal(comparison[..., 5:], candidate) + + +def test_video_writers_pass_normalized_float_frames(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: + from diffusers import utils + + captured: list[np.ndarray] = [] + + def capture_export(frames: list[np.ndarray], output_path: str, fps: int) -> str: + del output_path, fps + captured.extend(frames) + return "unused.mp4" + + monkeypatch.setattr(utils, "export_to_video", capture_export) + frames = torch.full((1, 3, 2, 2, 2), 0.5) + + write_side_by_side_video(frames, frames, tmp_path / "comparison.mp4") + _write_video_tensor(frames, tmp_path / "base.mp4") + + assert len(captured) == 4 + assert captured[0].dtype == np.float32 + assert float(captured[0][0, 0, 0]) == 0.5 + assert captured[2].dtype == np.float32 + assert float(captured[2][0, 0, 0]) == 0.5 diff --git a/tools/validation/benchmark_lingbot_video.py b/tools/validation/benchmark_lingbot_video.py new file mode 100644 index 0000000..e211ad1 --- /dev/null +++ b/tools/validation/benchmark_lingbot_video.py @@ -0,0 +1,285 @@ +"""Benchmark a native LingBot-Video base or base-plus-refiner run on one GPU.""" + +from __future__ import annotations + +import argparse +import json +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 + +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + build_lingbot_video_pipeline, + build_lingbot_video_refiner_stage, + 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 = _measure( + setup_samples, + "base_load", + lambda: build_lingbot_video_pipeline( + args.model_dir, + variant=args.variant, + 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) + ) + negative, negative_mask = _measure( + active_samples, "refiner_text_encode", lambda: pipeline.text_stage.encode(negative_caption) + ) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + negative = generation.prompt_conditions.negative_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative_mask = generation.prompt_conditions.negative_attention_mask + refiner_prompt_conditions_reused = True + _measure(active_samples, "base_release", pipeline.release_gpu_resources) + if refiner is None: + refiner = _measure( + setup_samples, + "refiner_load", + lambda: build_lingbot_video_refiner_stage(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..1f7f586 --- /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 UTC, datetime +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(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..808ea66 --- /dev/null +++ b/tools/validation/inspect_lingbot_video_checkpoint.py @@ -0,0 +1,162 @@ +"""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.pipelines.lingbot_video.loading import ( + checkpoint_key_coverage, + load_lingbot_video_dense_transformer, + load_lingbot_video_moe_transformer, +) + +_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 _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() + loader = load_lingbot_video_dense_transformer if args.variant == "dense" else load_lingbot_video_moe_transformer + model = loader(checkpoint_dir, device=args.device, torch_dtype=_dtype(args.dtype)) + 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..ae0381d --- /dev/null +++ b/tools/validation/replay_lingbot_video_dense_reference.py @@ -0,0 +1,383 @@ +"""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.pipelines.lingbot_video.data import preprocess_ti2v_image +from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from telefuser.pipelines.lingbot_video.loading import load_lingbot_video_dense_transformer +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() + ) + stage = LingBotVideoTextEncodingStage("text_encoder", text_encoder, processor, 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() + vae_encode = LingBotVideoVAEEncodeStage("vae_encode", vae, 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) + transformer = load_lingbot_video_dense_transformer(model_dir / "transformer", device=device) + denoising = LingBotVideoDenoisingStage("transformer", transformer, 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) + decoder = LingBotVideoVAEDecodeStage("vae_decode", vae, 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..90406b1 --- /dev/null +++ b/tools/validation/run_lingbot_video_dense_parity.py @@ -0,0 +1,80 @@ +"""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.pipelines.lingbot_video.loading import load_lingbot_video_dense_transformer + + +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)) + telefuser = load_lingbot_video_dense_transformer(args.transformer_dir) + 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_moe_parity.py b/tools/validation/run_lingbot_video_moe_parity.py new file mode 100644 index 0000000..44f08e3 --- /dev/null +++ b/tools/validation/run_lingbot_video_moe_parity.py @@ -0,0 +1,101 @@ +"""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.pipelines.lingbot_video.loading import load_lingbot_video_moe_transformer + + +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() + + telefuser = load_lingbot_video_moe_transformer(directory, device=device, torch_dtype=torch.bfloat16) + 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_refiner_core_parity.py b/tools/validation/run_lingbot_video_refiner_core_parity.py new file mode 100644 index 0000000..ce52026 --- /dev/null +++ b/tools/validation/run_lingbot_video_refiner_core_parity.py @@ -0,0 +1,172 @@ +"""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.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage +from telefuser.pipelines.lingbot_video.loading import load_lingbot_video_moe_transformer +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() + + local = load_lingbot_video_moe_transformer(DIRECTORY, device=DEVICE, torch_dtype=torch.bfloat16) + local_scheduler = FlowUniPCMultistepScheduler.from_pretrained(ROOT / "scheduler") + stage = LingBotVideoRefinerStage( + denoising_stage=LingBotVideoDenoisingStage( + "refiner", local, 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/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..0745e40 --- /dev/null +++ b/tools/validation/validate_lingbot_video_refiner_output_handoff.py @@ -0,0 +1,312 @@ +"""Compare final Refiner outputs from in-memory and MP4 base-video handoffs.""" + +from __future__ import annotations + +import argparse +import json +import tempfile +import time +from pathlib import Path +from typing import Any + +import torch +from PIL import Image + +from telefuser.pipelines.lingbot_video import ( + LingBotVideoRequest, + build_lingbot_video_pipeline, + build_lingbot_video_refiner_stage, + 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_lingbot_video_pipeline( + args.model_dir, + variant="moe", + 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 or pipeline.model_dir is None: + raise RuntimeError("LingBot refiner requires the base text stage and checkpoint metadata") + if generation.prompt_conditions.has_visual_condition: + positive, positive_mask = pipeline.text_stage.encode(caption) + negative, negative_mask = pipeline.text_stage.encode(negative_caption) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + negative = generation.prompt_conditions.negative_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative_mask = generation.prompt_conditions.negative_attention_mask + _, base_release_seconds = _measure(pipeline.release_gpu_resources) + refiner, refiner_load_seconds = _measure( + lambda: build_lingbot_video_refiner_stage(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() From 6baa9dc50375ceea0616fdfc583116bad2199c79 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 14:05:18 +0000 Subject: [PATCH 2/9] feat(lingbot-video): add distributed inference paths - add four-GPU FSDP and Ulysses execution for Dense and refiner stages - keep complete scheduler loops inside distributed stage workers - split the generic entrypoints into Dense and MoE model-specific examples - add native-parallel, distributed, and sequence-parallel validation tools Verification: - git diff --cached --check --- docs/en/lingbot_video.md | 97 +++- docs/zh/lingbot_video.md | 82 ++- examples/example_config.yaml | 18 + examples/lingbot_video/README.md | 57 +- .../lingbot_video/assets/t2v_5s.json.example | 77 +++ .../lingbot_video/lingbot_video_dense_1_3b.py | 351 ++++++++++++ .../lingbot_video/lingbot_video_generate.py | 142 ----- .../lingbot_video/lingbot_video_moe_30b.py | 536 ++++++++++++++++++ .../lingbot_video/lingbot_video_service.py | 205 ------- .../wan21_14b_image_to_video_h100.py | 25 +- telefuser/distributed/fsdp.py | 4 + telefuser/models/lingbot_video_dit.py | 90 ++- telefuser/pipelines/lingbot_video/__init__.py | 5 +- .../pipelines/lingbot_video/denoising.py | 81 ++- telefuser/pipelines/lingbot_video/pipeline.py | 104 ++-- telefuser/pipelines/lingbot_video/refiner.py | 56 +- telefuser/pipelines/lingbot_video/runtime.py | 129 ----- .../pipelines/lingbot_video/text_encoding.py | 9 +- telefuser/pipelines/lingbot_video/vae.py | 9 +- .../pipelines/lingbot_video/test_pipeline.py | 116 +++- .../lingbot_video/test_refiner_stage.py | 84 +++ .../pipelines/lingbot_video/test_runtime.py | 26 +- .../lingbot_video/test_service_contract.py | 223 ++++---- .../lingbot_video/test_text_encoding.py | 9 +- .../unit/pipelines/lingbot_video/test_vae.py | 7 +- tools/validation/benchmark_lingbot_video.py | 16 +- .../replay_lingbot_video_dense_reference.py | 17 +- .../run_lingbot_video_distributed.py | 130 +++++ .../run_lingbot_video_native_parallel.py | 167 ++++++ .../run_lingbot_video_refiner_core_parity.py | 5 +- .../validation/run_lingbot_video_sp_parity.py | 129 +++++ ...te_lingbot_video_refiner_output_handoff.py | 17 +- 32 files changed, 2242 insertions(+), 781 deletions(-) create mode 100644 examples/lingbot_video/assets/t2v_5s.json.example create mode 100644 examples/lingbot_video/lingbot_video_dense_1_3b.py delete mode 100644 examples/lingbot_video/lingbot_video_generate.py create mode 100644 examples/lingbot_video/lingbot_video_moe_30b.py delete mode 100644 examples/lingbot_video/lingbot_video_service.py delete mode 100644 telefuser/pipelines/lingbot_video/runtime.py create mode 100644 tools/validation/run_lingbot_video_distributed.py create mode 100644 tools/validation/run_lingbot_video_native_parallel.py create mode 100644 tools/validation/run_lingbot_video_sp_parity.py diff --git a/docs/en/lingbot_video.md b/docs/en/lingbot_video.md index ef402ca..b45e430 100644 --- a/docs/en/lingbot_video.md +++ b/docs/en/lingbot_video.md @@ -64,16 +64,18 @@ can materially alter color and image quality. checkpoint VAE's latent mean/std normalization. - `FlowUniPCMultistepScheduler` owns the sigma/timestep sequence. -Attach stages with `pipeline.set_runtime(...)` after initializing the pipeline -config. Provide a structured JSON caption, not casual unstructured text. +Load checkpoint components into `ModuleManager`, then call `pipeline.init(module_manager, config)`; `init` constructs all text, DiT, VAE, and scheduler stages from the manager +configuration. Provide a structured JSON caption, not casual unstructured text. -For the standard checkpoint layout, `build_lingbot_video_pipeline` loads those -components directly without importing the upstream runtime: +For the standard checkpoint layout, model loading and stage assembly live directly +in the public model-specific examples. This keeps the runtime wiring visible next +to `PPL_CONFIG`, `CONTRACT`, and the CLI/service entrypoints: ```python -from telefuser.pipelines.lingbot_video import LingBotVideoRequest, build_lingbot_video_pipeline +from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline +from telefuser.pipelines.lingbot_video import LingBotVideoRequest -pipeline = build_lingbot_video_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) +pipeline = build_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) ``` @@ -90,6 +92,38 @@ float frames directly to Diffusers `export_to_video`, which performs the uint8 conversion itself. Converting to uint8 before that call applies a second 255 scale and wraps channel values, producing a negative-like MP4. +## Four-GPU base inference + +Dense and MoE base DiTs support TeleFuser-native four-GPU inference through +per-block FSDP and Ulysses sequence parallelism. The whole joint +video-plus-text sequence is padded only when necessary, sharded across the four +ranks, then restored in source token order. CFG is evaluated as one batched +forward when positive and negative Qwen embeddings have matching shapes; +otherwise the runtime safely uses the source-order two-forward path. + +Start the API service normally; TeleFuser owns the four worker processes, so do +not wrap this command in `torchrun`: + +```bash +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 +``` + +This configuration keeps FSDP-enabled DiT weights resident on GPU and therefore +requires `cpu_offload=False`. Checkpoint FP32 modulation parameters are retained +as replicated FSDP ignored states, matching the upstream mixed-precision layout. +The MoE 30B model uses its own example; its sorted eager experts are still +replicated within each FSDP block rather than expert-parallelized: + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 +``` + +The validated 832x480, 121-frame, 40-step MoE base run used 68.1 GiB peak GPU +memory and 240.4 seconds for generation excluding checkpoint load and MP4 +encoding. The separately loaded refiner also has a validated four-GPU +FSDP/Ulysses path. Expert parallelism, grouped GEMM, and FP8 remain disabled +pending their own parity and throughput evidence. + ## TI2V TI2V has two independent first-frame condition paths: @@ -103,16 +137,27 @@ The current pipeline accepts a raw RGB condition image in the range [0,255] with ## Service -The service entrypoint exposes `t2i`, `t2v`, and `i2v` through TeleFuser service APIs and requires a structured JSON string as `prompt`: +The Dense and MoE examples each expose `PPL_CONFIG`, `CONTRACT`, `get_pipeline`, +`run`, and `run_with_file`. Both serve `t2i`, `t2v`, and `i2v` and require a +structured JSON string as `prompt`: + +Set each example's `PPL_CONFIG["model_root"]` before starting the service. Runtime options are read from `PPL_CONFIG`; direct CLI runs can override the exposed command-line options. ```bash -LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ - telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --port 8000 ``` -Set `LINGBOT_VIDEO_VARIANT=moe` and `LINGBOT_VIDEO_ENABLE_REFINER=1` to use the MoE checkpoint with its separate refiner. The service releases base-stage weights before loading the refiner; `refine` is also an explicit per-request boolean parameter. +Set `PPL_CONFIG["model_root"]` in `lingbot_video_moe_30b.py` before serving MoE. +Its contract exposes `refine`, enabled by default. The service releases base +weights before loading the refiner. `PPL_CONFIG["refiner_parallelism"] = 4` +selects an independent four-worker refiner; when omitted, it inherits service +parallelism. + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 +``` Requested service resolutions are rounded up to the LingBot VAE-and-DiT -sixteen-pixel spatial grid; for example, `480p` at `16:9` resolves to 864x480. +sixteen-pixel spatial grid; for example, `480p` at `16:9` uses the validated LingBot 832x480 landscape preset. Pass `negative_prompt` only to override the source-compatible default; an explicit empty string remains a supported override. @@ -123,15 +168,28 @@ mixes it with noise at `t_thresh`, and samples the low-noise sigma tail. It can also preserve a clean TI2V frame zero. Base and refiner are separate runtime stages. Call `base_pipeline.release_gpu_resources()` before loading the refiner when they share a GPU. + +On four GPUs, the refiner uses per-block FSDP and four-way Ulysses SP. Its CFG +defaults to sequential positive/negative forwards because batched CFG exceeds +four H100 80 GB cards at 1920x1088. Set `PPL_CONFIG["refiner_batch_cfg"] = True` +only after validating additional memory capacity. The runtime VAE-encodes the +RGB handoff, releases the VAE, runs the distributed refiner, closes all refiner +workers, and only then reloads the VAE decoder. + +The validated MoE run generated a 832x480, 121-frame base and refined it to +1920x1088 at 24 FPS. Base generation took 372.9 seconds and the eight-step +refiner stage took 886.8 seconds, excluding checkpoint deserialization. The +output contains 121 frames and is 5.0417 seconds long. The included CLI implements this lifecycle for MoE checkpoints: ```bash -python examples/lingbot_video/lingbot_video_generate.py \ - --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ - --caption-json /path/to/caption.json --output result.mp4 +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 ``` -With `--image first_frame.png`, the CLI also applies the upstream TI2V frame-zero geometry to the refiner condition. +With `--task i2v --first_image_path first_frame.png`, the CLI also applies the +upstream TI2V frame-zero geometry to the refiner condition. For an in-memory base-to-refiner handoff, call `prepare_refiner_video(...)` before `LingBotVideoRefinerStage.refine(...)`. It matches upstream training-aligned frame @@ -161,6 +219,8 @@ python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.m python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +python -m torch.distributed.run --standalone --nproc_per_node=4 tools/validation/run_lingbot_video_distributed.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output dense-sp4.mp4 --report dense-sp4.json +python tools/validation/run_lingbot_video_native_parallel.py --variant moe --refine --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --output moe-refiner-sp4.mp4 --report moe-refiner-sp4.json ``` `validate_lingbot_video_refiner_handoff.py` checks that the TeleFuser MP4 @@ -212,5 +272,8 @@ The numerical-oracle path requires CUDA, PyTorch, Diffusers, Transformers, and t checkpoint components `transformer/`, `text_encoder/`, `processor/`, `vae/`, and `scheduler/`. Dense runs source-equivalently on one GPU. The current MoE expert implementation is an eager correctness path; do not use it as a 30B production -backend. FSDP, Ulysses/CFG parallelism, FlashAttention, and FP8 experts are intentionally not -enabled. The service and CLI both support a serial single-process base-plus-refiner lifecycle; distributed alignment is outside the current support scope. +backend. Dense and MoE base DiTs have validated four-GPU FSDP/Ulysses SP paths. +The MoE refiner has a validated independent four-GPU FSDP/Ulysses SP stage. +FlashAttention, MoE expert parallelism, grouped GEMM, and FP8 experts are +intentionally not enabled. The base-plus-refiner stage lifecycle remains +serial so their 30B weights never coexist on GPU. diff --git a/docs/zh/lingbot_video.md b/docs/zh/lingbot_video.md index 6b1a440..350f2bf 100644 --- a/docs/zh/lingbot_video.md +++ b/docs/zh/lingbot_video.md @@ -44,7 +44,7 @@ caption。T2I 使用上游的静态图像版本；T2V 与 TI2V 使用包含时 ## 运行时组成 -`LingBotVideoPipeline` 通过 `set_runtime(...)` 组合独立加载的 stage： +`LingBotVideoPipeline` 通过标准的 `init(module_manager, config)` 接口统一初始化；checkpoint 组件先加载到 `ModuleManager`，再由 `init` 创建全部 stage： - `LingBotVideoTextEncodingStage`：Qwen3-VL structured JSON caption 编码。 - `LingBotVideoDenoisingStage`：source-order 的 two-forward CFG。 @@ -53,12 +53,14 @@ caption。T2I 使用上游的静态图像版本；T2V 与 TI2V 使用包含时 核心 pipeline 应输入结构化 JSON caption，不应将普通自然语言 prompt 直接替代 rewriter 输出。 -标准 checkpoint 目录可直接使用 `build_lingbot_video_pipeline` 装配，不依赖上游 runtime： +标准 checkpoint 的模型加载与 stage 装配直接放在公开的模型专用 example 中， +与 `PPL_CONFIG`、`CONTRACT` 以及 CLI/service 入口保持在同一层，方便阅读和修改： ```python -from telefuser.pipelines.lingbot_video import LingBotVideoRequest, build_lingbot_video_pipeline +from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline +from telefuser.pipelines.lingbot_video import LingBotVideoRequest -pipeline = build_lingbot_video_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) +pipeline = build_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) ``` @@ -72,17 +74,52 @@ VAE decode stage 的 RGB video 范围为 `[0,1]`。视频调用方必须将这 Diffusers 的 `export_to_video`，由它自行转换为 uint8。若在调用前先转 uint8，就会再次 乘以 255 并使通道值溢出，生成近似负片的 MP4。 +## 四卡 base 推理 + +Dense 与 MoE base DiT 均支持 TeleFuser 原生四卡：DiT block 使用 FSDP，联合 video/text +token 流使用 Ulysses sequence parallel。只有 token 总数不能被 rank 数整除时才补齐；输出会按 +上游 token 顺序恢复。正、负 Qwen embedding 形状一致时，CFG 使用一次 batched forward；形状 +不一致时会安全回退到保持上游顺序的两次 forward。 + +按普通方式启动 API 服务即可，TeleFuser 会管理四个 worker，不能再用 `torchrun` 包裹： + +```bash +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 +``` + +该模式要求 DiT 常驻 GPU，因此必须 `cpu_offload=False`。checkpoint 中保留 FP32 的 modulation +参数会作为 FSDP ignored state 在每个 rank 副本化，与上游 mixed-precision 布局一致。MoE 30B +使用独立 example；sorted eager expert 仍随 FSDP block 分片，而不是 expert parallel： + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 +``` + +已验证的 MoE base 为 832x480、121 帧、40 steps，生成阶段（不含 checkpoint 加载与 MP4 编码） +耗时 240.4 秒，峰值显存 68.1 GiB/卡。独立加载的 refiner 也已有经验证的四卡 +FSDP/Ulysses 路径；expert parallel、grouped GEMM 与 FP8 仍需独立 parity 与吞吐验证。 + ## 服务 -服务入口通过 TeleFuser API 暴露 `t2i`、`t2v`、`i2v`，`prompt` 必须是 structured JSON 字符串： +Dense 与 MoE example 都提供 `PPL_CONFIG`、`CONTRACT`、`get_pipeline`、`run` 与 +`run_with_file`，并通过 TeleFuser API 暴露 `t2i`、`t2v`、`i2v`；`prompt` 必须是 +structured JSON 字符串： + +启动服务前需要在对应 example 中设置 `PPL_CONFIG["model_root"]`。运行参数从 `PPL_CONFIG` 读取；直接执行 CLI 时可通过对应命令行参数覆盖。 ```bash -LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ - telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --port 8000 ``` -设置 `LINGBOT_VIDEO_VARIANT=moe` 与 `LINGBOT_VIDEO_ENABLE_REFINER=1` 可启用带独立 refiner 的 MoE checkpoint。服务会在加载 refiner 前释放 base stage 的权重；`refine` 也可作为显式的单请求布尔参数。 -服务请求的分辨率会向上对齐到 LingBot VAE 与 DiT 所需的 16 像素网格；例如 `480p`、`16:9` 会解析为 864x480。 +MoE 使用 `lingbot_video_moe_30b.py`，模型路径由其中的 `PPL_CONFIG["model_root"]` 配置。它的 contract +显式提供 `refine`，默认启用。服务会在加载 refiner 前释放 base stage 权重； +`PPL_CONFIG["refiner_parallelism"] = 4` 可为 refiner 单独选择四 worker，未设置时继承 +service parallelism。 + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 +``` +服务请求的分辨率会向上对齐到 LingBot VAE 与 DiT 所需的 16 像素网格；例如 `480p`、`16:9` 使用已验证的 LingBot 832x480 横屏预设。 仅在确实需要覆盖上游默认值时才传入 `negative_prompt`；显式传入空字符串仍是受支持的覆盖方式。 ## TI2V 与 Refiner @@ -91,15 +128,26 @@ TI2V 接收范围为 [0,255] 的 RGB tensor（`[B, 3, H, W]` 或 `[B, 3, F, H, W `LingBotVideoRefinerStage` 直接接收 base RGB tensor，避免 MP4/decord 往返；它按 `t_thresh` 与噪声混合，并使用 low-noise sigma tail 采样。base 与 refiner 是独立运行时 stage；共享 GPU 时，请在加载 refiner 前调用 `base_pipeline.release_gpu_resources()`。 +四卡模式下，refiner 使用 block FSDP 与四路 Ulysses SP。1920x1088 在四张 H100 80 GB +上默认采用正、负条件依次执行的 sequential CFG；batched CFG 会超出四卡显存，仅在额外 +显存容量已经验证时设置 `PPL_CONFIG["refiner_batch_cfg"] = True`。运行时顺序为：VAE encode、 +释放 VAE、分布式 refiner 去噪、关闭 refiner workers、重新加载 VAE decode，避免高显存 +stage 在 rank 0 重叠。 + +已验证的 MoE 运行先生成 832x480、121 帧 base，再以 8 steps 输出 1920x1088、24 FPS +视频。base 耗时 372.9 秒，refiner stage 耗时 886.8 秒（不含 checkpoint 反序列化）； +最终视频为 121 帧、5.0417 秒。 + 示例 CLI 已为 MoE checkpoint 自动完成这个生命周期： ```bash -python examples/lingbot_video/lingbot_video_generate.py \ - --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ - --caption-json /path/to/caption.json --output result.mp4 +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 ``` -附加 `--image first_frame.png` 时，CLI 也会把上游 TI2V frame-zero 几何规则应用到 refiner condition。 +附加 `--task i2v --first_image_path first_frame.png` 时，CLI 也会把上游 TI2V +frame-zero 几何规则应用到 refiner condition。 如需在内存中交接基座输出与 refiner，可先调用 `prepare_refiner_video(...)`，再调用 `LingBotVideoRefinerStage.refine(...)`。该函数复刻上游训练对齐抽帧与双三次缩放， @@ -126,6 +174,8 @@ python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.m python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 +python -m torch.distributed.run --standalone --nproc_per_node=4 tools/validation/run_lingbot_video_distributed.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output dense-sp4.mp4 --report dense-sp4.json +python tools/validation/run_lingbot_video_native_parallel.py --variant moe --refine --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --output moe-refiner-sp4.mp4 --report moe-refiner-sp4.json ``` `validate_lingbot_video_refiner_handoff.py` 验证 TeleFuser 的 MP4 兼容 loader 与上游 @@ -158,5 +208,7 @@ Refiner core CLI 会向上游与 TeleFuser 的 low-noise 路径注入相同的 l 数值 oracle 路径需要 CUDA、PyTorch、Diffusers、Transformers 以及 checkpoint 的 `transformer/`、`text_encoder/`、`processor/`、`vae/`、`scheduler/` 组件。Dense 可在单 GPU 上按 source-equivalent 路径运行。当前 MoE expert 是 eager correctness -实现，不应作为 30B 的生产后端。FSDP、Ulysses/CFG 并行、FlashAttention 与 FP8 expert -尚未启用。服务与 CLI 都支持串行单进程的 base+refiner 生命周期；当前不支持多卡对齐。 +实现，不应作为 30B 的生产后端。Dense 与 MoE base 都已有经验证的四卡 FSDP/Ulysses SP 路径； +MoE refiner 也已有独立的四卡 FSDP/Ulysses SP stage。FlashAttention、MoE expert parallel、 +grouped GEMM 与 FP8 expert 尚未启用。base+refiner 的 stage 生命周期保持串行，确保两个 +30B 权重不会同时驻留 GPU。 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 index 0ee74ed..71734ba 100644 --- a/examples/lingbot_video/README.md +++ b/examples/lingbot_video/README.md @@ -1,30 +1,61 @@ # LingBot-Video -Use a structured JSON caption produced by the LingBot rewriter. Dense T2V: +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. + +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_generate.py \ - --model-dir /path/to/lingbot-video-dense-1.3b \ - --caption-json /path/to/caption.json --output result.mp4 +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 `--image first_frame.png` for TI2V, or `--variant moe` for the MoE base. +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_generate.py \ - --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine \ - --caption-json /path/to/caption.json --output result.mp4 +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 +``` + +The command releases base-stage GPU weights before loading the separate refiner. +Use `--no-refine` for MoE base-only generation. The default sorted eager MoE +Use `--refiner_gpu_num` and `--refiner_batch_cfg` to override the corresponding MoE `PPL_CONFIG` defaults for a direct CLI run. +route path is source-equivalent with a diagnostic fallback, but it is not a +grouped-GEMM or FP8 production-performance backend. + +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 ``` -The command releases base-stage GPU weights before loading the separate refiner. Add `--image first_frame.png` to preserve the TI2V clean frame-zero condition. The default sorted eager MoE route path is source-equivalent with a diagnostic fallback, but it is not a grouped-GEMM or FP8 production-performance backend. +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 SP4/FSDP stage explicitly; +otherwise it inherits service parallelism. + +```bash +telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 +``` -Serve structured-caption T2I/T2V/TI2V requests with: +Dense and MoE base checkpoints support TeleFuser-native four-GPU FSDP plus Ulysses sequence parallelism: ```bash -LINGBOT_VIDEO_MODEL_ROOT=/path/to/lingbot-video-dense-1.3b \ - telefuser serve examples/lingbot_video/lingbot_video_service.py --port 8000 +telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 ``` -Set `LINGBOT_VIDEO_VARIANT=moe LINGBOT_VIDEO_ENABLE_REFINER=1` for the MoE base-plus-refiner service lifecycle. +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. The refiner starts as a separate worker group +after the base workers exit. Four-GPU 1920x1088 refinement uses sequential CFG +by default to stay within H100 80 GB memory. 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..135bd18 --- /dev/null +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -0,0 +1,351 @@ +"""LingBot-Video Dense 1.3B CLI and service example.""" + +from __future__ import annotations + +import json +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_dense_transformer, + 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") + +PPL_CONFIG: dict[str, Any] = { + "name": "lingbot_video_dense_1_3b", + "model_root": "model_zoo/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, +} + + +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, + ) + + +CONTRACT = _build_contract() +PIPELINE_CONTRACT = CONTRACT + + +def _resolve_runtime_device(device: str, parallel_config: ParallelConfig) -> tuple[str, int]: + """Resolve the torchrun-local device or the parent device for native workers.""" + if parallel_config.world_size == 1: + parsed = torch.device(device) + return device, parsed.index or 0 + if torch.device(device).type != "cuda": + raise ValueError("LingBot distributed inference requires CUDA devices") + local_rank = torch.cuda.current_device() + torch.cuda.set_device(local_rank) + return f"cuda:{local_rank}", local_rank + + +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 + if parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot FSDP inference requires cpu_offload=False") + device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, + device_id=device_id, + 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.add_module( + load_lingbot_video_dense_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype), + name="transformer", + path=str(transformer_dir), + ) + 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"]) -> 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") + return build_pipeline( + model_root, + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)), + sp_ulysses_degree=parallelism, + 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: int = PPL_CONFIG["target_video_length"], + 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, _ = parse_lingbot_video_prompt(json.loads(prompt)) + width, height = _resolve_video_size(aspect_ratio, resolution) + num_frames = ( + 1 if task == "t2i" else num_frames_from_duration(float(target_video_length), 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) + path.parent.mkdir(parents=True, exist_ok=True) + 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(np.uint8)).save(path) + else: + 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: int = PPL_CONFIG["target_video_length"], + 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("--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=PPL_CONFIG["target_video_length"], type=int) +@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, + model_root: str, + prompt: str, + negative_prompt: str | None, + seed: int, + resolution: str, + aspect_ratio: str, + target_video_length: int, + 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) + 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_generate.py b/examples/lingbot_video/lingbot_video_generate.py deleted file mode 100644 index 030b0b1..0000000 --- a/examples/lingbot_video/lingbot_video_generate.py +++ /dev/null @@ -1,142 +0,0 @@ -"""Generate a LingBot-Video T2I, T2V, or TI2V sample with native TeleFuser stages.""" - -from __future__ import annotations - -import argparse - -import torch -from PIL import Image - -from telefuser.pipelines.lingbot_video import ( - LingBotVideoRequest, - build_lingbot_video_pipeline, - build_lingbot_video_refiner_stage, - default_negative_caption, - load_lingbot_video_prompt, - load_refiner_first_frame, - num_frames_from_duration, - prepare_refiner_video, -) - - -def _image_to_tensor(path: str) -> torch.Tensor: - """Load an RGB image as the raw [0,255] tensor required by TI2V.""" - image = Image.open(path).convert("RGB") - return torch.from_numpy(__import__("numpy").asarray(image).copy()).permute(2, 0, 1).unsqueeze(0).float() - - -def main() -> None: - parser = argparse.ArgumentParser(description=__doc__) - parser.add_argument("--model-dir", required=True) - parser.add_argument("--caption-json", required=True, help="Path to the structured JSON caption.") - parser.add_argument("--output", required=True, help="Output MP4 path, or PNG when --num-frames=1.") - parser.add_argument("--height", type=int, default=480) - parser.add_argument("--width", type=int, default=832) - parser.add_argument( - "--num-frames", type=int, help="Override prompt-file duration with an explicit 4n+1 frame count." - ) - 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("--image", help="Optional first-frame image for TI2V.") - parser.add_argument("--variant", choices=("dense", "moe"), default="dense") - parser.add_argument("--negative-caption", default=None, help="Override the official structured negative caption.") - 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) - args = parser.parse_args() - if args.refine and args.variant != "moe": - raise ValueError("--refine requires --variant moe because the refiner is shipped with the MoE checkpoint") - caption_text, 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 = build_lingbot_video_pipeline( - args.model_dir, - variant=args.variant, - guidance_scale=args.guidance_scale, - num_inference_steps=args.steps, - ) - generator = torch.Generator("cuda").manual_seed(args.seed) - generation = pipeline.generate( - LingBotVideoRequest( - caption=caption_text, - height=args.height, - width=args.width, - num_frames=num_frames, - image=source_image, - ), - negative_caption=negative_caption, - generator=generator, - ) - frames = generation.output - if args.refine: - if pipeline.text_stage is None: - raise RuntimeError("LingBot-Video text stage is required for refiner prompt encoding") - # The upstream TI2V refiner consumes text-only prompt conditioning. Its - # frame-zero visual condition is injected through the VAE latent instead. - if generation.prompt_conditions.has_visual_condition: - positive, positive_mask = pipeline.text_stage.encode(caption_text) - negative, negative_mask = pipeline.text_stage.encode(negative_caption) - else: - positive = generation.prompt_conditions.positive_prompt_embeds - negative = generation.prompt_conditions.negative_prompt_embeds - positive_mask = generation.prompt_conditions.positive_attention_mask - negative_mask = generation.prompt_conditions.negative_attention_mask - pipeline.release_gpu_resources() - refiner = build_lingbot_video_refiner_stage(args.model_dir) - lowres_video, _ = 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 - ) - frames = 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), - ) - 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(args.output) - return - from diffusers.utils import export_to_video - - export_to_video(list(video), args.output, fps=24) - - -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..d576844 --- /dev/null +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -0,0 +1,536 @@ +"""LingBot-Video MoE 30B plus refiner CLI and service example.""" + +from __future__ import annotations + +import json +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_lingbot_video_moe_transformer, + 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") + +PPL_CONFIG: dict[str, Any] = { + "name": "lingbot_video_moe_30b", + "model_root": "model_zoo/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, + "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, +} + + +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"]}, + "refine": {"type": "boolean", "default": PPL_CONFIG["enable_refiner"]}, + } + 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, + ) + + +CONTRACT = _build_contract() +PIPELINE_CONTRACT = CONTRACT + + +def _resolve_runtime_device(device: str, parallel_config: ParallelConfig) -> tuple[str, int]: + """Resolve the torchrun-local device or the parent device for native workers.""" + if parallel_config.world_size == 1: + parsed = torch.device(device) + return device, parsed.index or 0 + if torch.device(device).type != "cuda": + raise ValueError("LingBot distributed inference requires CUDA devices") + local_rank = torch.cuda.current_device() + torch.cuda.set_device(local_rank) + return f"cuda:{local_rank}", local_rank + + +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, +) -> 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 + if parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot FSDP inference requires cpu_offload=False") + device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, + device_id=device_id, + 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.add_module( + load_lingbot_video_moe_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype), + name="transformer", + path=str(transformer_dir), + ) + 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, +) -> 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 parallel_config.enable_fsdp and cpu_offload: + raise ValueError("LingBot refiner FSDP inference requires cpu_offload=False") + device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, + device_id=device_id, + 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.add_module( + load_lingbot_video_moe_transformer(root / "refiner", device="cpu", torch_dtype=torch_dtype), + name="transformer", + path=str(root / "refiner"), + ) + 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, +) -> 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") + pipeline = build_pipeline( + model_root, + cpu_offload=parallelism == 1, + parallel_config=ParallelConfig( + device_ids=list(range(parallelism)), + sp_ulysses_degree=parallelism, + enable_fsdp=parallelism > 1, + ), + ) + 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") + pipeline.refiner_parallel_config = ParallelConfig( + device_ids=list(range(refiner_parallelism)), + sp_ulysses_degree=refiner_parallelism, + 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 + ) + 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: int = PPL_CONFIG["target_video_length"], + task: str = "t2v", + first_image_path: str = "", + refine: bool = PPL_CONFIG["enable_refiner"], + 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") + + caption, _ = parse_lingbot_video_prompt(json.loads(prompt)) + width, height = _resolve_video_size(aspect_ratio, resolution) + num_frames = ( + 1 if task == "t2i" else num_frames_from_duration(float(target_video_length), 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) + negative, negative_mask = pipeline.text_stage.encode(resolved_negative_prompt) + else: + positive = generation.prompt_conditions.positive_prompt_embeds + negative = generation.prompt_conditions.negative_prompt_embeds + positive_mask = generation.prompt_conditions.positive_attention_mask + negative_mask = generation.prompt_conditions.negative_attention_mask + + 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), + ) + 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), + ) + 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) + path.parent.mkdir(parents=True, exist_ok=True) + 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(np.uint8)).save(path) + else: + 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: int = PPL_CONFIG["target_video_length"], + resolution: str = PPL_CONFIG["resolution"], + aspect_ratio: str = PPL_CONFIG["aspect_ratio"], + task: str = "t2v", + first_image_path: str = "", + refine: bool = PPL_CONFIG["enable_refiner"], + 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("--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("--resolution", default=PPL_CONFIG["resolution"]) +@click.option("--aspect_ratio", default=PPL_CONFIG["aspect_ratio"]) +@click.option("--target_video_length", default=PPL_CONFIG["target_video_length"], type=int) +@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("--refine/--no-refine", default=PPL_CONFIG["enable_refiner"]) +@click.option("--output_path", default="output.mp4") +def main( + gpu_num: int, + model_root: str, + prompt: str, + negative_prompt: str | None, + seed: int, + resolution: str, + aspect_ratio: str, + target_video_length: int, + task: str, + first_image_path: str, + refiner_gpu_num: int, + refiner_batch_cfg: bool, + refine: bool, + 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, + ) + 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/lingbot_video/lingbot_video_service.py b/examples/lingbot_video/lingbot_video_service.py deleted file mode 100644 index 28a53f4..0000000 --- a/examples/lingbot_video/lingbot_video_service.py +++ /dev/null @@ -1,205 +0,0 @@ -"""Service entrypoint for LingBot-Video Dense generation.""" - -from __future__ import annotations - -import json -import os - -import numpy as np -import torch -from PIL import Image - -from telefuser.pipelines.lingbot_video import ( - DEFAULT_NEGATIVE_PROMPT, - DEFAULT_NEGATIVE_PROMPT_IMAGE, - LingBotVideoRequest, - build_lingbot_video_pipeline, - build_lingbot_video_refiner_stage, - default_negative_caption, - load_refiner_first_frame, - parse_lingbot_video_prompt, - prepare_refiner_video, -) -from telefuser.utils.video import get_target_video_size_from_ratio - -PPL_CONFIG = { - "model_root": os.environ.get( - "LINGBOT_VIDEO_MODEL_ROOT", "/hhb-data/aigc/model_zoo/lingbot/lingbot-video-dense-1.3b" - ), - "num_inference_steps": 40, - "guidance_scale": 3.0, - "fps": 24, - "variant": os.environ.get("LINGBOT_VIDEO_VARIANT", "dense"), - "enable_refiner": os.environ.get("LINGBOT_VIDEO_ENABLE_REFINER", "0").lower() in {"1", "true", "yes"}, - "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, -} - - -def get_pipeline_contract() -> dict: - """Declare the task contract consumed by TeleFuser service APIs.""" - parameters = { - "seed": {"type": "integer", "default": 42}, - "target_video_length": {"type": "integer", "default": 4}, - "resolution": {"type": "string", "default": "480p"}, - "aspect_ratio": {"type": "string", "default": "16:9"}, - "refine": {"type": "boolean", "default": PPL_CONFIG["enable_refiner"]}, - } - image_parameters = { - **parameters, - "negative_prompt": {"type": "string", "default": DEFAULT_NEGATIVE_PROMPT_IMAGE}, - } - video_parameters = { - **parameters, - "negative_prompt": {"type": "string", "default": DEFAULT_NEGATIVE_PROMPT}, - } - return { - "pipeline_name": "lingbot_video", - "supported_tasks": ["t2i", "t2v", "i2v"], - "supported_media_types": ["image", "video"], - "execution_mode": "serial_single_pipeline", - "effective_max_concurrent_tasks": 1, - "task_contracts": { - "t2i": {"media_type": "image", "parameters": image_parameters}, - "t2v": {"media_type": "video", "parameters": video_parameters}, - "i2v": { - "media_type": "video", - "required_inputs": ["first_image_path"], - "parameters": video_parameters, - }, - }, - } - - -def get_pipeline(parallelism: int = 1, model_root: str = PPL_CONFIG["model_root"]): - """Load a single-GPU source-equivalent LingBot base pipeline.""" - if parallelism != 1: - raise ValueError("LingBot-Video service currently supports parallelism=1") - return build_lingbot_video_pipeline( - model_root, - variant=PPL_CONFIG["variant"], - guidance_scale=PPL_CONFIG["guidance_scale"], - num_inference_steps=PPL_CONFIG["num_inference_steps"], - ) - - -def _load_condition_image(path: str) -> torch.Tensor: - """Load a user-uploaded RGB image as a raw [0,255] TI2V tensor.""" - array = np.asarray(Image.open(path).convert("RGB")).copy() - return torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0).float() - - -def _num_frames(seconds: int, fps: int) -> int: - """Convert service duration to the nearest valid ``4n+1`` LingBot count.""" - return max(1, 4 * round(seconds * fps / 4) + 1) - - -def _lingbot_video_size(aspect_ratio: str, resolution: str) -> tuple[int, int]: - """Resolve a service resolution for the combined VAE and DiT spatial contract.""" - 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 run_with_file( - pipeline, - prompt: str, - *, - first_image_path: str = "", - negative_prompt: str | None = None, - seed: int = 42, - output_path: str = "output.mp4", - target_video_length: int = 4, - resolution: str = "480p", - aspect_ratio: str = "16:9", - task: str = "t2v", - refine: bool | None = None, - **_: object, -) -> dict[str, str]: - """Generate and encode a T2I, T2V, or TI2V service result.""" - width, height = _lingbot_video_size(aspect_ratio, resolution) - caption, _ = parse_lingbot_video_prompt(json.loads(prompt)) - if task not in {"t2i", "t2v", "i2v"}: - raise ValueError(f"unsupported LingBot-Video service task: {task}") - if task == "i2v" and not first_image_path: - raise ValueError("LingBot-Video i2v requires first_image_path") - refine_enabled = PPL_CONFIG["enable_refiner"] if refine is None else refine - if refine_enabled and pipeline.variant != "moe": - raise ValueError("LingBot-Video refiner requires a pipeline loaded with variant=moe") - num_frames = 1 if task == "t2i" else _num_frames(target_video_length, 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=torch.Generator("cuda").manual_seed(seed), - ) - frames = generation.output - if refine_enabled: - if pipeline.text_stage is None or pipeline.model_dir is None: - raise RuntimeError("LingBot-Video refiner requires text-stage and checkpoint metadata") - if generation.prompt_conditions.has_visual_condition: - positive, positive_mask = pipeline.text_stage.encode(caption) - negative, negative_mask = pipeline.text_stage.encode(resolved_negative_prompt) - else: - positive = generation.prompt_conditions.positive_prompt_embeds - negative = generation.prompt_conditions.negative_prompt_embeds - positive_mask = generation.prompt_conditions.positive_attention_mask - negative_mask = generation.prompt_conditions.negative_attention_mask - pipeline.release_gpu_resources() - refiner = build_lingbot_video_refiner_stage(pipeline.model_dir) - lowres_video, _ = prepare_refiner_video( - frames, - source_fps=PPL_CONFIG["fps"], - height=PPL_CONFIG["refiner_height"], - width=PPL_CONFIG["refiner_width"], - ) - clean_first_frame = ( - load_refiner_first_frame( - first_image_path, - target_height=PPL_CONFIG["refiner_height"], - target_width=PPL_CONFIG["refiner_width"], - geometry_height=height, - geometry_width=width, - ) - if first_image_path - else None - ) - frames = refiner.refine( - lowres_video, - positive, - negative, - positive_mask, - negative_mask, - num_inference_steps=PPL_CONFIG["refiner_steps"], - guidance_scale=PPL_CONFIG["refiner_guidance_scale"], - shift=PPL_CONFIG["refiner_shift"], - t_thresh=PPL_CONFIG["refiner_t_thresh"], - tail_steps=PPL_CONFIG["refiner_tail_steps"], - clean_first_frame=clean_first_frame, - generator=torch.Generator("cuda").manual_seed(seed), - ) - video = frames[0].permute(1, 2, 3, 0).float().clamp(0.0, 1.0).cpu().numpy() - if task == "t2i": - Image.fromarray((video[0] * 255).round().astype("uint8")).save(output_path) - return {"output_path": output_path} - from diffusers.utils import export_to_video - - export_to_video(list(video), output_path, fps=PPL_CONFIG["fps"]) - return {"output_path": output_path} 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/distributed/fsdp.py b/telefuser/distributed/fsdp.py index 866e0dd..3b9c81b 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, ) diff --git a/telefuser/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py index 7f6c5d6..4cfa968 100644 --- a/telefuser/models/lingbot_video_dit.py +++ b/telefuser/models/lingbot_video_dit.py @@ -10,9 +10,11 @@ from typing import Any import torch +import torch.distributed as dist import torch.nn.functional as F from torch import nn +from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads from telefuser.ops.attention import attention @@ -122,11 +124,16 @@ def __init__(self, hidden_size: int, num_heads: int, norm_eps: float, qkv_bias: 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 ) -> torch.Tensor: @@ -136,6 +143,14 @@ def forward( 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: + query = ulysses_scatter_heads(query, group)() + key = ulysses_scatter_heads(key, group)() + value = ulysses_scatter_heads(value, group)() output = attention( query.transpose(1, 2), key.transpose(1, 2), @@ -147,7 +162,10 @@ def forward( ) if not isinstance(output, torch.Tensor): raise RuntimeError("LingBot attention does not support log-sum-exp outputs") - return self.to_out(output.transpose(1, 2).reshape(batch, sequence, -1).to(hidden_states.dtype)) + output = output.transpose(1, 2) + if use_ulysses: + output = ulysses_gather_heads(output, group, num_heads=self.num_heads)() + return self.to_out(output.reshape(batch, sequence, -1).to(hidden_states.dtype)) class LingBotVideoBlock(nn.Module): @@ -288,12 +306,61 @@ def __init__( 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) + 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, @@ -323,16 +390,31 @@ def forward( ) time_embedding = self.time_embedder(timestep) temb_input = time_embedding.unsqueeze(1).expand(-1, joint.shape[1], -1) + text_mask = ( + encoder_attention_mask.bool() + if encoder_attention_mask is not None + else torch.ones(batch, text_length, dtype=torch.bool, device=hidden_states.device) + ) + valid_mask = torch.cat( + (torch.ones(batch, video_tokens, dtype=torch.bool, device=hidden_states.device), text_mask), dim=1 + ) + joint, rotary, temb_input, local_valid_mask, padding = self._ulysses_shard_joint( + joint, rotary, temb_input, valid_mask + ) temb6 = self.time_modulation(temb_input) + global_valid_mask = self._ulysses_gather_sequence(local_valid_mask) attention_mask = None - if encoder_attention_mask is not None and not bool(encoder_attention_mask.bool().all()): - video_mask = torch.ones(batch, video_tokens, dtype=torch.bool, device=hidden_states.device) - attention_mask = torch.cat((video_mask, encoder_attention_mask.bool()), dim=1)[:, None, None, :] + if not bool(global_valid_mask.all()): + attention_mask = global_valid_mask[:, None, None, :] for block in self.blocks: joint = block(joint, temb6, rotary, attention_mask) 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] 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) diff --git a/telefuser/pipelines/lingbot_video/__init__.py b/telefuser/pipelines/lingbot_video/__init__.py index 48892c8..6a7f751 100644 --- a/telefuser/pipelines/lingbot_video/__init__.py +++ b/telefuser/pipelines/lingbot_video/__init__.py @@ -1,4 +1,4 @@ -"""LingBot-Video pipeline contracts and runtime components.""" +"""LingBot-Video reusable pipeline contracts and components.""" from .data import ( DEFAULT_NEGATIVE_PROMPT, @@ -32,7 +32,6 @@ prepare_refiner_latent, prepare_refiner_video, ) -from .runtime import build_lingbot_video_pipeline, build_lingbot_video_refiner_stage from .text_encoding import LingBotVideoTextEncodingStage from .vae import ( LingBotVideoVAEDecodeStage, @@ -50,8 +49,6 @@ "DEFAULT_NEGATIVE_PROMPT_IMAGE", "LingBotVideoPromptConditions", "LingBotVideoRefinerStage", - "build_lingbot_video_pipeline", - "build_lingbot_video_refiner_stage", "compute_refiner_sigmas", "load_refiner_first_frame", "load_refiner_video_file", diff --git a/telefuser/pipelines/lingbot_video/denoising.py b/telefuser/pipelines/lingbot_video/denoising.py index 05958ca..dee14a8 100644 --- a/telefuser/pipelines/lingbot_video/denoising.py +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -7,12 +7,17 @@ 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 -from telefuser.models.lingbot_video_dit import LingBotVideoTransformer3DModel +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_ulysses_group +from telefuser.distributed.fsdp import shard_model +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: @@ -70,14 +75,51 @@ class LingBotVideoDenoisingStage(BaseStage): """Dense LingBot-Video denoising stage with source-equivalent CFG order.""" def __init__( - self, name: str, transformer: LingBotVideoTransformer3DModel, model_runtime_config: ModelRuntimeConfig + 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 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 + + def parallel_models(self) -> None: + """Attach Ulysses SP and optional per-block FSDP to the DiT stage.""" + 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) + 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") + # ParallelWorker receives a CPU checkpoint. Move all state first so + # intentionally unsharded FP32 parameters do not remain on CPU. + self.transformer.to(self.device) + ignored_states = [ + parameter for parameter in self.transformer.parameters() if parameter.dtype != self.torch_dtype + ] + 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 FSDP for {self.name}") + self.transformer = shard_model( + module=self.transformer, + device_id=self.device.index if self.device.index is not None else torch.cuda.current_device(), + 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() @@ -96,14 +138,33 @@ def predict_noise_with_cfg( 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): - positive = self.transformer( - latents, model_timestep, positive_prompt_embeds, positive_attention_mask - ).float() if guidance_scale == 1.0: - return positive + 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 not 1") - negative = self.transformer( - latents, model_timestep, negative_prompt_embeds, negative_attention_mask - ).float() + matching_masks = (positive_attention_mask is None) == (negative_attention_mask is None) + matching_shapes = positive_prompt_embeds.shape == negative_prompt_embeds.shape + if self.batch_cfg and matching_masks and matching_shapes: + combined_latents = torch.cat((latents, latents)) + combined_timestep = torch.cat((model_timestep, model_timestep)) + combined_embeds = torch.cat((positive_prompt_embeds, negative_prompt_embeds)) + combined_mask = ( + None + if positive_attention_mask is None + else torch.cat((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) diff --git a/telefuser/pipelines/lingbot_video/pipeline.py b/telefuser/pipelines/lingbot_video/pipeline.py index e8d6c34..d1a2a57 100644 --- a/telefuser/pipelines/lingbot_video/pipeline.py +++ b/telefuser/pipelines/lingbot_video/pipeline.py @@ -6,12 +6,16 @@ from __future__ import annotations -from dataclasses import dataclass +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, @@ -30,14 +34,19 @@ ) -@dataclass(frozen=True) +@dataclass class LingBotVideoPipelineConfig: - """Runtime configuration shared by Dense and MoE pipelines.""" + """Configuration for LingBot-Video model stages and sampling.""" - model: LingBotVideoModelConfig = LingBotVideoModelConfig() + 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: @@ -64,7 +73,7 @@ class LingBotVideoGeneration: class LingBotVideoPipeline(BasePipeline): - """Lifecycle-aware facade for independently loaded LingBot-Video stages.""" + """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) @@ -74,12 +83,50 @@ def __init__(self, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat self.vae_encode_stage: LingBotVideoVAEEncodeStage | None = None self.vae_decode_stage: LingBotVideoVAEDecodeStage | None = None self.scheduler: FlowUniPCMultistepScheduler | None = None - self.model_dir: str | None = None self.variant: str = "dense" - def init(self, module_manager: object, config: LingBotVideoPipelineConfig) -> None: - del module_manager + 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.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.denoising_stage = ParallelWorker(denoising_stage) + else: + 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.""" + 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: @@ -87,25 +134,13 @@ def validate_request(request: LingBotVideoRequest) -> LingBotVideoRequest: validate_frame_count(request.num_frames) return request - def set_runtime( - self, - *, - text_stage: LingBotVideoTextEncodingStage, - denoising_stage: LingBotVideoDenoisingStage, - scheduler: FlowUniPCMultistepScheduler, - vae_decode_stage: LingBotVideoVAEDecodeStage | None = None, - vae_encode_stage: LingBotVideoVAEEncodeStage | None = None, - ) -> None: - """Attach independently loaded stages to this pipeline instance.""" - self.text_stage = text_stage - self.denoising_stage = denoising_stage - self.scheduler = scheduler - self.vae_decode_stage = vae_decode_stage - self.vae_encode_stage = vae_encode_stage - 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() @@ -179,14 +214,16 @@ def generate( for step in self.scheduler.timesteps: if condition is not None: latent = reinject_ti2v_condition(latent, condition, condition_mask) - prediction = self.denoising_stage.predict_noise_with_cfg( - latent, - step.expand(latent.shape[0]), - positive, - negative, - positive_mask, - negative_mask, - self.config.guidance_scale, + 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: @@ -195,4 +232,5 @@ def generate( return LingBotVideoGeneration(latent, prompt_conditions) if self.vae_decode_stage is None: raise RuntimeError("decode=True requires a configured VAE decode stage") - return LingBotVideoGeneration(self.vae_decode_stage.decode(latent), prompt_conditions) + 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 index 9216922..537cc9a 100644 --- a/telefuser/pipelines/lingbot_video/refiner.py +++ b/telefuser/pipelines/lingbot_video/refiner.py @@ -10,8 +10,10 @@ import torch import torch.nn.functional as F -from telefuser.pipelines.lingbot_video.denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition -from telefuser.pipelines.lingbot_video.vae import ( +from telefuser.platforms import current_platform + +from .denoising import LingBotVideoDenoisingStage, reinject_ti2v_condition +from .vae import ( LingBotVideoVAEDecodeStage, LingBotVideoVAEEncodeStage, first_frame_condition_mask, @@ -211,6 +213,28 @@ def __init__( 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, @@ -238,6 +262,10 @@ def refine( 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: @@ -263,16 +291,22 @@ def refine( for timestep in self.scheduler.timesteps: if condition is not None: latent = reinject_ti2v_condition(latent, condition, condition_mask) - prediction = 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, + 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) - return self.vae_decode_stage.decode(latent) + # 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() + return self._resolve_stage_result(self.vae_decode_stage.decode(latent)) diff --git a/telefuser/pipelines/lingbot_video/runtime.py b/telefuser/pipelines/lingbot_video/runtime.py deleted file mode 100644 index 31c35bd..0000000 --- a/telefuser/pipelines/lingbot_video/runtime.py +++ /dev/null @@ -1,129 +0,0 @@ -"""Native checkpoint assembly for the LingBot-Video pipeline.""" - -from __future__ import annotations - -from pathlib import Path - -import torch - -from telefuser.core.config import AttentionConfig, AttnImplType, ModelRuntimeConfig, OffloadConfig, WeightOffloadType -from telefuser.schedulers.unipc import FlowUniPCMultistepScheduler - -from .data import load_lingbot_video_model_config -from .denoising import LingBotVideoDenoisingStage -from .loading import load_lingbot_video_dense_transformer, load_lingbot_video_moe_transformer -from .pipeline import LingBotVideoPipeline, LingBotVideoPipelineConfig -from .refiner import LingBotVideoRefinerStage -from .text_encoding import LingBotVideoTextEncodingStage -from .vae import LingBotVideoVAEDecodeStage, LingBotVideoVAEEncodeStage - - -def build_lingbot_video_pipeline( - model_dir: str | Path, - *, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - variant: str = "dense", - cpu_offload: bool | None = None, - guidance_scale: float = 3.0, - num_inference_steps: int = 50, - shift: float = 3.0, - attention_config: AttentionConfig | None = None, -) -> LingBotVideoPipeline: - """Load official checkpoint components and assemble a native pipeline runtime. - - ``variant`` selects a Dense transformer or the MoE base transformer. The - refiner remains a separately loaded stage so its weights do not need to - coexist with the base model. - """ - if cpu_offload is None: - cpu_offload = variant == "moe" - if variant not in {"dense", "moe"}: - raise ValueError("variant must be 'dense' or 'moe'") - try: - from diffusers import AutoencoderKLWan - from transformers import AutoProcessor, Qwen3VLForConditionalGeneration - except ImportError as exc: - raise RuntimeError("LingBot-Video runtime requires diffusers and transformers") from exc - - root = Path(model_dir) - transformer_dir = root / "transformer" - runtime_config = ModelRuntimeConfig( - device_type=torch.device(device).type, - torch_dtype=torch_dtype, - attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), - offload_config=OffloadConfig( - offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD - ), - ) - load_device = "cpu" if cpu_offload else device - transformer = ( - load_lingbot_video_dense_transformer(transformer_dir, device=load_device, torch_dtype=torch_dtype) - if variant == "dense" - else load_lingbot_video_moe_transformer(transformer_dir, device=load_device, torch_dtype=torch_dtype) - ) - processor = AutoProcessor.from_pretrained(root / "processor") - text_encoder = ( - Qwen3VLForConditionalGeneration.from_pretrained( - root / "text_encoder", dtype=torch_dtype, attn_implementation="sdpa" - ) - .to(load_device) - .eval() - ) - vae = AutoencoderKLWan.from_pretrained(root / "vae", torch_dtype=torch.float32).to(load_device).eval() - scheduler = FlowUniPCMultistepScheduler.from_pretrained(root / "scheduler") - pipeline = LingBotVideoPipeline(device=device, torch_dtype=torch_dtype) - pipeline.model_dir = str(root) - pipeline.variant = variant - pipeline.init( - None, - LingBotVideoPipelineConfig( - model=load_lingbot_video_model_config(transformer_dir, variant=variant), - guidance_scale=guidance_scale, - num_inference_steps=num_inference_steps, - shift=shift, - ), - ) - pipeline.set_runtime( - text_stage=LingBotVideoTextEncodingStage("text_encoder", text_encoder, processor, runtime_config), - denoising_stage=LingBotVideoDenoisingStage("transformer", transformer, runtime_config), - vae_encode_stage=LingBotVideoVAEEncodeStage("vae_encode", vae, runtime_config), - vae_decode_stage=LingBotVideoVAEDecodeStage("vae_decode", vae, runtime_config), - scheduler=scheduler, - ) - return pipeline - - -def build_lingbot_video_refiner_stage( - model_dir: str | Path, - *, - device: str = "cuda", - torch_dtype: torch.dtype = torch.bfloat16, - cpu_offload: bool = True, - attention_config: AttentionConfig | None = None, -) -> LingBotVideoRefinerStage: - """Load the separately checkpointed MoE refiner and its VAE/scheduler stages.""" - try: - from diffusers import AutoencoderKLWan - except ImportError as exc: - raise RuntimeError("LingBot-Video refiner runtime requires diffusers") from exc - - root = Path(model_dir) - runtime_config = ModelRuntimeConfig( - device_type=torch.device(device).type, - torch_dtype=torch_dtype, - attention_config=attention_config or AttentionConfig.dense_attention(AttnImplType.TORCH_SDPA), - offload_config=OffloadConfig( - offload_type=WeightOffloadType.MODEL_CPU_OFFLOAD if cpu_offload else WeightOffloadType.NO_CPU_OFFLOAD - ), - ) - load_device = "cpu" if cpu_offload else device - transformer = load_lingbot_video_moe_transformer(root / "refiner", device=load_device, torch_dtype=torch_dtype) - vae = AutoencoderKLWan.from_pretrained(root / "vae", torch_dtype=torch.float32).to(load_device).eval() - scheduler = FlowUniPCMultistepScheduler.from_pretrained(root / "scheduler") - return LingBotVideoRefinerStage( - denoising_stage=LingBotVideoDenoisingStage("refiner", transformer, runtime_config), - vae_encode_stage=LingBotVideoVAEEncodeStage("refiner_vae_encode", vae, runtime_config), - vae_decode_stage=LingBotVideoVAEDecodeStage("refiner_vae_decode", vae, runtime_config), - scheduler=scheduler, - ) diff --git a/telefuser/pipelines/lingbot_video/text_encoding.py b/telefuser/pipelines/lingbot_video/text_encoding.py index fa00284..aabe86c 100644 --- a/telefuser/pipelines/lingbot_video/text_encoding.py +++ b/telefuser/pipelines/lingbot_video/text_encoding.py @@ -10,10 +10,10 @@ import torch from PIL import Image -from torch import nn 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 @@ -42,16 +42,15 @@ class LingBotVideoTextEncodingStage(BaseStage): def __init__( self, name: str, - text_encoder: nn.Module, - processor: Any, + 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 = text_encoder - self.processor = processor + 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 diff --git a/telefuser/pipelines/lingbot_video/vae.py b/telefuser/pipelines/lingbot_video/vae.py index 528ac9e..98fd7b8 100644 --- a/telefuser/pipelines/lingbot_video/vae.py +++ b/telefuser/pipelines/lingbot_video/vae.py @@ -10,6 +10,7 @@ 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( @@ -60,9 +61,9 @@ def first_frame_condition_mask(latent_frames: int, *, device: torch.device | str class LingBotVideoVAEEncodeStage(BaseStage): """Encode RGB video tensors using official LingBot latent normalization.""" - def __init__(self, name: str, vae: torch.nn.Module, model_runtime_config: ModelRuntimeConfig) -> None: + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: super().__init__(name, model_runtime_config) - self.vae = vae + self.vae = module_manager.fetch_module("vae") self.model_names = ["vae"] @with_model_offload(["vae"]) @@ -84,9 +85,9 @@ def encode(self, pixel_values: torch.Tensor, generator: torch.Generator | None = class LingBotVideoVAEDecodeStage(BaseStage): """Decode normalized LingBot diffusion latents to RGB video tensors.""" - def __init__(self, name: str, vae: torch.nn.Module, model_runtime_config: ModelRuntimeConfig) -> None: + def __init__(self, name: str, module_manager: ModuleManager, model_runtime_config: ModelRuntimeConfig) -> None: super().__init__(name, model_runtime_config) - self.vae = vae + self.vae = module_manager.fetch_module("vae") self.model_names = ["vae"] @with_model_offload(["vae"]) diff --git a/tests/unit/pipelines/lingbot_video/test_pipeline.py b/tests/unit/pipelines/lingbot_video/test_pipeline.py index 46e3b80..4318adf 100644 --- a/tests/unit/pipelines/lingbot_video/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -25,6 +25,40 @@ def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object) -> torch.Ten return torch.ones_like(latents) +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]) @@ -36,10 +70,11 @@ def step(self, prediction: torch.Tensor, timestep: torch.Tensor, latent: torch.T return latent - prediction -def test_pipeline_runs_t2v_sampling_with_injected_runtime() -> None: - pipeline = LingBotVideoPipeline(device="cpu") - pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) - pipeline.set_runtime(text_stage=_TextStage(), denoising_stage=_DenoisingStage(), scheduler=_Scheduler()) +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) @@ -69,7 +104,9 @@ def forward( def test_denoising_stage_returns_fp32_scheduler_prediction() -> None: stage = LingBotVideoDenoisingStage( - "denoising", _DtypeProbeTransformer(), ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32) + "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) @@ -81,6 +118,41 @@ def test_denoising_stage_returns_fp32_scheduler_prediction() -> None: assert torch.equal(prediction, torch.full_like(latents, 5.0)) +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, mask + self.calls += 1 + offset = embeds.mean(dim=(1, 2)).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 + + class _RecordingTextStage(_TextStage): def __init__(self) -> None: self.images: list[object | None] = [] @@ -100,9 +172,10 @@ def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tens return super().encode(caption, images) text_stage = _CapturingTextStage() - pipeline = LingBotVideoPipeline(device="cpu") - pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) - pipeline.set_runtime(text_stage=text_stage, denoising_stage=_DenoisingStage(), scheduler=_Scheduler()) + 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) @@ -122,14 +195,11 @@ def encode(self, pixel_values: torch.Tensor, generator: torch.Generator | None = def test_pipeline_ti2v_reuses_one_visual_condition_for_both_cfg_branches() -> None: text_stage = _RecordingTextStage() vae_encode_stage = _VAEEncodeStage() - pipeline = LingBotVideoPipeline(device="cpu") - pipeline.init(None, LingBotVideoPipelineConfig(num_inference_steps=2)) - pipeline.set_runtime( - text_stage=text_stage, - denoising_stage=_DenoisingStage(), - scheduler=_Scheduler(), - vae_encode_stage=vae_encode_stage, - ) + 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)), @@ -167,14 +237,12 @@ def test_pipeline_releases_stages_before_a_separate_refiner_is_loaded() -> None: denoising_stage = _OffloadStage() vae_encode_stage = _OffloadStage() vae_decode_stage = _OffloadStage() - pipeline = LingBotVideoPipeline(device="cpu") - pipeline.set_runtime( - text_stage=text_stage, - denoising_stage=denoising_stage, - scheduler=_Scheduler(), - vae_encode_stage=vae_encode_stage, - vae_decode_stage=vae_decode_stage, - ) + 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() diff --git a/tests/unit/pipelines/lingbot_video/test_refiner_stage.py b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py index a5f56a5..57aaa97 100644 --- a/tests/unit/pipelines/lingbot_video/test_refiner_stage.py +++ b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py @@ -21,6 +21,35 @@ def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object) -> torch.Ten 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 @@ -87,3 +116,58 @@ def test_refiner_accepts_captured_noise() -> None: ) 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 diff --git a/tests/unit/pipelines/lingbot_video/test_runtime.py b/tests/unit/pipelines/lingbot_video/test_runtime.py index 7218eb2..aa6bff7 100644 --- a/tests/unit/pipelines/lingbot_video/test_runtime.py +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -5,8 +5,9 @@ import sys import types +from examples.lingbot_video import lingbot_video_moe_30b as moe_example from telefuser.core.config import WeightOffloadType -from telefuser.pipelines.lingbot_video import runtime +from telefuser.pipelines.lingbot_video import LingBotVideoModelConfig class _FakeModule: @@ -24,14 +25,9 @@ def set_attention_config(self, attention_config) -> None: self.attention_config = attention_config -def test_moe_runtime_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: +def test_moe_example_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: loaded_transformers: list[tuple[object, str]] = [] - def load_dense(path, *, device: str, torch_dtype): - del torch_dtype - loaded_transformers.append((path, device)) - return _FakeModule() - def load_moe(path, *, device: str, torch_dtype): del torch_dtype loaded_transformers.append((path, device)) @@ -68,16 +64,18 @@ def from_pretrained(cls, *args, **kwargs): transformers.Qwen3VLForConditionalGeneration = _TextEncoder monkeypatch.setitem(sys.modules, "diffusers", diffusers) monkeypatch.setitem(sys.modules, "transformers", transformers) - monkeypatch.setattr(runtime, "load_lingbot_video_dense_transformer", load_dense) - monkeypatch.setattr(runtime, "load_lingbot_video_moe_transformer", load_moe) - monkeypatch.setattr(runtime, "load_lingbot_video_model_config", lambda *args, **kwargs: None) - monkeypatch.setattr(runtime.FlowUniPCMultistepScheduler, "from_pretrained", _Scheduler.from_pretrained) + monkeypatch.setattr(moe_example, "load_lingbot_video_moe_transformer", load_moe) + 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) - base = runtime.build_lingbot_video_pipeline(tmp_path, device="cuda", variant="moe") - refiner = runtime.build_lingbot_video_refiner_stage(tmp_path, device="cuda") + base = moe_example.build_pipeline(tmp_path, device="cuda") + refiner = moe_example.build_refiner(tmp_path, device="cuda") assert loaded_transformers == [(tmp_path / "transformer", "cpu"), (tmp_path / "refiner", "cpu")] - assert base.model_dir == str(tmp_path) assert base.variant == "moe" assert ( base.denoising_stage.transformer.attention_config is base.denoising_stage.model_runtime_config.attention_config diff --git a/tests/unit/pipelines/lingbot_video/test_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py index bd1d306..bb89bcb 100644 --- a/tests/unit/pipelines/lingbot_video/test_service_contract.py +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -1,4 +1,4 @@ -"""Tests for the LingBot-Video service entrypoint contract.""" +"""Tests for the split LingBot-Video CLI and service examples.""" from __future__ import annotations @@ -6,54 +6,104 @@ import json import sys from pathlib import Path -from types import SimpleNamespace +from types import ModuleType, SimpleNamespace import numpy as np import pytest import torch -from PIL import Image +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_service_module(): - path = Path("examples/lingbot_video/lingbot_video_service.py") - spec = importlib.util.spec_from_file_location("lingbot_video_service_test", path) + +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[spec.name] = module + sys.modules[module_name] = module spec.loader.exec_module(module) return module -def test_service_contract_declares_t2i_t2v_and_i2v() -> None: - module = _load_service_module() +@pytest.mark.parametrize("variant", ["dense", "moe"]) +def test_examples_expose_cli_and_service_entrypoints(variant: str) -> None: + module = _load_example(variant) - contract, is_explicit = load_pipeline_contract(module, ppl_file="lingbot_video_service.py", default_task="t2v") + for name in ("PPL_CONFIG", "CONTRACT", "get_pipeline", "run", "run_with_file"): + assert hasattr(module, name) + assert module.PIPELINE_CONTRACT is module.CONTRACT + 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").media_type == "image" assert contract.get_task_contract("t2i").parameters["negative_prompt"].default == DEFAULT_NEGATIVE_PROMPT_IMAGE - assert contract.get_task_contract("t2v").parameters["refine"].type == "boolean" 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_service_writes_single_image_for_t2i(tmp_path: Path) -> None: - module = _load_service_module() +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 moe_contract.get_task_contract("t2v").parameters["refine"].default is True + + +@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 + + +def test_dense_run_with_file_supports_service_t2i(tmp_path: Path) -> None: + module = _load_example("dense") class FakePipeline: - variant = "dense" + device = "cpu" def generate(self, request, **_: object) -> SimpleNamespace: assert request.num_frames == 1 assert request.caption == '{"scene":"test"}' - assert (request.height, request.width) == (480, 864) + assert (request.height, request.width) == (480, 832) return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2)) - output_path = tmp_path / "result.png" + output_path = tmp_path / "dense.png" result = module.run_with_file( FakePipeline(), json.dumps({"caption": {"scene": "test"}, "duration": 5}), @@ -65,83 +115,49 @@ def generate(self, request, **_: object) -> SimpleNamespace: assert output_path.is_file() -def test_service_writes_video_for_t2v(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - module = _load_service_module() - exported: list[object] = [] - - import diffusers.utils +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, output_path: str, fps: int) -> str: - del output_path, fps + def capture_export(frames: list[np.ndarray], output_path: str, fps: int) -> None: + del output_path + assert fps == 24 exported.extend(frames) - return "unused.mp4" - monkeypatch.setattr(diffusers.utils, "export_to_video", capture_export) + monkeypatch.setattr(module, "export_to_video", capture_export) class FakePipeline: - variant = "dense" + device = "cpu" def generate(self, request, **_: object) -> SimpleNamespace: - assert request.num_frames == 97 - assert (request.height, request.width) == (480, 864) + assert request.num_frames == 121 return SimpleNamespace(output=torch.full((1, 3, 5, 2, 2), 0.5)) - output_path = tmp_path / "result.mp4" result = module.run_with_file( FakePipeline(), json.dumps({"caption": {"scene": "test"}, "duration": 5}), task="t2v", - output_path=str(output_path), + output_path=str(tmp_path / "dense.mp4"), ) - assert result == {"output_path": str(output_path)} + 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_service_resolution_is_aligned_to_lingbot_vae_and_dit() -> None: - module = _load_service_module() - - for resolution in ("480p", "720p", "1080p", "2k", "4k"): - width, height = module._lingbot_video_size("16:9", resolution) - assert width % 16 == 0 - assert height % 16 == 0 - - -def test_service_rejects_i2v_without_first_image_path() -> None: - module = _load_service_module() - - with pytest.raises(ValueError, match="requires first_image_path"): - module.run_with_file( - object(), - json.dumps({"caption": {"scene": "test"}, "duration": 5}), - task="i2v", - ) - - -def test_service_runs_base_then_refiner_on_one_pipeline(tmp_path: Path, monkeypatch) -> None: - module = _load_service_module() - - class FakeTextStage: - def __init__(self) -> None: - self.calls = 0 - - def encode(self, caption: str) -> tuple[torch.Tensor, torch.Tensor]: - self.calls += 1 - assert caption in {'{"scene":"test"}', DEFAULT_NEGATIVE_PROMPT_IMAGE} - return torch.zeros(1, 1, 2), torch.ones(1, 1, dtype=torch.bool) +def test_moe_run_releases_base_then_runs_refiner(monkeypatch: pytest.MonkeyPatch) -> None: + module = _load_example("moe") class FakePipeline: - variant = "moe" - model_dir = "fake-model" - text_stage = FakeTextStage() + device = "cpu" + text_stage = object() def __init__(self) -> None: self.released = False def generate(self, request, **_: object) -> SimpleNamespace: - assert request.num_frames == 1 + del request conditions = SimpleNamespace( has_visual_condition=False, positive_prompt_embeds=torch.zeros(1, 1, 2), @@ -157,78 +173,53 @@ def release_gpu_resources(self) -> None: class FakeRefiner: def __init__(self) -> None: self.called = False + self.closed = False def refine(self, lowres_video, *args, **kwargs) -> torch.Tensor: del args, kwargs self.called = True - assert lowres_video.shape == (1, 3, 1, 2, 2) return lowres_video + def close(self) -> None: + self.closed = True + pipeline = FakePipeline() refiner = FakeRefiner() - monkeypatch.setattr(module, "build_lingbot_video_refiner_stage", lambda _: refiner) + 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) - output_path = tmp_path / "refined.png" - result = module.run_with_file( + frames = module.run( pipeline, json.dumps({"caption": {"scene": "test"}, "duration": 5}), task="t2i", refine=True, - output_path=str(output_path), ) - assert result == {"output_path": str(output_path)} + assert frames.shape == (1, 3, 1, 2, 2) assert pipeline.released assert refiner.called - assert pipeline.text_stage.calls == 0 - assert output_path.is_file() - - -def test_service_reencodes_text_only_conditions_for_ti2v_refiner(tmp_path: Path, monkeypatch) -> None: - module = _load_service_module() - - class FakeTextStage: - def __init__(self) -> None: - self.calls = 0 - - def encode(self, caption: str) -> tuple[torch.Tensor, torch.Tensor]: - self.calls += 1 - assert caption in {'{"scene":"test"}', DEFAULT_NEGATIVE_PROMPT} - return torch.zeros(1, 1, 2), torch.ones(1, 1, dtype=torch.bool) - - class FakePipeline: - variant = "moe" - model_dir = "fake-model" - text_stage = FakeTextStage() + assert refiner.closed - def generate(self, request, **_: object) -> SimpleNamespace: - assert request.image is not None - conditions = SimpleNamespace(has_visual_condition=True) - return SimpleNamespace(output=torch.zeros(1, 3, 1, 2, 2), prompt_conditions=conditions) - def release_gpu_resources(self) -> None: - return None +@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]]] = [] - class FakeRefiner: - def refine(self, lowres_video, *args, **kwargs) -> torch.Tensor: - del args, kwargs - return lowres_video + monkeypatch.setattr(module, "get_pipeline", lambda *args, **kwargs: pipeline) - monkeypatch.setattr(module, "build_lingbot_video_refiner_stage", lambda _: FakeRefiner()) - monkeypatch.setitem(module.PPL_CONFIG, "refiner_height", 2) - monkeypatch.setitem(module.PPL_CONFIG, "refiner_width", 2) - image_path = tmp_path / "first.png" - Image.new("RGB", (2, 2)).save(image_path) + 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"} - module.run_with_file( - FakePipeline(), - json.dumps({"caption": {"scene": "test"}, "duration": 5}), - task="i2v", - first_image_path=str(image_path), - refine=True, - output_path=str(tmp_path / "refined.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 FakePipeline.text_stage.calls == 2 + 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 index 1deed0c..0289573 100644 --- a/tests/unit/pipelines/lingbot_video/test_text_encoding.py +++ b/tests/unit/pipelines/lingbot_video/test_text_encoding.py @@ -5,6 +5,7 @@ import torch from telefuser.core.config import ModelRuntimeConfig +from telefuser.core.module_manager import ModuleManager from telefuser.pipelines.lingbot_video.text_encoding import LingBotVideoTextEncodingStage @@ -13,11 +14,11 @@ class _VisionEncoder: 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", - _VisionEncoder(), - processor=object(), - model_runtime_config=ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32), + "text", module_manager, ModelRuntimeConfig(device_type="cpu", torch_dtype=torch.float32) ) image = stage.prepare_ti2v_vlm_image(torch.zeros(1, 3, 1, 192, 320)) diff --git a/tests/unit/pipelines/lingbot_video/test_vae.py b/tests/unit/pipelines/lingbot_video/test_vae.py index 0ed6af5..8e3001c 100644 --- a/tests/unit/pipelines/lingbot_video/test_vae.py +++ b/tests/unit/pipelines/lingbot_video/test_vae.py @@ -6,6 +6,7 @@ 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, @@ -36,8 +37,10 @@ def decode(self, latents: torch.Tensor) -> SimpleNamespace: def test_vae_stages_apply_checkpoint_normalization() -> None: runtime = ModelRuntimeConfig(device_type="cpu") vae = _VAE() - encoded = LingBotVideoVAEEncodeStage("encode", vae, runtime).encode(torch.ones(1, 3, 1, 2, 2)) - decoded = LingBotVideoVAEDecodeStage("decode", vae, runtime).decode(encoded) + 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 diff --git a/tools/validation/benchmark_lingbot_video.py b/tools/validation/benchmark_lingbot_video.py index e211ad1..b43ae8b 100644 --- a/tools/validation/benchmark_lingbot_video.py +++ b/tools/validation/benchmark_lingbot_video.py @@ -4,6 +4,7 @@ import argparse import json +import sys import time from collections import defaultdict from collections.abc import Callable @@ -13,10 +14,15 @@ 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, - build_lingbot_video_pipeline, - build_lingbot_video_refiner_stage, default_negative_caption, load_lingbot_video_prompt, load_refiner_first_frame, @@ -135,12 +141,12 @@ def get_active_samples() -> BenchmarkSamples: 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: build_lingbot_video_pipeline( + lambda: pipeline_builder( args.model_dir, - variant=args.variant, guidance_scale=args.guidance_scale, num_inference_steps=args.steps, cpu_offload=False if args.disable_cpu_offload else None, @@ -197,7 +203,7 @@ def generate_once() -> torch.Tensor: refiner = _measure( setup_samples, "refiner_load", - lambda: build_lingbot_video_refiner_stage(args.model_dir, cpu_offload=not args.disable_cpu_offload), + lambda: build_refiner(args.model_dir, cpu_offload=not args.disable_cpu_offload), ) _instrument_method( get_active_samples, diff --git a/tools/validation/replay_lingbot_video_dense_reference.py b/tools/validation/replay_lingbot_video_dense_reference.py index ae0381d..ecb230b 100644 --- a/tools/validation/replay_lingbot_video_dense_reference.py +++ b/tools/validation/replay_lingbot_video_dense_reference.py @@ -17,6 +17,7 @@ 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.loading import load_lingbot_video_dense_transformer @@ -130,7 +131,10 @@ def replay_text( .to(device) .eval() ) - stage = LingBotVideoTextEncodingStage("text_encoder", text_encoder, processor, runtime_config) + 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 @@ -157,7 +161,8 @@ def replay_text( 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() - vae_encode = LingBotVideoVAEEncodeStage("vae_encode", vae, runtime_config) + 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, @@ -196,7 +201,9 @@ def _load_replay_runtime( device = torch.device("cuda") runtime_config = ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) transformer = load_lingbot_video_dense_transformer(model_dir / "transformer", device=device) - denoising = LingBotVideoDenoisingStage("transformer", transformer, runtime_config) + module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) + module_manager.add_module(transformer, name="transformer") + 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 @@ -273,7 +280,9 @@ def _replay_with_runtime( torch.tensor(vae.config.latents_std, device=device), ) _record(results, "vae_decode_input", _load_tensor(reference_dir, metadata, "vae_decode_input"), raw_latent) - decoder = LingBotVideoVAEDecodeStage("vae_decode", vae, runtime_config) + 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() 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_native_parallel.py b/tools/validation/run_lingbot_video_native_parallel.py new file mode 100644 index 0000000..d9a4734 --- /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, + generation.prompt_conditions.negative_prompt_embeds, + generation.prompt_conditions.positive_attention_mask, + generation.prompt_conditions.negative_attention_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, + 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 index ce52026..af50f4d 100644 --- a/tools/validation/run_lingbot_video_refiner_core_parity.py +++ b/tools/validation/run_lingbot_video_refiner_core_parity.py @@ -10,6 +10,7 @@ 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.loading import load_lingbot_video_moe_transformer from telefuser.pipelines.lingbot_video.refiner import ( @@ -130,9 +131,11 @@ def main() -> None: local = load_lingbot_video_moe_transformer(DIRECTORY, device=DEVICE, torch_dtype=torch.bfloat16) local_scheduler = FlowUniPCMultistepScheduler.from_pretrained(ROOT / "scheduler") + module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) + module_manager.add_module(local, name="transformer") stage = LingBotVideoRefinerStage( denoising_stage=LingBotVideoDenoisingStage( - "refiner", local, ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) + "refiner", module_manager, ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) ), vae_encode_stage=_Encode(), vae_decode_stage=_Decode(), 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_output_handoff.py b/tools/validation/validate_lingbot_video_refiner_output_handoff.py index 0745e40..cfe385b 100644 --- a/tools/validation/validate_lingbot_video_refiner_output_handoff.py +++ b/tools/validation/validate_lingbot_video_refiner_output_handoff.py @@ -4,6 +4,7 @@ import argparse import json +import sys import tempfile import time from pathlib import Path @@ -12,10 +13,13 @@ 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, - build_lingbot_video_pipeline, - build_lingbot_video_refiner_stage, default_negative_caption, load_lingbot_video_prompt, load_refiner_first_frame, @@ -196,9 +200,8 @@ def main() -> None: args.negative_caption if args.negative_caption is not None else default_negative_caption(num_frames) ) pipeline, base_load_seconds = _measure( - lambda: build_lingbot_video_pipeline( + lambda: build_pipeline( args.model_dir, - variant="moe", guidance_scale=args.guidance_scale, num_inference_steps=args.steps, cpu_offload=False if args.disable_cpu_offload else None, @@ -218,8 +221,8 @@ def main() -> None: ) ) frames = generation.output - if pipeline.text_stage is None or pipeline.model_dir is None: - raise RuntimeError("LingBot refiner requires the base text stage and checkpoint metadata") + 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) negative, negative_mask = pipeline.text_stage.encode(negative_caption) @@ -230,7 +233,7 @@ def main() -> None: negative_mask = generation.prompt_conditions.negative_attention_mask _, base_release_seconds = _measure(pipeline.release_gpu_resources) refiner, refiner_load_seconds = _measure( - lambda: build_lingbot_video_refiner_stage(args.model_dir, cpu_offload=not args.disable_cpu_offload) + 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 From 3b61107ee768b0acf0e99d70a1a15f773d87c520 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 14:05:36 +0000 Subject: [PATCH 3/9] perf(lingbot-video): align four-gpu MoE inference - add inference-only FSDP2 sharding and source-style Ulysses communication - support packed SDPA and distributed MoE routing with grouped-GEMM fallback behavior - cap per-rank CPU intra-op threads and avoid unnecessary worker cache flushes - document and test the validated four-GPU execution constraints Verification: - git diff --cached --check --- CLAUDE.md | 3 +- docs/en/configuration.md | 1 + docs/en/lingbot_video.md | 45 +++--- docs/en/parallel.md | 8 +- docs/zh/configuration.md | 1 + docs/zh/lingbot_video.md | 32 +++-- docs/zh/parallel.md | 6 +- examples/lingbot_video/README.md | 10 +- .../lingbot_video/lingbot_video_moe_30b.py | 23 ++- telefuser/core/config.py | 3 + telefuser/distributed/fsdp.py | 30 ++++ telefuser/distributed/ulysses_comm.py | 15 ++ telefuser/models/lingbot_video_dit.py | 117 +++++++++++---- telefuser/models/lingbot_video_moe.py | 8 +- telefuser/ops/__init__.py | 3 +- telefuser/ops/attention/attention_impl.py | 40 +++++- telefuser/ops/moe.py | 136 +++++++++++++++++- .../pipelines/lingbot_video/denoising.py | 127 ++++++++++++++-- telefuser/pipelines/lingbot_video/pipeline.py | 36 +++-- telefuser/worker/parallel_worker.py | 9 +- tests/unit/core/test_config.py | 7 + tests/unit/core/test_config_serializer.py | 1 + tests/unit/models/test_lingbot_video_dit.py | 29 ++++ tests/unit/models/test_lingbot_video_moe.py | 35 +++++ .../pipelines/lingbot_video/test_pipeline.py | 75 +++++++++- .../pipelines/lingbot_video/test_runtime.py | 5 + 26 files changed, 698 insertions(+), 107 deletions(-) diff --git a/CLAUDE.md b/CLAUDE.md index 007bb78..9d562fb 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -88,7 +88,8 @@ telefuser/ - 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. Grouped-GEMM, FP8, and distributed execution require separate parity and benchmark evidence before being enabled. +- 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 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/lingbot_video.md b/docs/en/lingbot_video.md index b45e430..863e8dc 100644 --- a/docs/en/lingbot_video.md +++ b/docs/en/lingbot_video.md @@ -20,8 +20,10 @@ transformer = load_lingbot_video_dense_transformer( The MoE and refiner directories are sharded. Use `load_lingbot_video_moe_transformer` for either `transformer/` or `refiner/`. The default sorted eager expert path preserves upstream route ordering and keeps -a `where`-based diagnostic fallback. It is a validated single-GPU BF16 path, -but not a grouped-GEMM, FP8, or distributed production-throughput backend. For +a `where`-based diagnostic fallback. It remains the validated single-GPU BF16 +path. Four-GPU MoE execution defaults to native grouped GEMM when the installed +CUDA PyTorch build exposes `torch._grouped_mm`; the example accepts +`expert_backend=sorted` as an explicit fallback. For `variant="moe"`, the runtime defaults to stage CPU offload so the base DiT, text encoder, VAE, and separately loaded refiner do not need to reside on one GPU. Set `cpu_offload=False` only when GPU capacity is known to be sufficient. @@ -111,18 +113,24 @@ telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 - This configuration keeps FSDP-enabled DiT weights resident on GPU and therefore requires `cpu_offload=False`. Checkpoint FP32 modulation parameters are retained as replicated FSDP ignored states, matching the upstream mixed-precision layout. -The MoE 30B model uses its own example; its sorted eager experts are still -replicated within each FSDP block rather than expert-parallelized: +The MoE 30B model uses its own example. Its grouped expert weights are sharded +with their containing FSDP blocks rather than expert-parallelized: ```bash telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 ``` -The validated 832x480, 121-frame, 40-step MoE base run used 68.1 GiB peak GPU -memory and 240.4 seconds for generation excluding checkpoint load and MP4 -encoding. The separately loaded refiner also has a validated four-GPU -FSDP/Ulysses path. Expert parallelism, grouped GEMM, and FP8 remain disabled -pending their own parity and throughput evidence. +On PyTorch 2.11 with CUDA 13, the grouped-GEMM 832x480, 121-frame, 40-step MoE +base run uses one worker-owned scheduler loop and one intra-op CPU thread per +rank, matching `torchrun`. On four H100 GPUs, steady sampling is approximately +2.3 seconds per step and 93 seconds for 40 steps, compared with 92 seconds for +the source runner. The first TeleFuser request also pays about 15 seconds for +FSDP materialization. A cold CLI invocation took 231.9 seconds because the +parent load and spawn-time 30B model transfer are not included in steady +service latency; persistent service requests amortize that setup. The +separately loaded refiner also has a validated four-GPU FSDP/Ulysses path. +Expert parallelism and FP8 remain disabled pending their own parity and +throughput evidence. ## TI2V @@ -256,9 +264,9 @@ When `--negative-caption` is omitted, the benchmark uses the same source-compati T2I or video negative caption as the pipeline, CLI, and service. Pass an explicit empty string only when intentionally benchmarking that semantic override. For a base-plus-refiner run, it also records the serial base release and refiner -load phases. The default sorted eager MoE path is source-equivalent and has an -explicit diagnostic fallback, but it is not a grouped-GEMM or FP8 -production-throughput backend. +load phases. Single-GPU benchmarks use the source-equivalent sorted eager MoE +path. Four-GPU benchmarks select grouped GEMM by default; use +`expert_backend=sorted` only when measuring the correctness fallback. For T2I/T2V, the refiner reuses the exact CFG text conditions from the base generation and reports this as `refiner_prompt_conditions_reused`; TI2V keeps the source-compatible text-only refiner encoding path. @@ -270,10 +278,11 @@ zero-drift numerical-oracle gate instead of only writing metrics. The numerical-oracle path requires CUDA, PyTorch, Diffusers, Transformers, and the checkpoint components `transformer/`, `text_encoder/`, `processor/`, `vae/`, and -`scheduler/`. Dense runs source-equivalently on one GPU. The current MoE expert -implementation is an eager correctness path; do not use it as a 30B production -backend. Dense and MoE base DiTs have validated four-GPU FSDP/Ulysses SP paths. +`scheduler/`. Dense runs source-equivalently on one GPU. MoE keeps sorted eager +for single-GPU correctness and uses native grouped GEMM for four-GPU throughput; +the latter requires a CUDA PyTorch build exposing `torch._grouped_mm`. Dense and +MoE base DiTs have validated four-GPU FSDP/Ulysses SP paths. The MoE refiner has a validated independent four-GPU FSDP/Ulysses SP stage. -FlashAttention, MoE expert parallelism, grouped GEMM, and FP8 experts are -intentionally not enabled. The base-plus-refiner stage lifecycle remains -serial so their 30B weights never coexist on GPU. +External FlashAttention, MoE expert parallelism, and FP8 experts are not enabled. +The base-plus-refiner stage lifecycle remains serial so their 30B weights never +coexist on GPU. 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/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/lingbot_video.md b/docs/zh/lingbot_video.md index 350f2bf..e9da96e 100644 --- a/docs/zh/lingbot_video.md +++ b/docs/zh/lingbot_video.md @@ -16,7 +16,9 @@ transformer = load_lingbot_video_dense_transformer("/path/to/transformer") MoE 与 refiner 使用分片 safetensors，调用 `load_lingbot_video_moe_transformer` 加载 `transformer/` 或 `refiner/`。 默认的 sorted eager expert 路径保持上游 route 顺序，并保留 `where` 诊断 fallback。 -它是经过验证的单卡 BF16 路径，但不是 grouped-GEMM、FP8 或分布式生产吞吐后端。 +它仍是经过验证的单卡 BF16 路径。四卡 MoE 默认在 CUDA PyTorch 提供 +`torch._grouped_mm` 时启用原生 grouped GEMM；可通过 +`expert_backend=sorted` 显式回退。 `variant="moe"` 默认启用 stage CPU offload，避免 base DiT、text encoder、VAE 与 独立 refiner 同时驻留单张 GPU。仅在显存明确充足时才设置 `cpu_offload=False`。 @@ -89,15 +91,20 @@ telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 - 该模式要求 DiT 常驻 GPU，因此必须 `cpu_offload=False`。checkpoint 中保留 FP32 的 modulation 参数会作为 FSDP ignored state 在每个 rank 副本化，与上游 mixed-precision 布局一致。MoE 30B -使用独立 example；sorted eager expert 仍随 FSDP block 分片，而不是 expert parallel： +使用独立 example；grouped expert 权重仍随 FSDP block 分片，而不是 expert parallel： ```bash telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 ``` -已验证的 MoE base 为 832x480、121 帧、40 steps，生成阶段（不含 checkpoint 加载与 MP4 编码） -耗时 240.4 秒，峰值显存 68.1 GiB/卡。独立加载的 refiner 也已有经验证的四卡 -FSDP/Ulysses 路径；expert parallel、grouped GEMM 与 FP8 仍需独立 parity 与吞吐验证。 +在 PyTorch 2.11/CUDA 13 环境中，已验证的 grouped-GEMM MoE base 为 832x480、 +121 帧、40 steps。四卡路径由 worker 持有完整 scheduler loop，并与 `torchrun` +一致地为每个 rank 使用一个 CPU intra-op 线程。四张 H100 上稳态约为 2.3 秒/step， +40 steps 约 93 秒；官方 runner 为 92 秒。TeleFuser 首次请求还包含约 15 秒的 FSDP +materialization。冷启动 CLI 总耗时为 231.9 秒，因为父进程加载和 spawn 时传输 30B +模型不属于稳态服务延迟；常驻服务的后续请求会摊销这部分开销。独立加载的 refiner +也已有经验证的四卡 FSDP/Ulysses 路径；expert parallel 与 FP8 仍需独立 parity 与 +吞吐验证。 ## 服务 @@ -196,8 +203,9 @@ checkpoint inspection 工具会执行正常的严格加载，并记录已消费 基准工具会将一次性 setup 与 warmup、measured 指标分开记录，覆盖 checkpoint load、text encoding、每个 denoising step、VAE、refiner、输出编码与峰值 GPU 显存。启用优化前应先用它建立基线，并将全分辨率测量与 smoke run 分开记录。 未传入 `--negative-caption` 时，benchmark 会与 pipeline、CLI、service 一样使用上游兼容的 T2I 或视频负向 caption。只有需要特意评测该语义覆盖时，才显式传入空字符串。 -对于 base+refiner，它还会记录串行的 base release 与 refiner load 阶段。默认 sorted eager -MoE 路径与 source 等价，并有显式诊断 fallback，但不是 grouped-GEMM 或 FP8 生产吞吐后端。 +对于 base+refiner，它还会记录串行的 base release 与 refiner load 阶段。单卡 benchmark +使用 source-equivalent sorted eager MoE；四卡 benchmark 默认选择 grouped GEMM，仅在评测 +正确性 fallback 时显式设置 `expert_backend=sorted`。 T2I/T2V 的 Refiner 会复用 base generation 生成的完全相同 CFG text condition，并在报告中 记录 `refiner_prompt_conditions_reused`；TI2V 保持上游兼容的 text-only Refiner 编码路径。 Refiner core CLI 会向上游与 TeleFuser 的 low-noise 路径注入相同的 latent、noise、prompt 与 frame-zero condition；它会先 offload 上游 DiT，再加载 TeleFuser DiT，避免两个 30B 模型同时占用 GPU。 @@ -207,8 +215,8 @@ Refiner core CLI 会向上游与 TeleFuser 的 low-noise 路径注入相同的 l 数值 oracle 路径需要 CUDA、PyTorch、Diffusers、Transformers 以及 checkpoint 的 `transformer/`、`text_encoder/`、`processor/`、`vae/`、`scheduler/` 组件。Dense -可在单 GPU 上按 source-equivalent 路径运行。当前 MoE expert 是 eager correctness -实现，不应作为 30B 的生产后端。Dense 与 MoE base 都已有经验证的四卡 FSDP/Ulysses SP 路径； -MoE refiner 也已有独立的四卡 FSDP/Ulysses SP stage。FlashAttention、MoE expert parallel、 -grouped GEMM 与 FP8 expert 尚未启用。base+refiner 的 stage 生命周期保持串行，确保两个 -30B 权重不会同时驻留 GPU。 +可在单 GPU 上按 source-equivalent 路径运行。MoE 单卡保留 sorted eager correctness 路径， +四卡使用原生 grouped GEMM；后者要求 CUDA PyTorch 提供 `torch._grouped_mm`。Dense 与 +MoE base 都已有经验证的四卡 FSDP/Ulysses SP 路径；MoE refiner 也已有独立的四卡 +FSDP/Ulysses SP stage。外部 FlashAttention、MoE expert parallel 与 FP8 expert 尚未启用。 +base+refiner 的 stage 生命周期保持串行，确保两个 30B 权重不会同时驻留 GPU。 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/examples/lingbot_video/README.md b/examples/lingbot_video/README.md index 71734ba..04504ac 100644 --- a/examples/lingbot_video/README.md +++ b/examples/lingbot_video/README.md @@ -28,10 +28,12 @@ python examples/lingbot_video/lingbot_video_moe_30b.py \ ``` The command releases base-stage GPU weights before loading the separate refiner. -Use `--no-refine` for MoE base-only generation. The default sorted eager MoE -Use `--refiner_gpu_num` and `--refiner_batch_cfg` to override the corresponding MoE `PPL_CONFIG` defaults for a direct CLI run. -route path is source-equivalent with a diagnostic fallback, but it is not a -grouped-GEMM or FP8 production-performance backend. +Use `--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. +Use `--refiner_gpu_num` and `--refiner_batch_cfg` to override the corresponding +MoE `PPL_CONFIG` defaults for a direct CLI run. Set `PPL_CONFIG["model_root"]` in the selected example, then serve structured-caption T2I/T2V/TI2V requests with: diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py index d576844..bcd0cfd 100644 --- a/examples/lingbot_video/lingbot_video_moe_30b.py +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -59,6 +59,7 @@ "resolution": "480p", "aspect_ratio": "16:9", "seed": 42, + "expert_backend": "auto", "enable_refiner": True, "refiner_height": 1088, "refiner_width": 1920, @@ -121,6 +122,7 @@ def build_pipeline( 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: @@ -149,9 +151,18 @@ def build_pipeline( 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 {"grouped_mm", "sorted"}: + raise ValueError("LingBot MoE expert_backend must be 'auto', '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") + transformer = load_lingbot_video_moe_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype) + transformer.set_expert_execution_backend(expert_backend) + module_manager = ModuleManager(torch_dtype=torch_dtype, device="cpu") module_manager.add_module( - load_lingbot_video_moe_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype), + transformer, name="transformer", path=str(transformer_dir), ) @@ -275,6 +286,7 @@ def get_pipeline( model_root: str = PPL_CONFIG["model_root"], refiner_parallelism: int | None = None, refiner_batch_cfg: 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}: @@ -287,6 +299,7 @@ def get_pipeline( sp_ulysses_degree=parallelism, enable_fsdp=parallelism > 1, ), + expert_backend=expert_backend, ) refiner_parallelism = ( int(PPL_CONFIG["refiner_parallelism"]) if refiner_parallelism is None else refiner_parallelism @@ -476,6 +489,12 @@ def run_with_file( @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", "grouped_mm", "sorted"]), + help="MoE expert backend; auto selects grouped_mm for four GPUs", +) @click.option("--resolution", default=PPL_CONFIG["resolution"]) @click.option("--aspect_ratio", default=PPL_CONFIG["aspect_ratio"]) @click.option("--target_video_length", default=PPL_CONFIG["target_video_length"], type=int) @@ -493,6 +512,7 @@ def main( prompt: str, negative_prompt: str | None, seed: int, + expert_backend: str, resolution: str, aspect_ratio: str, target_video_length: int, @@ -509,6 +529,7 @@ def main( model_root, refiner_parallelism=refiner_gpu_num, refiner_batch_cfg=refiner_batch_cfg, + expert_backend=expert_backend, ) try: result = run_with_file( 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/distributed/fsdp.py b/telefuser/distributed/fsdp.py index 3b9c81b..69ec0a4 100644 --- a/telefuser/distributed/fsdp.py +++ b/telefuser/distributed/fsdp.py @@ -146,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 index 4cfa968..6e04af0 100644 --- a/telefuser/models/lingbot_video_dit.py +++ b/telefuser/models/lingbot_video_dit.py @@ -14,7 +14,7 @@ import torch.nn.functional as F from torch import nn -from telefuser.distributed.ulysses_comm import ulysses_gather_heads, ulysses_scatter_heads +from telefuser.distributed.ulysses_comm import ulysses_all_to_all_split_cat from telefuser.ops.attention import attention @@ -135,7 +135,11 @@ def set_ulysses_group(self, group: dist.ProcessGroup | None) -> None: self.ulysses_group = group def forward( - self, hidden_states: torch.Tensor, rotary_emb: torch.Tensor, attention_mask: torch.Tensor | None = None + 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) @@ -148,15 +152,33 @@ def forward( group is not None and dist.is_available() and dist.is_initialized() and dist.get_world_size(group) > 1 ) if use_ulysses: - query = ulysses_scatter_heads(query, group)() - key = ulysses_scatter_heads(key, group)() - value = ulysses_scatter_heads(value, group)() + 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", ) @@ -164,7 +186,12 @@ def forward( raise RuntimeError("LingBot attention does not support log-sum-exp outputs") output = output.transpose(1, 2) if use_ulysses: - output = ulysses_gather_heads(output, group, num_heads=self.num_heads)() + 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)) @@ -195,6 +222,7 @@ def forward( 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): @@ -202,7 +230,12 @@ def forward( 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) + 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 ) @@ -379,35 +412,60 @@ def forward( 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) - text = self.text_embedder(encoder_hidden_states) - joint = torch.cat((video, text), dim=1) - text_length = text.shape[1] - positions = make_lingbot_video_joint_position_ids(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) - .expand(batch, -1, -1) - ) - time_embedding = self.time_embedder(timestep) - temb_input = time_embedding.unsqueeze(1).expand(-1, joint.shape[1], -1) + 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, text_length, dtype=torch.bool, device=hidden_states.device) - ) - valid_mask = torch.cat( - (torch.ones(batch, video_tokens, dtype=torch.bool, device=hidden_states.device), text_mask), dim=1 + 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) - global_valid_mask = self._ulysses_gather_sequence(local_valid_mask) attention_mask = None - if not bool(global_valid_mask.all()): - attention_mask = global_valid_mask[:, None, 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) + 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)) @@ -415,6 +473,11 @@ def forward( projected = self._ulysses_gather_sequence(projected) if padding: projected = projected[:, :-padding] - projected = projected[:, :video_tokens] + 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) diff --git a/telefuser/models/lingbot_video_moe.py b/telefuser/models/lingbot_video_moe.py index 22a07c6..5aa6bac 100644 --- a/telefuser/models/lingbot_video_moe.py +++ b/telefuser/models/lingbot_video_moe.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from torch import nn -from telefuser.ops.moe import route_topk +from telefuser.ops.moe import grouped_expert_forward, route_topk from .lingbot_video_dit import LingBotVideoBlock, LingBotVideoMLP, LingBotVideoTransformer3DModel @@ -64,11 +64,13 @@ def __init__( def set_execution_backend(self, backend: str) -> None: """Select the source-style sorted path or the diagnostic where fallback.""" - if backend not in {"sorted", "where"}: - raise ValueError("execution backend must be 'sorted' or 'where'") + if backend not in {"grouped_mm", "sorted", "where"}: + raise ValueError("execution backend must be 'grouped_mm', 'sorted', or 'where'") self.execution_backend = backend def forward(self, tokens: torch.Tensor, indices: torch.Tensor, weights: torch.Tensor) -> torch.Tensor: + 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) diff --git a/telefuser/ops/__init__.py b/telefuser/ops/__init__.py index d99594b..2b76b31 100644 --- a/telefuser/ops/__init__.py +++ b/telefuser/ops/__init__.py @@ -12,7 +12,7 @@ from .base import CustomOp, CustomOpFunction from .custom_op import TritonKernelWrapper, register_custom_op -from .moe import route_topk +from .moe import grouped_expert_forward, route_topk from .normalization import AdaLayerNormContinuous, LayerNorm, RMSNorm, fused_scale_shift, modulate from .rotary import apply_rotary_emb @@ -30,6 +30,7 @@ "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..5ad7be9 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. @@ -206,7 +232,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 +296,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 index d5069de..32c3ee2 100644 --- a/telefuser/ops/moe.py +++ b/telefuser/ops/moe.py @@ -1,8 +1,9 @@ -"""Correctness-first Mixture-of-Experts routing primitives.""" +"""Mixture-of-Experts routing and grouped expert primitives.""" from __future__ import annotations import torch +import torch.nn.functional as F def route_topk( @@ -38,12 +39,133 @@ def route_topk( 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).values.sum(dim=-1) - selected_groups = group_values.topk(top_k_groups, dim=-1).indices - group_mask = torch.zeros_like(grouped, dtype=torch.bool) - group_mask.scatter_(1, selected_groups.unsqueeze(-1).expand(-1, -1, experts_per_group), True) - masked = selection_scores.masked_fill(~group_mask.reshape(tokens, experts), float("-inf")) - indices = masked.topk(top_k, dim=-1).indices + 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]) + + sorted_weights = flat_weights[sorted_positions].to(expert_output.dtype) + weighted_output = expert_output * sorted_weights.unsqueeze(-1) + routed = torch.zeros( + num_tokens * top_k, + hidden_size, + dtype=expert_output.dtype, + device=expert_output.device, + ) + routed[sorted_positions] = weighted_output + return routed.reshape(num_tokens, top_k, hidden_size).sum(dim=1).to(tokens.dtype) + + +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/denoising.py b/telefuser/pipelines/lingbot_video/denoising.py index dee14a8..3c99ade 100644 --- a/telefuser/pipelines/lingbot_video/denoising.py +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -15,7 +15,7 @@ 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_ulysses_group -from telefuser.distributed.fsdp import shard_model +from telefuser.distributed.fsdp import shard_model_fsdp2_inference from telefuser.platforms import current_platform from telefuser.utils.logging import logger @@ -48,6 +48,39 @@ def transformer_timestep(timestep: torch.Tensor, transformer_dtype: torch.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, @@ -85,14 +118,26 @@ def __init__( 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 @@ -103,18 +148,15 @@ def parallel_models(self) -> None: 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") - # ParallelWorker receives a CPU checkpoint. Move all state first so - # intentionally unsharded FP32 parameters do not remain on CPU. - self.transformer.to(self.device) ignored_states = [ parameter for parameter in self.transformer.parameters() if parameter.dtype != self.torch_dtype ] 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 FSDP for {self.name}") - self.transformer = shard_model( + logger.info(f"enabled LingBot block FSDP2 for {self.name}") + self.transformer = shard_model_fsdp2_inference( module=self.transformer, - device_id=self.device.index if self.device.index is not None else torch.cuda.current_device(), + device_mesh=self.transformer.device_mesh, wrap_module_names=self.transformer.get_fsdp_module_names(), ignored_states=ignored_states, ) @@ -134,6 +176,27 @@ def predict_noise_with_cfg( 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} @@ -144,16 +207,14 @@ def predict_noise_with_cfg( ).float() if negative_prompt_embeds is None: raise ValueError("negative_prompt_embeds is required when guidance_scale is not 1") - matching_masks = (positive_attention_mask is None) == (negative_attention_mask is None) - matching_shapes = positive_prompt_embeds.shape == negative_prompt_embeds.shape - if self.batch_cfg and matching_masks and matching_shapes: + if self.batch_cfg: combined_latents = torch.cat((latents, latents)) combined_timestep = torch.cat((model_timestep, model_timestep)) - combined_embeds = torch.cat((positive_prompt_embeds, negative_prompt_embeds)) - combined_mask = ( - None - if positive_attention_mask is None - else torch.cat((positive_attention_mask, negative_attention_mask)) + 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) @@ -168,3 +229,39 @@ def predict_noise_with_cfg( 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, + 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 index d1a2a57..03a59f1 100644 --- a/telefuser/pipelines/lingbot_video/pipeline.py +++ b/telefuser/pipelines/lingbot_video/pipeline.py @@ -210,24 +210,40 @@ def generate( dtype=torch.float32, generator=generator, ) - 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( + if self.config.enable_denoising_parallel: + latent = self._resolve_stage_result( + self.denoising_stage.denoise( latent, - step.expand(latent.shape[0]), positive, negative, positive_mask, negative_mask, self.config.guidance_scale, + self.config.num_inference_steps, + self.config.shift, + condition, + condition_mask, ) ) - latent = self.scheduler.step(prediction, step, latent) - if condition is not None: - latent = reinject_ti2v_condition(latent, 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: 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 index 845d3b2..bbe16e3 100644 --- a/tests/unit/models/test_lingbot_video_dit.py +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -66,3 +66,32 @@ def test_transformer_rejects_non_patch_aligned_latent_geometry() -> None: 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) diff --git a/tests/unit/models/test_lingbot_video_moe.py b/tests/unit/models/test_lingbot_video_moe.py index b8ff12b..2c0fa66 100644 --- a/tests/unit/models/test_lingbot_video_moe.py +++ b/tests/unit/models/test_lingbot_video_moe.py @@ -1,5 +1,6 @@ from __future__ import annotations +import pytest import torch import torch.nn.functional as F @@ -94,3 +95,37 @@ def test_grouped_expert_backends_are_exactly_equivalent() -> None: 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), + ) + + +@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) + + torch.testing.assert_close(grouped_output, sorted_output, rtol=2e-2, atol=2e-2) diff --git a/tests/unit/pipelines/lingbot_video/test_pipeline.py b/tests/unit/pipelines/lingbot_video/test_pipeline.py index 4318adf..6456515 100644 --- a/tests/unit/pipelines/lingbot_video/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -25,6 +25,15 @@ def predict_noise_with_cfg(self, latents: torch.Tensor, *_: object) -> torch.Ten 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 @@ -90,6 +99,20 @@ def test_initialized_pipeline_runs_t2v_sampling() -> None: ) +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__() @@ -127,9 +150,14 @@ def __init__(self) -> None: def forward( self, latents: torch.Tensor, timestep: torch.Tensor, embeds: torch.Tensor, mask: torch.Tensor | None ) -> torch.Tensor: - del timestep, mask + del timestep self.calls += 1 - offset = embeds.mean(dim=(1, 2)).reshape(-1, 1, 1, 1, 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 @@ -153,6 +181,49 @@ def test_denoising_stage_batched_cfg_matches_two_source_order_forwards() -> None assert batched_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] = [] diff --git a/tests/unit/pipelines/lingbot_video/test_runtime.py b/tests/unit/pipelines/lingbot_video/test_runtime.py index aa6bff7..0f9fc08 100644 --- a/tests/unit/pipelines/lingbot_video/test_runtime.py +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -13,6 +13,7 @@ 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 @@ -24,6 +25,9 @@ def eval(self) -> _FakeModule: 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 test_moe_example_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: loaded_transformers: list[tuple[object, str]] = [] @@ -77,6 +81,7 @@ def from_pretrained(cls, *args, **kwargs): assert loaded_transformers == [(tmp_path / "transformer", "cpu"), (tmp_path / "refiner", "cpu")] 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 ) From 3e7686553263f9bb38bdfd9ee09d0cbdee3d61d2 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 15:19:42 +0000 Subject: [PATCH 4/9] fix(lingbot-video): align runtime edge cases with upstream Fix distributed base-worker restoration after refiner handoff, official zero-null refiner CFG defaults, T2I refiner rejection, prompt-duration precedence, and single-frame output suffix handling. Restore grouped-MoE outputs with FP32 route reduction, disable CFG for guidance scales at or below one, and reject packed sequences on attention backends that do not implement packed semantics. Keep validation tools aligned and add regression coverage for the corrected lifecycle and contracts. Verification: - .venv/bin/python -m pytest (119 passed) - H100 grouped_mm versus sorted eager exact parity (0 mismatches) - ruff check on all changed files - ruff format --check on all changed files - git diff --check - Dense 1.3B 4-GPU example, 121 frames at 832x480 - MoE 30B plus refiner 4-GPU example, 121 frames at 1920x1088 --- .../lingbot_video/lingbot_video_dense_1_3b.py | 22 +++--- .../lingbot_video/lingbot_video_moe_30b.py | 48 ++++++++---- telefuser/ops/attention/attention_impl.py | 2 + telefuser/ops/moe.py | 10 ++- .../pipelines/lingbot_video/denoising.py | 6 +- telefuser/pipelines/lingbot_video/pipeline.py | 26 ++++++- tests/unit/models/test_lingbot_video_dit.py | 23 ++++++ tests/unit/models/test_lingbot_video_moe.py | 2 +- .../pipelines/lingbot_video/test_pipeline.py | 76 +++++++++++++++++++ .../lingbot_video/test_service_contract.py | 52 ++++++++++--- tools/validation/benchmark_lingbot_video.py | 7 +- .../run_lingbot_video_native_parallel.py | 4 +- ...te_lingbot_video_refiner_output_handoff.py | 5 +- 13 files changed, 226 insertions(+), 57 deletions(-) diff --git a/examples/lingbot_video/lingbot_video_dense_1_3b.py b/examples/lingbot_video/lingbot_video_dense_1_3b.py index 135bd18..147f7d7 100644 --- a/examples/lingbot_video/lingbot_video_dense_1_3b.py +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -231,7 +231,7 @@ def run( seed: int = PPL_CONFIG["seed"], resolution: str = PPL_CONFIG["resolution"], aspect_ratio: str = PPL_CONFIG["aspect_ratio"], - target_video_length: int = PPL_CONFIG["target_video_length"], + target_video_length: float | None = None, task: str = "t2v", first_image_path: str = "", ) -> torch.Tensor: @@ -241,11 +241,12 @@ def run( if task == "i2v" and not first_image_path: raise ValueError("LingBot-Video i2v requires first_image_path") - caption, _ = parse_lingbot_video_prompt(json.loads(prompt)) + caption, prompt_duration = parse_lingbot_video_prompt(json.loads(prompt)) width, height = _resolve_video_size(aspect_ratio, resolution) - num_frames = ( - 1 if task == "t2i" else num_frames_from_duration(float(target_video_length), fps=int(PPL_CONFIG["fps"])) - ) + 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( @@ -263,11 +264,14 @@ def run( 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) - path.parent.mkdir(parents=True, exist_ok=True) 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)} @@ -278,7 +282,7 @@ def run_with_file( negative_prompt: str | None = None, seed: int = PPL_CONFIG["seed"], output_path: str = "output.mp4", - target_video_length: int = PPL_CONFIG["target_video_length"], + target_video_length: float | None = None, resolution: str = PPL_CONFIG["resolution"], aspect_ratio: str = PPL_CONFIG["aspect_ratio"], task: str = "t2v", @@ -308,7 +312,7 @@ def run_with_file( @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=PPL_CONFIG["target_video_length"], type=int) +@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") @@ -320,7 +324,7 @@ def main( seed: int, resolution: str, aspect_ratio: str, - target_video_length: int, + target_video_length: float | None, task: str, first_image_path: str, output_path: str, diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py index bcd0cfd..89352ca 100644 --- a/examples/lingbot_video/lingbot_video_moe_30b.py +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -70,6 +70,7 @@ "refiner_tail_steps": 2, "refiner_parallelism": 0, "refiner_batch_cfg": False, + "refiner_null_cond_clone_zero": True, } @@ -82,10 +83,10 @@ def _build_contract() -> dict[str, Any]: "seed": {"default": PPL_CONFIG["seed"]}, "resolution": {"default": PPL_CONFIG["resolution"]}, "aspect_ratio": {"default": PPL_CONFIG["aspect_ratio"]}, - "refine": {"type": "boolean", "default": PPL_CONFIG["enable_refiner"]}, } 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"], @@ -352,10 +353,10 @@ def run( seed: int = PPL_CONFIG["seed"], resolution: str = PPL_CONFIG["resolution"], aspect_ratio: str = PPL_CONFIG["aspect_ratio"], - target_video_length: int = PPL_CONFIG["target_video_length"], + target_video_length: float | None = None, task: str = "t2v", first_image_path: str = "", - refine: bool = PPL_CONFIG["enable_refiner"], + refine: bool | None = None, model_root: str = PPL_CONFIG["model_root"], ) -> torch.Tensor: """Run MoE T2I, T2V, or I2V with an optional refiner stage.""" @@ -363,12 +364,17 @@ def run( 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, _ = parse_lingbot_video_prompt(json.loads(prompt)) + caption, prompt_duration = parse_lingbot_video_prompt(json.loads(prompt)) width, height = _resolve_video_size(aspect_ratio, resolution) - num_frames = ( - 1 if task == "t2i" else num_frames_from_duration(float(target_video_length), fps=int(PPL_CONFIG["fps"])) - ) + 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( @@ -389,11 +395,18 @@ def run( if generation.prompt_conditions.has_visual_condition: positive, positive_mask = pipeline.text_stage.encode(caption) - negative, negative_mask = pipeline.text_stage.encode(resolved_negative_prompt) else: positive = generation.prompt_conditions.positive_prompt_embeds - negative = generation.prompt_conditions.negative_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 pipeline.release_gpu_resources() @@ -442,11 +455,14 @@ def run( 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) - path.parent.mkdir(parents=True, exist_ok=True) 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)} @@ -457,12 +473,12 @@ def run_with_file( negative_prompt: str | None = None, seed: int = PPL_CONFIG["seed"], output_path: str = "output.mp4", - target_video_length: int = PPL_CONFIG["target_video_length"], + 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 = PPL_CONFIG["enable_refiner"], + refine: bool | None = None, model_root: str = PPL_CONFIG["model_root"], **_: object, ) -> dict[str, str]: @@ -497,14 +513,14 @@ def run_with_file( ) @click.option("--resolution", default=PPL_CONFIG["resolution"]) @click.option("--aspect_ratio", default=PPL_CONFIG["aspect_ratio"]) -@click.option("--target_video_length", default=PPL_CONFIG["target_video_length"], type=int) +@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("--refine/--no-refine", default=PPL_CONFIG["enable_refiner"]) +@click.option("--refine/--no-refine", default=None) @click.option("--output_path", default="output.mp4") def main( gpu_num: int, @@ -515,12 +531,12 @@ def main( expert_backend: str, resolution: str, aspect_ratio: str, - target_video_length: int, + target_video_length: float | None, task: str, first_image_path: str, refiner_gpu_num: int, refiner_batch_cfg: bool, - refine: bool, + refine: bool | None, output_path: str, ) -> None: """Generate with the LingBot-Video MoE 30B checkpoint and refiner.""" diff --git a/telefuser/ops/attention/attention_impl.py b/telefuser/ops/attention/attention_impl.py index 5ad7be9..d2bc4d2 100755 --- a/telefuser/ops/attention/attention_impl.py +++ b/telefuser/ops/attention/attention_impl.py @@ -157,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): diff --git a/telefuser/ops/moe.py b/telefuser/ops/moe.py index 32c3ee2..b4965dd 100644 --- a/telefuser/ops/moe.py +++ b/telefuser/ops/moe.py @@ -119,16 +119,18 @@ def grouped_expert_forward( ).to(tokens.dtype) expert_output = _unpad_grouped_tokens(expert_output, padded_indices, sorted_tokens.shape[0]) - sorted_weights = flat_weights[sorted_positions].to(expert_output.dtype) - weighted_output = expert_output * sorted_weights.unsqueeze(-1) routed = torch.zeros( num_tokens * top_k, hidden_size, dtype=expert_output.dtype, device=expert_output.device, ) - routed[sorted_positions] = weighted_output - return routed.reshape(num_tokens, top_k, hidden_size).sum(dim=1).to(tokens.dtype) + 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 _pad_grouped_tokens( diff --git a/telefuser/pipelines/lingbot_video/denoising.py b/telefuser/pipelines/lingbot_video/denoising.py index 3c99ade..cc3f863 100644 --- a/telefuser/pipelines/lingbot_video/denoising.py +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -201,12 +201,12 @@ def _predict_noise_with_cfg( 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: + 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 not 1") + raise ValueError("negative_prompt_embeds is required when guidance_scale is greater than 1") if self.batch_cfg: combined_latents = torch.cat((latents, latents)) combined_timestep = torch.cat((model_timestep, model_timestep)) @@ -236,7 +236,7 @@ def denoise( self, latent: torch.Tensor, positive_prompt_embeds: torch.Tensor, - negative_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, diff --git a/telefuser/pipelines/lingbot_video/pipeline.py b/telefuser/pipelines/lingbot_video/pipeline.py index 03a59f1..0b359bb 100644 --- a/telefuser/pipelines/lingbot_video/pipeline.py +++ b/telefuser/pipelines/lingbot_video/pipeline.py @@ -58,9 +58,9 @@ class LingBotVideoPromptConditions: """Text conditions produced for one source-order CFG generation.""" positive_prompt_embeds: torch.Tensor - negative_prompt_embeds: torch.Tensor + negative_prompt_embeds: torch.Tensor | None positive_attention_mask: torch.Tensor - negative_attention_mask: torch.Tensor + negative_attention_mask: torch.Tensor | None has_visual_condition: bool @@ -84,12 +84,15 @@ def __init__(self, device: str = "cuda", torch_dtype: torch.dtype = torch.bfloat 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( @@ -101,8 +104,10 @@ def init(self, module_manager: ModuleManager, config: LingBotVideoPipelineConfig 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]: @@ -115,6 +120,7 @@ def _get_stages(self) -> list[object]: 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): @@ -148,6 +154,16 @@ def release_gpu_resources(self) -> None: 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, @@ -169,6 +185,7 @@ def generate( ) -> 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: @@ -186,7 +203,10 @@ def generate( 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) - negative, negative_mask = self.text_stage.encode(negative_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, diff --git a/tests/unit/models/test_lingbot_video_dit.py b/tests/unit/models/test_lingbot_video_dit.py index bbe16e3..255edfb 100644 --- a/tests/unit/models/test_lingbot_video_dit.py +++ b/tests/unit/models/test_lingbot_video_dit.py @@ -95,3 +95,26 @@ def test_transformer_packed_batch_matches_independent_samples_with_different_tex ) 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 index 2c0fa66..bffa8cc 100644 --- a/tests/unit/models/test_lingbot_video_moe.py +++ b/tests/unit/models/test_lingbot_video_moe.py @@ -128,4 +128,4 @@ def test_grouped_mm_backend_matches_sorted_backend_on_cuda() -> None: experts.set_execution_backend("grouped_mm") grouped_output = experts(tokens, indices, weights) - torch.testing.assert_close(grouped_output, sorted_output, rtol=2e-2, atol=2e-2) + assert torch.equal(grouped_output, sorted_output) diff --git a/tests/unit/pipelines/lingbot_video/test_pipeline.py b/tests/unit/pipelines/lingbot_video/test_pipeline.py index 6456515..e60800f 100644 --- a/tests/unit/pipelines/lingbot_video/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -4,6 +4,7 @@ import torch.nn as nn from telefuser.core.config import ModelRuntimeConfig +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 @@ -141,6 +142,26 @@ def test_denoising_stage_returns_fp32_scheduler_prediction() -> None: 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__() @@ -253,6 +274,29 @@ def encode(self, caption: str, images: object | None = None) -> tuple[torch.Tens 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 @@ -320,3 +364,35 @@ def test_pipeline_releases_stages_before_a_separate_refiner_is_loaded() -> None: 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_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py index bb89bcb..02f6317 100644 --- a/tests/unit/pipelines/lingbot_video/test_service_contract.py +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -66,6 +66,7 @@ def test_only_moe_contract_exposes_refiner() -> None: 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 @@ -115,6 +116,17 @@ def generate(self, request, **_: object) -> SimpleNamespace: 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] = [] @@ -130,12 +142,12 @@ class FakePipeline: device = "cpu" def generate(self, request, **_: object) -> SimpleNamespace: - assert request.num_frames == 121 + 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": 5}), + json.dumps({"caption": {"scene": "test"}, "duration": 2}), task="t2v", output_path=str(tmp_path / "dense.mp4"), ) @@ -160,12 +172,12 @@ def generate(self, request, **_: object) -> SimpleNamespace: del request conditions = SimpleNamespace( has_visual_condition=False, - positive_prompt_embeds=torch.zeros(1, 1, 2), - negative_prompt_embeds=torch.zeros(1, 1, 2), - positive_attention_mask=torch.ones(1, 1, dtype=torch.bool), - negative_attention_mask=torch.ones(1, 1, dtype=torch.bool), + 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, 1, 2, 2), prompt_conditions=conditions) + return SimpleNamespace(output=torch.zeros(1, 3, 5, 2, 2), prompt_conditions=conditions) def release_gpu_resources(self) -> None: self.released = True @@ -174,10 +186,14 @@ 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, *args, **kwargs) -> torch.Tensor: - del args, kwargs + 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: @@ -192,14 +208,28 @@ def close(self) -> None: frames = module.run( pipeline, json.dumps({"caption": {"scene": "test"}, "duration": 5}), - task="t2i", + task="t2v", refine=True, ) - assert frames.shape == (1, 3, 1, 2, 2) + 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_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"]) diff --git a/tools/validation/benchmark_lingbot_video.py b/tools/validation/benchmark_lingbot_video.py index b43ae8b..79dd6c1 100644 --- a/tools/validation/benchmark_lingbot_video.py +++ b/tools/validation/benchmark_lingbot_video.py @@ -189,15 +189,12 @@ def generate_once() -> torch.Tensor: positive, positive_mask = _measure( active_samples, "refiner_text_encode", lambda: pipeline.text_stage.encode(caption) ) - negative, negative_mask = _measure( - active_samples, "refiner_text_encode", lambda: pipeline.text_stage.encode(negative_caption) - ) else: positive = generation.prompt_conditions.positive_prompt_embeds - negative = generation.prompt_conditions.negative_prompt_embeds positive_mask = generation.prompt_conditions.positive_attention_mask - negative_mask = generation.prompt_conditions.negative_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( diff --git a/tools/validation/run_lingbot_video_native_parallel.py b/tools/validation/run_lingbot_video_native_parallel.py index d9a4734..31eb6c4 100644 --- a/tools/validation/run_lingbot_video_native_parallel.py +++ b/tools/validation/run_lingbot_video_native_parallel.py @@ -119,9 +119,9 @@ def main() -> None: output = refiner.refine( lowres_video, generation.prompt_conditions.positive_prompt_embeds, - generation.prompt_conditions.negative_prompt_embeds, + torch.zeros_like(generation.prompt_conditions.positive_prompt_embeds), generation.prompt_conditions.positive_attention_mask, - generation.prompt_conditions.negative_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, diff --git a/tools/validation/validate_lingbot_video_refiner_output_handoff.py b/tools/validation/validate_lingbot_video_refiner_output_handoff.py index cfe385b..9242eed 100644 --- a/tools/validation/validate_lingbot_video_refiner_output_handoff.py +++ b/tools/validation/validate_lingbot_video_refiner_output_handoff.py @@ -225,12 +225,11 @@ def main() -> 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) - negative, negative_mask = pipeline.text_stage.encode(negative_caption) else: positive = generation.prompt_conditions.positive_prompt_embeds - negative = generation.prompt_conditions.negative_prompt_embeds positive_mask = generation.prompt_conditions.positive_attention_mask - negative_mask = generation.prompt_conditions.negative_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) From b05af23be575d0b424757eda2f78896caed4e5f9 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 16:19:21 +0000 Subject: [PATCH 5/9] feat(lingbot-video): add hybrid CFG and FP8 runtime modes - add CFG2 x SP2 execution for Dense, MoE, and refiner stages with batch-CFG conflict validation - add native per-channel W8A8 expert execution while retaining grouped BF16 as the performance default - keep base and refiner DiT workers co-resident during refinement and release both before high-resolution VAE decode - expose CLI/config controls, document memory and backend tradeoffs, and cover lifecycle and numerical behavior Verification: - ruff format --check on all changed Python files - ruff check on all changed Python files - git diff --check - pytest LingBot-Video unit/model/tool suite: 89 passed - 4x H100 Dense CFG2 x SP2, 832x480, 25 frames - 4x H100 MoE FP8 CFG2 x SP2 with co-resident refiner, 1920x1088, 25 frames - 4x H100 SP4 versus CFG2 x SP2 Dense and MoE grouped_mm benchmarks --- examples/lingbot_video/README.md | 37 +++++-- .../lingbot_video/lingbot_video_dense_1_3b.py | 21 +++- .../lingbot_video/lingbot_video_moe_30b.py | 78 +++++++++++++-- telefuser/models/lingbot_video_moe.py | 42 +++++++- telefuser/ops/moe.py | 99 +++++++++++++++++++ .../pipelines/lingbot_video/denoising.py | 25 ++++- telefuser/pipelines/lingbot_video/refiner.py | 4 + tests/unit/models/test_lingbot_video_moe.py | 51 ++++++++++ .../pipelines/lingbot_video/test_pipeline.py | 31 ++++++ .../lingbot_video/test_refiner_stage.py | 38 +++++++ .../lingbot_video/test_service_contract.py | 69 +++++++++++++ 11 files changed, 467 insertions(+), 28 deletions(-) diff --git a/examples/lingbot_video/README.md b/examples/lingbot_video/README.md index 04504ac..ea89f6e 100644 --- a/examples/lingbot_video/README.md +++ b/examples/lingbot_video/README.md @@ -27,13 +27,31 @@ python examples/lingbot_video/lingbot_video_moe_30b.py \ --prompt "$(cat /path/to/caption.json)" --output_path result.mp4 ``` -The command releases base-stage GPU weights before loading the separate refiner. -Use `--no-refine` for MoE base-only generation. `expert_backend=auto` keeps +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. -Use `--refiner_gpu_num` and `--refiner_batch_cfg` to override the corresponding -MoE `PPL_CONFIG` defaults for a direct CLI run. +`--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: @@ -43,8 +61,9 @@ 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 SP4/FSDP stage explicitly; -otherwise it inherits service parallelism. +`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 @@ -58,6 +77,6 @@ telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 - 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. The refiner starts as a separate worker group -after the base workers exit. Four-GPU 1920x1088 refinement uses sequential CFG -by default to stay within H100 80 GB memory. +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/lingbot_video_dense_1_3b.py b/examples/lingbot_video/lingbot_video_dense_1_3b.py index 147f7d7..48041fe 100644 --- a/examples/lingbot_video/lingbot_video_dense_1_3b.py +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -52,6 +52,7 @@ "resolution": "480p", "aspect_ratio": "16:9", "seed": 42, + "cfg_parallel_degree": 1, } @@ -113,7 +114,9 @@ def build_pipeline( parallel_config = parallel_config or ParallelConfig() if batch_cfg is None: - batch_cfg = parallel_config.world_size > 1 + 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") device, device_id = _resolve_runtime_device(device, parallel_config) @@ -183,15 +186,23 @@ def build_pipeline( return pipeline -def get_pipeline(parallelism: int = 1, model_root: str = PPL_CONFIG["model_root"]) -> object: +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)), - sp_ulysses_degree=parallelism, + cfg_degree=cfg_parallel_degree, + sp_ulysses_degree=parallelism // cfg_parallel_degree, enable_fsdp=parallelism > 1, ), ) @@ -306,6 +317,7 @@ def run_with_file( @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") @@ -318,6 +330,7 @@ def run_with_file( @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, @@ -330,7 +343,7 @@ def main( output_path: str, ) -> None: """Generate with the LingBot-Video Dense 1.3B checkpoint.""" - pipeline = get_pipeline(gpu_num, model_root) + pipeline = get_pipeline(gpu_num, model_root, cfg_parallel_degree) try: result = run_with_file( pipeline, diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py index 89352ca..ff807d0 100644 --- a/examples/lingbot_video/lingbot_video_moe_30b.py +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -71,6 +71,9 @@ "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, } @@ -134,7 +137,9 @@ def build_pipeline( parallel_config = parallel_config or ParallelConfig() if batch_cfg is None: - batch_cfg = parallel_config.world_size > 1 + 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") device, device_id = _resolve_runtime_device(device, parallel_config) @@ -154,11 +159,15 @@ def build_pipeline( if expert_backend == "auto": expert_backend = "grouped_mm" if parallel_config.world_size > 1 else "sorted" - if expert_backend not in {"grouped_mm", "sorted"}: - raise ValueError("LingBot MoE expert_backend must be 'auto', 'grouped_mm', or '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") transformer = load_lingbot_video_moe_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype) + 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 = ModuleManager(torch_dtype=torch_dtype, device="cpu") @@ -222,6 +231,7 @@ def build_refiner( 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: @@ -232,6 +242,8 @@ def build_refiner( 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") device, device_id = _resolve_runtime_device(device, parallel_config) @@ -248,9 +260,20 @@ def build_refiner( 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'") + transformer = load_lingbot_video_moe_transformer(root / "refiner", device="cpu", torch_dtype=torch_dtype) + 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 = ModuleManager(torch_dtype=torch_dtype, device="cpu") module_manager.add_module( - load_lingbot_video_moe_transformer(root / "refiner", device="cpu", torch_dtype=torch_dtype), + transformer, name="transformer", path=str(root / "refiner"), ) @@ -287,17 +310,24 @@ def get_pipeline( 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)), - sp_ulysses_degree=parallelism, + cfg_degree=cfg_parallel_degree, + sp_ulysses_degree=parallelism // cfg_parallel_degree, enable_fsdp=parallelism > 1, ), expert_backend=expert_backend, @@ -307,15 +337,29 @@ def get_pipeline( ) 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)), - sp_ulysses_degree=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 @@ -409,12 +453,14 @@ def run( negative = generation.prompt_conditions.negative_prompt_embeds negative_mask = generation.prompt_conditions.negative_attention_mask - pipeline.release_gpu_resources() + 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( @@ -447,6 +493,7 @@ def run( 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() @@ -501,6 +548,7 @@ def run_with_file( @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") @@ -508,8 +556,8 @@ def run_with_file( @click.option( "--expert_backend", default=PPL_CONFIG["expert_backend"], - type=click.Choice(["auto", "grouped_mm", "sorted"]), - help="MoE expert backend; auto selects grouped_mm for four GPUs", + 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"]) @@ -520,10 +568,17 @@ def run_with_file( "--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, @@ -536,6 +591,8 @@ def main( 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: @@ -545,6 +602,9 @@ def main( 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: diff --git a/telefuser/models/lingbot_video_moe.py b/telefuser/models/lingbot_video_moe.py index 5aa6bac..759c1a0 100644 --- a/telefuser/models/lingbot_video_moe.py +++ b/telefuser/models/lingbot_video_moe.py @@ -12,7 +12,7 @@ import torch.nn.functional as F from torch import nn -from telefuser.ops.moe import grouped_expert_forward, route_topk +from telefuser.ops.moe import fp8_expert_forward, grouped_expert_forward, quantize_expert_weight_fp8, route_topk from .lingbot_video_dit import LingBotVideoBlock, LingBotVideoMLP, LingBotVideoTransformer3DModel @@ -60,15 +60,45 @@ def __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 {"grouped_mm", "sorted", "where"}: - raise ValueError("execution backend must be 'grouped_mm', 'sorted', or 'where'") + 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": @@ -249,3 +279,9 @@ def set_expert_execution_backend(self, backend: str) -> None: for module in self.modules(): if isinstance(module, LingBotVideoGroupedExperts): module.set_execution_backend(backend) + + 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_() diff --git a/telefuser/ops/moe.py b/telefuser/ops/moe.py index b4965dd..ce69555 100644 --- a/telefuser/ops/moe.py +++ b/telefuser/ops/moe.py @@ -133,6 +133,105 @@ def grouped_expert_forward( ) +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, diff --git a/telefuser/pipelines/lingbot_video/denoising.py b/telefuser/pipelines/lingbot_video/denoising.py index cc3f863..68d97c9 100644 --- a/telefuser/pipelines/lingbot_video/denoising.py +++ b/telefuser/pipelines/lingbot_video/denoising.py @@ -14,8 +14,14 @@ 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_ulysses_group +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 @@ -142,6 +148,11 @@ def parallel_models(self) -> None: 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}") @@ -149,7 +160,7 @@ def parallel_models(self) -> None: 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 != self.torch_dtype + 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") @@ -207,7 +218,15 @@ def _predict_noise_with_cfg( ).float() if negative_prompt_embeds is None: raise ValueError("negative_prompt_embeds is required when guidance_scale is greater than 1") - if self.batch_cfg: + 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( diff --git a/telefuser/pipelines/lingbot_video/refiner.py b/telefuser/pipelines/lingbot_video/refiner.py index 537cc9a..88f1bc6 100644 --- a/telefuser/pipelines/lingbot_video/refiner.py +++ b/telefuser/pipelines/lingbot_video/refiner.py @@ -4,6 +4,7 @@ from __future__ import annotations +from collections.abc import Callable from pathlib import Path import numpy as np @@ -251,6 +252,7 @@ def refine( 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) @@ -309,4 +311,6 @@ def refine( # 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/tests/unit/models/test_lingbot_video_moe.py b/tests/unit/models/test_lingbot_video_moe.py index bffa8cc..3ed6d8f 100644 --- a/tests/unit/models/test_lingbot_video_moe.py +++ b/tests/unit/models/test_lingbot_video_moe.py @@ -109,6 +109,34 @@ def test_grouped_mm_backend_requires_cuda() -> None: ) +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"): @@ -129,3 +157,26 @@ def test_grouped_mm_backend_matches_sorted_backend_on_cuda() -> None: 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_pipeline.py b/tests/unit/pipelines/lingbot_video/test_pipeline.py index e60800f..19e7325 100644 --- a/tests/unit/pipelines/lingbot_video/test_pipeline.py +++ b/tests/unit/pipelines/lingbot_video/test_pipeline.py @@ -4,6 +4,7 @@ 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 @@ -202,6 +203,36 @@ def test_denoising_stage_batched_cfg_matches_two_source_order_forwards() -> None 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) diff --git a/tests/unit/pipelines/lingbot_video/test_refiner_stage.py b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py index 57aaa97..fa06c52 100644 --- a/tests/unit/pipelines/lingbot_video/test_refiner_stage.py +++ b/tests/unit/pipelines/lingbot_video/test_refiner_stage.py @@ -171,3 +171,41 @@ def test_refiner_releases_vae_before_denoising_and_worker_before_decode() -> Non 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_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py index 02f6317..9cca21b 100644 --- a/tests/unit/pipelines/lingbot_video/test_service_contract.py +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -92,6 +92,28 @@ def capture_build(model_root: str, **kwargs: object) -> object: 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") @@ -220,6 +242,53 @@ def close(self) -> None: 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") From 2b7e69d1c64047a489764fb5e66875084fe76b0a Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 16:48:12 +0000 Subject: [PATCH 6/9] docs(lingbot-video): consolidate usage documentation - keep examples/lingbot_video/README.md as the single user-facing guide\n- remove duplicate English and Chinese MkDocs pages and stale navigation entries\n- update model-zoo defaults, CFG/SP, FP8, and refiner lifecycle references\n- update supported pipeline matrices and service contract tests\n\nVerification: git diff --check --- README.md | 3 +- README_zh.md | 2 + docs/en/index.md | 3 +- docs/en/lingbot_video.md | 288 ------------------ docs/en/service.md | 2 + docs/zh/index.md | 3 +- docs/zh/lingbot_video.md | 222 -------------- docs/zh/service.md | 2 + examples/lingbot_video/README.md | 2 + .../lingbot_video/lingbot_video_dense_1_3b.py | 4 +- .../lingbot_video/lingbot_video_moe_30b.py | 4 +- .../lingbot_video/test_service_contract.py | 17 ++ 12 files changed, 37 insertions(+), 515 deletions(-) delete mode 100644 docs/en/lingbot_video.md delete mode 100644 docs/zh/lingbot_video.md diff --git a/README.md b/README.md index 4af7fa9..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,7 +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) | -| `LingBot-Video` | T2I, T2V, TI2V, MoE refiner | Precision-first Dense/MoE generation with a single-process base-to-refiner path; see [docs/en/lingbot_video.md](docs/en/lingbot_video.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/index.md b/docs/en/index.md index 88b01ec..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 diff --git a/docs/en/lingbot_video.md b/docs/en/lingbot_video.md deleted file mode 100644 index 863e8dc..0000000 --- a/docs/en/lingbot_video.md +++ /dev/null @@ -1,288 +0,0 @@ -# LingBot-Video - -LingBot-Video supports Dense and MoE base DiTs for T2I, T2V, and TI2V. The -MoE checkpoint also includes a separate low-noise refiner. The integration is -precision-first: use the upstream reference capture artifacts before enabling -backend or distributed optimizations. - -## Checkpoints - -The Dense loader consumes a Diffusers `transformer/` directory directly: - -```python -from telefuser.pipelines.lingbot_video import load_lingbot_video_dense_transformer - -transformer = load_lingbot_video_dense_transformer( - "/path/to/lingbot-video-dense-1.3b/transformer" -) -``` - -The MoE and refiner directories are sharded. Use -`load_lingbot_video_moe_transformer` for either `transformer/` or `refiner/`. -The default sorted eager expert path preserves upstream route ordering and keeps -a `where`-based diagnostic fallback. It remains the validated single-GPU BF16 -path. Four-GPU MoE execution defaults to native grouped GEMM when the installed -CUDA PyTorch build exposes `torch._grouped_mm`; the example accepts -`expert_backend=sorted` as an explicit fallback. For -`variant="moe"`, the runtime defaults to stage CPU offload so the base DiT, -text encoder, VAE, and separately loaded refiner do not need to reside on one -GPU. Set `cpu_offload=False` only when GPU capacity is known to be sufficient. - -## Prompt preparation - -The generation pipeline intentionally consumes the structured JSON caption, while -prompt rewriting remains an optional, separately deployable workflow. Preserve the -official two-stage contract: EXPAND uses the base VLM with its LoRA disabled, and -MAP uses the same VLM with the LingBot rewriter LoRA enabled. For TI2V, provide -the identical first frame to both the rewriter and TeleFuser generation. - -```bash -REWRITER_BASE_MODEL=/path/to/Qwen3.6-27B \ -REWRITER_ADAPTER=/path/to/lingbot-video-rewriter-lora \ -python work_dirs/lingbot-video-master/rewriter/inference.py \ - --mode t2v --prompt "" --duration 5 --output prompt.json -``` - -Pass the resulting `prompt.json` to `--caption-json`, or serialize its -`caption` object as the service `prompt`. The rewriter is deliberately not loaded -inside a DiT service process because the two models have separate deployment and -capacity requirements. - -Unless explicitly overridden, the pipeline, CLI, and service use the checkpoint's -structured negative CFG caption. T2I uses the source still-image variant, while -T2V and TI2V use the source video variant, including its temporal-stability -constraints. Do not replace an omitted negative prompt with an empty string when -reproducing an upstream sample: it changes the Qwen3-VL negative condition and -can materially alter color and image quality. - -## Runtime composition - -`LingBotVideoPipeline` composes independently loaded stages: - -- `LingBotVideoTextEncodingStage` encodes the structured JSON caption with - Qwen3-VL. -- `LingBotVideoDenoisingStage` runs source-order, two-forward CFG. -- `LingBotVideoVAEEncodeStage` and `LingBotVideoVAEDecodeStage` apply the - checkpoint VAE's latent mean/std normalization. -- `FlowUniPCMultistepScheduler` owns the sigma/timestep sequence. - -Load checkpoint components into `ModuleManager`, then call `pipeline.init(module_manager, config)`; `init` constructs all text, DiT, VAE, and scheduler stages from the manager -configuration. Provide a structured JSON caption, not casual unstructured text. - -For the standard checkpoint layout, model loading and stage assembly live directly -in the public model-specific examples. This keeps the runtime wiring visible next -to `PPL_CONFIG`, `CONTRACT`, and the CLI/service entrypoints: - -```python -from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline -from telefuser.pipelines.lingbot_video import LingBotVideoRequest - -pipeline = build_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) -frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) -``` - -Direct API and CLI heights and widths must be divisible by 16: the Wan VAE -downsamples by eight and the DiT spatially patchifies the resulting latents by two. - -The default `AttentionConfig` uses TeleFuser's SDPA dispatcher and remains the -source-equivalent numerical path. Alternative attention backends are opt-in -through `attention_config=` and require a separate L2 parity report; they are -not enabled by the service or CLI defaults. - -The VAE decode stage returns RGB video in `[0,1]`. Video callers must pass these -float frames directly to Diffusers `export_to_video`, which performs the uint8 -conversion itself. Converting to uint8 before that call applies a second 255 -scale and wraps channel values, producing a negative-like MP4. - -## Four-GPU base inference - -Dense and MoE base DiTs support TeleFuser-native four-GPU inference through -per-block FSDP and Ulysses sequence parallelism. The whole joint -video-plus-text sequence is padded only when necessary, sharded across the four -ranks, then restored in source token order. CFG is evaluated as one batched -forward when positive and negative Qwen embeddings have matching shapes; -otherwise the runtime safely uses the source-order two-forward path. - -Start the API service normally; TeleFuser owns the four worker processes, so do -not wrap this command in `torchrun`: - -```bash -telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 -``` - -This configuration keeps FSDP-enabled DiT weights resident on GPU and therefore -requires `cpu_offload=False`. Checkpoint FP32 modulation parameters are retained -as replicated FSDP ignored states, matching the upstream mixed-precision layout. -The MoE 30B model uses its own example. Its grouped expert weights are sharded -with their containing FSDP blocks rather than expert-parallelized: - -```bash -telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 -``` - -On PyTorch 2.11 with CUDA 13, the grouped-GEMM 832x480, 121-frame, 40-step MoE -base run uses one worker-owned scheduler loop and one intra-op CPU thread per -rank, matching `torchrun`. On four H100 GPUs, steady sampling is approximately -2.3 seconds per step and 93 seconds for 40 steps, compared with 92 seconds for -the source runner. The first TeleFuser request also pays about 15 seconds for -FSDP materialization. A cold CLI invocation took 231.9 seconds because the -parent load and spawn-time 30B model transfer are not included in steady -service latency; persistent service requests amortize that setup. The -separately loaded refiner also has a validated four-GPU FSDP/Ulysses path. -Expert parallelism and FP8 remain disabled pending their own parity and -throughput evidence. - -## TI2V - -TI2V has two independent first-frame condition paths: - -1. The image is supplied to Qwen3-VL as visual input with the caption. -2. The image is VAE-encoded to a clean temporal-prefix latent. It is written - before every denoising step and once after the final scheduler step. - -The current pipeline accepts a raw RGB condition image in the range [0,255] with shape -`[B, 3, H, W]` (or `[B, 3, F, H, W]`, using frame zero). It resizes and center-crops the image before separately passing it to Qwen3-VL and the VAE. - -## Service - -The Dense and MoE examples each expose `PPL_CONFIG`, `CONTRACT`, `get_pipeline`, -`run`, and `run_with_file`. Both serve `t2i`, `t2v`, and `i2v` and require a -structured JSON string as `prompt`: - -Set each example's `PPL_CONFIG["model_root"]` before starting the service. Runtime options are read from `PPL_CONFIG`; direct CLI runs can override the exposed command-line options. - -```bash -telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --port 8000 -``` - -Set `PPL_CONFIG["model_root"]` in `lingbot_video_moe_30b.py` before serving MoE. -Its contract exposes `refine`, enabled by default. The service releases base -weights before loading the refiner. `PPL_CONFIG["refiner_parallelism"] = 4` -selects an independent four-worker refiner; when omitted, it inherits service -parallelism. - -```bash -telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 -``` -Requested service resolutions are rounded up to the LingBot VAE-and-DiT -sixteen-pixel spatial grid; for example, `480p` at `16:9` uses the validated LingBot 832x480 landscape preset. -Pass `negative_prompt` only to override the source-compatible default; an -explicit empty string remains a supported override. - -## Refiner - -`LingBotVideoRefinerStage` takes the base RGB output in memory, VAE-encodes it, -mixes it with noise at `t_thresh`, and samples the low-noise sigma tail. It can -also preserve a clean TI2V frame zero. Base and refiner are separate runtime -stages. Call `base_pipeline.release_gpu_resources()` before loading the refiner -when they share a GPU. - -On four GPUs, the refiner uses per-block FSDP and four-way Ulysses SP. Its CFG -defaults to sequential positive/negative forwards because batched CFG exceeds -four H100 80 GB cards at 1920x1088. Set `PPL_CONFIG["refiner_batch_cfg"] = True` -only after validating additional memory capacity. The runtime VAE-encodes the -RGB handoff, releases the VAE, runs the distributed refiner, closes all refiner -workers, and only then reloads the VAE decoder. - -The validated MoE run generated a 832x480, 121-frame base and refined it to -1920x1088 at 24 FPS. Base generation took 372.9 seconds and the eight-step -refiner stage took 886.8 seconds, excluding checkpoint deserialization. The -output contains 121 frames and is 5.0417 seconds long. -The included CLI implements this lifecycle for MoE checkpoints: - -```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 -``` - -With `--task i2v --first_image_path first_frame.png`, the CLI also applies the -upstream TI2V frame-zero geometry to the refiner condition. - -For an in-memory base-to-refiner handoff, call `prepare_refiner_video(...)` before -`LingBotVideoRefinerStage.refine(...)`. It matches upstream training-aligned frame -selection and bicubic resize without an MP4 write/read round trip; pass the base -output FPS explicitly. Validate this path against the corresponding source MP4 baseline. -The MP4 compatibility test uses the upstream Diffusers writer and compares the upstream loader (through a PyAV-backed decord adapter when decord is unavailable) with this loader tensor-for-tensor. - - -## Validation - -Capture upstream artifacts before comparing a TeleFuser run: - -```bash -python tools/validation/capture_lingbot_video_reference.py --dry-run -python tools/validation/capture_lingbot_video_reference.py --all-cases --mode t2i --mode t2v --mode ti2v --trace sampled -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-dense-1.3b --variant dense --output dense-load-report.json -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --output moe-load-report.json -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant refiner --output refiner-load-report.json -python tools/validation/compare_lingbot_video_parity.py REFERENCE CANDIDATE -python tools/validation/replay_lingbot_video_dense_reference.py --reference-dir work_dirs/lingbot_video_reference/t2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --validate-ti2v-vae --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --reference-root work_dirs/lingbot_video_reference_all_cases --assert-exact --output dense-all-cases-replay.json -PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_moe_parity.py --transformer-dir /path/to/lingbot-video-moe-30b-a3b/transformer --assert-exact -PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_refiner_core_parity.py --model-root /path/to/lingbot-video-moe-30b-a3b --assert-exact -python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.mp4 --height 1088 --width 1920 --assert-exact -python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 -python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 -python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 -python -m torch.distributed.run --standalone --nproc_per_node=4 tools/validation/run_lingbot_video_distributed.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output dense-sp4.mp4 --report dense-sp4.json -python tools/validation/run_lingbot_video_native_parallel.py --variant moe --refine --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --output moe-refiner-sp4.mp4 --report moe-refiner-sp4.json -``` - -`validate_lingbot_video_refiner_handoff.py` checks that the TeleFuser MP4 -compatibility loader is tensor-identical to the source loader. Add -`--in-memory-video` and `--in-memory-fps` to quantify the input difference -introduced by MP4 encoding; this comparison does not replace a final refiner -output-quality evaluation. -Use `--assert-exact` to make the source MP4 compatibility comparison fail on -metadata, tensor shape, dtype, or value drift. It intentionally does not judge -the native in-memory handoff, whose difference from lossy MP4 is expected. -`validate_lingbot_video_refiner_output_handoff.py` generates one MoE base -sample, drives the refiner with both the native RGB tensor and a temporary MP4 -round trip using identical prompt conditions and RNG state, then reports the -final-output L2 difference. This is intentionally a quality comparison, not an -equivalence test: the upstream refiner uses the lossy MP4 round trip, while the -native path removes that intermediate encoding. The report includes decoded-frame -PSNR and local SSIM; `--comparison-output` writes memory output on the left and -MP4-round-trip output on the right for human review. - -The capture tool records scheduler tensors, prompt tensors, selected denoising steps, latent inputs/outputs, the generation seed, RNG state hashes, and decoded frame hashes. Add `--validate-text` to the replay command to compare Qwen3-VL processor inputs and final embeddings; TI2V validation also compares the preprocessed first frame. Add `--validate-ti2v-vae` to compare the sampled clean condition latent; use `--seed` only with pre-existing captures that lack seed metadata. -Use `--reference-root` for an all-case Dense DiT/VAE replay. It retains one -loaded Dense transformer and VAE across captured runs, while instantiating a -fresh scheduler for each run so each capture preserves its own sampling setup. -Add `--assert-exact` to make the command fail when any recorded tensor differs -in shape or value, so it can serve as a CI parity gate rather than a report-only -diagnostic. -The checkpoint-inspection tool performs the normal strict load and records the -consumed config fields, checkpoint-key coverage, component/block parameter -counts, dtype/device distribution, retained FP32 parameter count, and model -memory allocation evidence. -The benchmark tool reports one-time setup separately from warmup and measured metrics for checkpoint load, text encoding, each denoising step, VAE, refiner, output encoding, and peak GPU memory. Use it to establish a baseline before enabling an optimization; record full-resolution measurements separately from smoke runs. -When `--negative-caption` is omitted, the benchmark uses the same source-compatible -T2I or video negative caption as the pipeline, CLI, and service. Pass an explicit -empty string only when intentionally benchmarking that semantic override. -For a base-plus-refiner run, it also records the serial base release and refiner -load phases. Single-GPU benchmarks use the source-equivalent sorted eager MoE -path. Four-GPU benchmarks select grouped GEMM by default; use -`expert_backend=sorted` only when measuring the correctness fallback. -For T2I/T2V, the refiner reuses the exact CFG text conditions from the base -generation and reports this as `refiner_prompt_conditions_reused`; TI2V keeps -the source-compatible text-only refiner encoding path. -The refiner core CLI injects identical latent, noise, prompt, and frame-zero condition tensors into the upstream and TeleFuser low-noise paths. It offloads the upstream DiT before loading the TeleFuser DiT, so both 30B models do not overlap in GPU memory. -Add `--assert-exact` to the MoE or refiner core validator to enforce the -zero-drift numerical-oracle gate instead of only writing metrics. - -## Requirements and limitations - -The numerical-oracle path requires CUDA, PyTorch, Diffusers, Transformers, and the -checkpoint components `transformer/`, `text_encoder/`, `processor/`, `vae/`, and -`scheduler/`. Dense runs source-equivalently on one GPU. MoE keeps sorted eager -for single-GPU correctness and uses native grouped GEMM for four-GPU throughput; -the latter requires a CUDA PyTorch build exposing `torch._grouped_mm`. Dense and -MoE base DiTs have validated four-GPU FSDP/Ulysses SP paths. -The MoE refiner has a validated independent four-GPU FSDP/Ulysses SP stage. -External FlashAttention, MoE expert parallelism, and FP8 experts are not enabled. -The base-plus-refiner stage lifecycle remains serial so their 30B weights never -coexist on GPU. 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/index.md b/docs/zh/index.md index 5a051e4..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 交互式世界模型 | ### 视频生成 diff --git a/docs/zh/lingbot_video.md b/docs/zh/lingbot_video.md deleted file mode 100644 index e9da96e..0000000 --- a/docs/zh/lingbot_video.md +++ /dev/null @@ -1,222 +0,0 @@ -# LingBot-Video - -LingBot-Video 支持 Dense 与 MoE base DiT 的 T2I、T2V、TI2V；MoE checkpoint -还包含独立的低噪 refiner。本集成按精度优先实现，在启用优化或分布式后端前应先完成上游参考对比。 - -## Checkpoint - -Dense checkpoint 可直接加载 Diffusers `transformer/` 目录： - -```python -from telefuser.pipelines.lingbot_video import load_lingbot_video_dense_transformer - -transformer = load_lingbot_video_dense_transformer("/path/to/transformer") -``` - -MoE 与 refiner 使用分片 safetensors，调用 -`load_lingbot_video_moe_transformer` 加载 `transformer/` 或 `refiner/`。 -默认的 sorted eager expert 路径保持上游 route 顺序，并保留 `where` 诊断 fallback。 -它仍是经过验证的单卡 BF16 路径。四卡 MoE 默认在 CUDA PyTorch 提供 -`torch._grouped_mm` 时启用原生 grouped GEMM；可通过 -`expert_backend=sorted` 显式回退。 -`variant="moe"` 默认启用 stage CPU offload，避免 base DiT、text encoder、VAE 与 -独立 refiner 同时驻留单张 GPU。仅在显存明确充足时才设置 `cpu_offload=False`。 - -## Prompt 准备 - -生成管线有意只消费结构化 JSON caption；prompt rewriter 是可独立部署的可选流程。 -必须保留上游两阶段语义：EXPAND 使用不挂 LoRA 的 base VLM，MAP 使用同一 base VLM -并启用 LingBot rewriter LoRA。TI2V 必须将同一张首帧同时传给 rewriter 与 TeleFuser。 - -```bash -REWRITER_BASE_MODEL=/path/to/Qwen3.6-27B \ -REWRITER_ADAPTER=/path/to/lingbot-video-rewriter-lora \ -python work_dirs/lingbot-video-master/rewriter/inference.py \ - --mode t2v --prompt "" --duration 5 --output prompt.json -``` - -将输出的 `prompt.json` 传给 `--caption-json`，或将其中的 `caption` 对象序列化为 -服务请求的 `prompt`。由于 rewriter 与 DiT 的部署和显存需求独立，TeleFuser 不会在 -同一 DiT 服务进程内加载 rewriter。 - -未显式覆盖时，pipeline、CLI 与 service 都会使用 checkpoint 的结构化负向 CFG -caption。T2I 使用上游的静态图像版本；T2V 与 TI2V 使用包含时序稳定性约束的上游视频 -版本。复现上游样本时，不能以空字符串替代未传入的负向 prompt：这会改变 Qwen3-VL 的 -负向 condition，并可能显著影响颜色与画质。 - -## 运行时组成 - -`LingBotVideoPipeline` 通过标准的 `init(module_manager, config)` 接口统一初始化；checkpoint 组件先加载到 `ModuleManager`，再由 `init` 创建全部 stage： - -- `LingBotVideoTextEncodingStage`：Qwen3-VL structured JSON caption 编码。 -- `LingBotVideoDenoisingStage`：source-order 的 two-forward CFG。 -- `LingBotVideoVAEEncodeStage` / `LingBotVideoVAEDecodeStage`：使用 checkpoint 的 latent mean/std。 -- `FlowUniPCMultistepScheduler`：sigma/timestep 调度。 - -核心 pipeline 应输入结构化 JSON caption，不应将普通自然语言 prompt 直接替代 rewriter 输出。 - -标准 checkpoint 的模型加载与 stage 装配直接放在公开的模型专用 example 中， -与 `PPL_CONFIG`、`CONTRACT` 以及 CLI/service 入口保持在同一层，方便阅读和修改： - -```python -from examples.lingbot_video.lingbot_video_dense_1_3b import build_pipeline -from telefuser.pipelines.lingbot_video import LingBotVideoRequest - -pipeline = build_pipeline("/path/to/lingbot-video-dense-1.3b", num_inference_steps=40) -frames = pipeline(LingBotVideoRequest(caption=structured_caption, height=480, width=832, num_frames=121)) -``` - -直接使用 API 或 CLI 时，高和宽必须能被 16 整除：Wan VAE 会以 8 倍下采样，DiT 会再以 2 倍对 latent 作空间 patchify。 - -默认 `AttentionConfig` 使用 TeleFuser 的 SDPA dispatcher，仍属于 source-equivalent -数值路径。其他 attention backend 必须通过 `attention_config=` 显式启用并单独提供 L2 -parity 报告；service 与 CLI 默认不会启用它们。 - -VAE decode stage 的 RGB video 范围为 `[0,1]`。视频调用方必须将这些浮点帧直接传给 -Diffusers 的 `export_to_video`，由它自行转换为 uint8。若在调用前先转 uint8，就会再次 -乘以 255 并使通道值溢出，生成近似负片的 MP4。 - -## 四卡 base 推理 - -Dense 与 MoE base DiT 均支持 TeleFuser 原生四卡：DiT block 使用 FSDP，联合 video/text -token 流使用 Ulysses sequence parallel。只有 token 总数不能被 rank 数整除时才补齐；输出会按 -上游 token 顺序恢复。正、负 Qwen embedding 形状一致时，CFG 使用一次 batched forward；形状 -不一致时会安全回退到保持上游顺序的两次 forward。 - -按普通方式启动 API 服务即可，TeleFuser 会管理四个 worker，不能再用 `torchrun` 包裹： - -```bash -telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --gpu-num 4 --port 8000 -``` - -该模式要求 DiT 常驻 GPU，因此必须 `cpu_offload=False`。checkpoint 中保留 FP32 的 modulation -参数会作为 FSDP ignored state 在每个 rank 副本化，与上游 mixed-precision 布局一致。MoE 30B -使用独立 example；grouped expert 权重仍随 FSDP block 分片，而不是 expert parallel： - -```bash -telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --gpu-num 4 --port 8000 -``` - -在 PyTorch 2.11/CUDA 13 环境中，已验证的 grouped-GEMM MoE base 为 832x480、 -121 帧、40 steps。四卡路径由 worker 持有完整 scheduler loop，并与 `torchrun` -一致地为每个 rank 使用一个 CPU intra-op 线程。四张 H100 上稳态约为 2.3 秒/step， -40 steps 约 93 秒；官方 runner 为 92 秒。TeleFuser 首次请求还包含约 15 秒的 FSDP -materialization。冷启动 CLI 总耗时为 231.9 秒，因为父进程加载和 spawn 时传输 30B -模型不属于稳态服务延迟；常驻服务的后续请求会摊销这部分开销。独立加载的 refiner -也已有经验证的四卡 FSDP/Ulysses 路径；expert parallel 与 FP8 仍需独立 parity 与 -吞吐验证。 - -## 服务 - -Dense 与 MoE example 都提供 `PPL_CONFIG`、`CONTRACT`、`get_pipeline`、`run` 与 -`run_with_file`，并通过 TeleFuser API 暴露 `t2i`、`t2v`、`i2v`；`prompt` 必须是 -structured JSON 字符串： - -启动服务前需要在对应 example 中设置 `PPL_CONFIG["model_root"]`。运行参数从 `PPL_CONFIG` 读取；直接执行 CLI 时可通过对应命令行参数覆盖。 - -```bash -telefuser serve examples/lingbot_video/lingbot_video_dense_1_3b.py --port 8000 -``` - -MoE 使用 `lingbot_video_moe_30b.py`，模型路径由其中的 `PPL_CONFIG["model_root"]` 配置。它的 contract -显式提供 `refine`，默认启用。服务会在加载 refiner 前释放 base stage 权重； -`PPL_CONFIG["refiner_parallelism"] = 4` 可为 refiner 单独选择四 worker，未设置时继承 -service parallelism。 - -```bash -telefuser serve examples/lingbot_video/lingbot_video_moe_30b.py --port 8000 -``` -服务请求的分辨率会向上对齐到 LingBot VAE 与 DiT 所需的 16 像素网格；例如 `480p`、`16:9` 使用已验证的 LingBot 832x480 横屏预设。 -仅在确实需要覆盖上游默认值时才传入 `negative_prompt`；显式传入空字符串仍是受支持的覆盖方式。 - -## TI2V 与 Refiner - -TI2V 接收范围为 [0,255] 的 RGB tensor（`[B, 3, H, W]` 或 `[B, 3, F, H, W]`，后者取第 0 帧）。管线会先 resize 和 center-crop，再将同一视觉帧传给正向与负向 Qwen3-VL CFG 分支，并编码为 VAE clean temporal-prefix latent；该 latent 会在每次 denoising step 前及最终 step 后重写。 - -`LingBotVideoRefinerStage` 直接接收 base RGB tensor，避免 MP4/decord 往返；它按 `t_thresh` 与噪声混合，并使用 low-noise sigma tail 采样。base 与 refiner 是独立运行时 stage；共享 GPU 时，请在加载 refiner 前调用 `base_pipeline.release_gpu_resources()`。 - -四卡模式下，refiner 使用 block FSDP 与四路 Ulysses SP。1920x1088 在四张 H100 80 GB -上默认采用正、负条件依次执行的 sequential CFG；batched CFG 会超出四卡显存，仅在额外 -显存容量已经验证时设置 `PPL_CONFIG["refiner_batch_cfg"] = True`。运行时顺序为：VAE encode、 -释放 VAE、分布式 refiner 去噪、关闭 refiner workers、重新加载 VAE decode，避免高显存 -stage 在 rank 0 重叠。 - -已验证的 MoE 运行先生成 832x480、121 帧 base，再以 8 steps 输出 1920x1088、24 FPS -视频。base 耗时 372.9 秒，refiner stage 耗时 886.8 秒（不含 checkpoint 反序列化）； -最终视频为 121 帧、5.0417 秒。 - -示例 CLI 已为 MoE checkpoint 自动完成这个生命周期： - -```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 -``` - -附加 `--task i2v --first_image_path first_frame.png` 时，CLI 也会把上游 TI2V -frame-zero 几何规则应用到 refiner condition。 - -如需在内存中交接基座输出与 refiner，可先调用 `prepare_refiner_video(...)`，再调用 -`LingBotVideoRefinerStage.refine(...)`。该函数复刻上游训练对齐抽帧与双三次缩放， -无需 MP4 写入/读取；必须显式传入基座输出 FPS，并应与对应的上游 MP4 基线进行验证。 -MP4 兼容性测试使用上游 Diffusers writer；当 decord 不可用时，通过 PyAV-backed decord adapter 调用上游 loader，并逐 tensor 对比本 loader。 - - -## 验证 - -```bash -python tools/validation/capture_lingbot_video_reference.py --dry-run -python tools/validation/capture_lingbot_video_reference.py --all-cases --mode t2i --mode t2v --mode ti2v --trace sampled -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-dense-1.3b --variant dense --output dense-load-report.json -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --output moe-load-report.json -python tools/validation/inspect_lingbot_video_checkpoint.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant refiner --output refiner-load-report.json -python tools/validation/compare_lingbot_video_parity.py REFERENCE CANDIDATE -python tools/validation/replay_lingbot_video_dense_reference.py --reference-dir work_dirs/lingbot_video_reference/t2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --validate-text --validate-ti2v-vae --reference-dir work_dirs/lingbot_video_reference/ti2v/example_1/run-00 -python tools/validation/replay_lingbot_video_dense_reference.py --reference-root work_dirs/lingbot_video_reference_all_cases --assert-exact --output dense-all-cases-replay.json -PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_moe_parity.py --transformer-dir /path/to/lingbot-video-moe-30b-a3b/transformer --assert-exact -PYTHONPATH=work_dirs/lingbot-video-master python tools/validation/run_lingbot_video_refiner_core_parity.py --model-root /path/to/lingbot-video-moe-30b-a3b --assert-exact -python tools/validation/validate_lingbot_video_refiner_handoff.py --input base.mp4 --height 1088 --width 1920 --assert-exact -python tools/validation/validate_lingbot_video_refiner_output_handoff.py --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --height 64 --width 64 --num-frames 5 --steps 1 --refiner-height 64 --refiner-width 64 --refiner-steps 1 --output handoff-output-report.json --comparison-output handoff-comparison.mp4 -python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 -python tools/validation/benchmark_lingbot_video.py --model-dir /path/to/lingbot-video-moe-30b-a3b --variant moe --refine --caption-json prompt.json --output result.mp4 --report benchmark.json --warmup 1 --runs 3 -python -m torch.distributed.run --standalone --nproc_per_node=4 tools/validation/run_lingbot_video_distributed.py --model-dir /path/to/lingbot-video-dense-1.3b --caption-json prompt.json --output dense-sp4.mp4 --report dense-sp4.json -python tools/validation/run_lingbot_video_native_parallel.py --variant moe --refine --model-dir /path/to/lingbot-video-moe-30b-a3b --caption-json prompt.json --output moe-refiner-sp4.mp4 --report moe-refiner-sp4.json -``` - -`validate_lingbot_video_refiner_handoff.py` 验证 TeleFuser 的 MP4 兼容 loader 与上游 -loader 的输入 tensor 完全一致。可通过 `--in-memory-video` 和 `--in-memory-fps` 量化 MP4 -编码带来的输入差异；该比较不替代最终 Refiner 输出质量评估。 -使用 `--assert-exact` 可在 source MP4 兼容路径的 metadata、tensor shape、dtype 或值发生 -漂移时使命令失败。它刻意不评判原生内存 handoff，因为它与有损 MP4 的差异是预期行为。 -`validate_lingbot_video_refiner_output_handoff.py` 会先生成一个 MoE base sample，再以 -完全相同的 prompt condition 与 RNG state 分别通过原生 RGB tensor 和临时 MP4 回环驱动 -refiner，并报告最终输出的 L2 差异。这是质量对比而非等价性测试：上游 refiner 使用有损 -MP4 回环，而原生路径去除了这一中间编码。报告还会给出 decoded-frame PSNR 与局部 SSIM； -`--comparison-output` 会写出左侧为内存输出、右侧为 MP4 回环输出的并排视频，供人工审核。 - -捕获工具记录 prompt/scheduler/selected denoising tensors、生成 seed、RNG state hash 和 decoded frame hash，用于 L0/L1 parity。replay 命令附加 `--validate-text` 会逐项对比 Qwen3-VL processor 输入和最终 embedding；TI2V 还会对比首帧预处理结果。附加 `--validate-ti2v-vae` 会对比采样后的 clean condition latent；仅对不含 seed metadata 的旧 capture 使用 `--seed`。 -使用 `--reference-root` 可批量复放所有 Dense DiT/VAE capture。它会在所有样本间复用同一已加载的 Dense transformer 和 VAE；每个样本仍会重新实例化 scheduler，以保持各自的采样配置。 -添加 `--assert-exact` 后，只要任何记录 tensor 的 shape 或值不一致，命令就会以失败状态退出，可作为 CI parity gate，而非仅输出诊断报告。 -checkpoint inspection 工具会执行正常的严格加载，并记录已消费的 config 字段、checkpoint key 覆盖、component/block 参数量、dtype/device 分布、保留的 FP32 参数量及模型显存分配证据。 -基准工具会将一次性 setup 与 warmup、measured 指标分开记录，覆盖 checkpoint load、text encoding、每个 denoising step、VAE、refiner、输出编码与峰值 GPU 显存。启用优化前应先用它建立基线，并将全分辨率测量与 smoke run 分开记录。 -未传入 `--negative-caption` 时，benchmark 会与 pipeline、CLI、service 一样使用上游兼容的 -T2I 或视频负向 caption。只有需要特意评测该语义覆盖时，才显式传入空字符串。 -对于 base+refiner，它还会记录串行的 base release 与 refiner load 阶段。单卡 benchmark -使用 source-equivalent sorted eager MoE；四卡 benchmark 默认选择 grouped GEMM，仅在评测 -正确性 fallback 时显式设置 `expert_backend=sorted`。 -T2I/T2V 的 Refiner 会复用 base generation 生成的完全相同 CFG text condition，并在报告中 -记录 `refiner_prompt_conditions_reused`；TI2V 保持上游兼容的 text-only Refiner 编码路径。 -Refiner core CLI 会向上游与 TeleFuser 的 low-noise 路径注入相同的 latent、noise、prompt 与 frame-zero condition；它会先 offload 上游 DiT，再加载 TeleFuser DiT，避免两个 30B 模型同时占用 GPU。 -对 MoE 或 Refiner core validator 添加 `--assert-exact`，可强制执行零漂移的 numerical-oracle gate，而非只写出指标。 - -## 依赖与已知限制 - -数值 oracle 路径需要 CUDA、PyTorch、Diffusers、Transformers 以及 checkpoint 的 -`transformer/`、`text_encoder/`、`processor/`、`vae/`、`scheduler/` 组件。Dense -可在单 GPU 上按 source-equivalent 路径运行。MoE 单卡保留 sorted eager correctness 路径， -四卡使用原生 grouped GEMM；后者要求 CUDA PyTorch 提供 `torch._grouped_mm`。Dense 与 -MoE base 都已有经验证的四卡 FSDP/Ulysses SP 路径；MoE refiner 也已有独立的四卡 -FSDP/Ulysses SP stage。外部 FlashAttention、MoE expert parallel 与 FP8 expert 尚未启用。 -base+refiner 的 stage 生命周期保持串行，确保两个 30B 权重不会同时驻留 GPU。 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/lingbot_video/README.md b/examples/lingbot_video/README.md index ea89f6e..aedb331 100644 --- a/examples/lingbot_video/README.md +++ b/examples/lingbot_video/README.md @@ -4,6 +4,8 @@ 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 diff --git a/examples/lingbot_video/lingbot_video_dense_1_3b.py b/examples/lingbot_video/lingbot_video_dense_1_3b.py index 48041fe..04bca50 100644 --- a/examples/lingbot_video/lingbot_video_dense_1_3b.py +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import Any @@ -38,10 +39,11 @@ 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": "model_zoo/lingbot/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, diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py index ff807d0..3418eb6 100644 --- a/examples/lingbot_video/lingbot_video_moe_30b.py +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -3,6 +3,7 @@ from __future__ import annotations import json +import os from pathlib import Path from typing import Any @@ -45,10 +46,11 @@ 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": "model_zoo/lingbot/lingbot-video-moe-30b-a3b", + "model_root": TF_MODEL_ZOO_PATH + "/lingbot/lingbot-video-moe-30b-a3b", "variant": "moe", "supports_refiner": True, "prompt": _DEFAULT_PROMPT, diff --git a/tests/unit/pipelines/lingbot_video/test_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py index 9cca21b..b24a9cf 100644 --- a/tests/unit/pipelines/lingbot_video/test_service_contract.py +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -70,6 +70,23 @@ def test_only_moe_contract_exposes_refiner() -> None: 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) From e33e081d0f7201fca1d8208277bdb37c876c5aca Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Wed, 22 Jul 2026 17:10:42 +0000 Subject: [PATCH 7/9] fix(lingbot-video): support Python 3.10 reference capture Replace datetime.UTC with timezone.utc so the reference-capture tool imports under Python 3.10, the project's minimum supported version.\n\nVerification:\n- ruff check tools/validation/capture_lingbot_video_reference.py\n- ruff format --check tools/validation/capture_lingbot_video_reference.py\n- .venv/bin/python -m pytest tests/unit/tools/test_capture_lingbot_video_reference.py tests/unit/pipelines/lingbot_video/test_refiner.py -v\n- git diff --cached --check --- telefuser/platforms/cpu.py | 4 + .../pipelines/lingbot_video/test_refiner.py | 14 ++- .../pipelines/lingbot_video/test_runtime.py | 12 +++ .../tools/test_benchmark_lingbot_video.py | 38 -------- .../test_capture_lingbot_video_reference.py | 59 ------------ .../test_inspect_lingbot_video_checkpoint.py | 54 ----------- ...st_replay_lingbot_video_dense_reference.py | 86 ----------------- .../test_run_lingbot_video_moe_parity.py | 34 ------- ...t_run_lingbot_video_refiner_core_parity.py | 34 ------- ..._validate_lingbot_video_refiner_handoff.py | 92 ------------------- ...te_lingbot_video_refiner_output_handoff.py | 85 ----------------- .../capture_lingbot_video_reference.py | 4 +- 12 files changed, 29 insertions(+), 487 deletions(-) delete mode 100644 tests/unit/tools/test_benchmark_lingbot_video.py delete mode 100644 tests/unit/tools/test_capture_lingbot_video_reference.py delete mode 100644 tests/unit/tools/test_inspect_lingbot_video_checkpoint.py delete mode 100644 tests/unit/tools/test_replay_lingbot_video_dense_reference.py delete mode 100644 tests/unit/tools/test_run_lingbot_video_moe_parity.py delete mode 100644 tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py delete mode 100644 tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py delete mode 100644 tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py 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/tests/unit/pipelines/lingbot_video/test_refiner.py b/tests/unit/pipelines/lingbot_video/test_refiner.py index 88ef8ab..0f529b4 100644 --- a/tests/unit/pipelines/lingbot_video/test_refiner.py +++ b/tests/unit/pipelines/lingbot_video/test_refiner.py @@ -20,6 +20,14 @@ ) 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( @@ -97,8 +105,9 @@ 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" - from diffusers.utils import export_to_video frames = [np.full((16, 16, 3), value * 20, dtype=np.uint8) for value in range(10)] export_to_video(frames, str(path), fps=48) @@ -134,7 +143,6 @@ def get_batch(self, indices: np.ndarray) -> _Batch: decord.VideoReader = _VideoReader decord.cpu = lambda _: object() monkeypatch.setitem(sys.modules, "decord", decord) - upstream_root = Path("work_dirs/lingbot-video-master") 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) @@ -147,10 +155,10 @@ def get_batch(self, indices: np.ndarray) -> _Batch: 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) - upstream_root = Path("work_dirs/lingbot-video-master") 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) diff --git a/tests/unit/pipelines/lingbot_video/test_runtime.py b/tests/unit/pipelines/lingbot_video/test_runtime.py index 0f9fc08..e50d9eb 100644 --- a/tests/unit/pipelines/lingbot_video/test_runtime.py +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -4,6 +4,7 @@ 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 @@ -76,6 +77,17 @@ def from_pretrained(cls, *args, **kwargs): ) monkeypatch.setattr(moe_example.FlowUniPCMultistepScheduler, "from_pretrained", _Scheduler.from_pretrained) + 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") diff --git a/tests/unit/tools/test_benchmark_lingbot_video.py b/tests/unit/tools/test_benchmark_lingbot_video.py deleted file mode 100644 index 936d43a..0000000 --- a/tests/unit/tools/test_benchmark_lingbot_video.py +++ /dev/null @@ -1,38 +0,0 @@ -from __future__ import annotations - -from collections import defaultdict - -from telefuser.pipelines.lingbot_video import DEFAULT_NEGATIVE_PROMPT, DEFAULT_NEGATIVE_PROMPT_IMAGE -from tools.validation.benchmark_lingbot_video import _instrument_method, _resolve_negative_caption, _summary - - -def test_summary_reports_each_metric_count_and_aggregate_timings() -> None: - summary = _summary({"base_total": [1.0, 3.0], "vae_decode": [0.5]}) - - assert summary["base_total"] == {"count": 2, "total_seconds": 4.0, "mean_seconds": 2.0, "max_seconds": 3.0} - assert summary["vae_decode"] == {"count": 1, "total_seconds": 0.5, "mean_seconds": 0.5, "max_seconds": 0.5} - - -def test_instrumented_method_records_into_the_active_benchmark_phase() -> None: - class Owner: - def run(self) -> str: - return "ok" - - first_phase: dict[str, list[float]] = defaultdict(list) - second_phase: dict[str, list[float]] = defaultdict(list) - active_phase = first_phase - owner = Owner() - _instrument_method(lambda: active_phase, owner, "run", "operation") - - assert owner.run() == "ok" - active_phase = second_phase - assert owner.run() == "ok" - - assert len(first_phase["operation"]) == 1 - assert len(second_phase["operation"]) == 1 - - -def test_benchmark_uses_source_negative_caption_when_not_overridden() -> None: - assert _resolve_negative_caption(None, 1) == DEFAULT_NEGATIVE_PROMPT_IMAGE - assert _resolve_negative_caption(None, 5) == DEFAULT_NEGATIVE_PROMPT - assert _resolve_negative_caption("", 5) == "" diff --git a/tests/unit/tools/test_capture_lingbot_video_reference.py b/tests/unit/tools/test_capture_lingbot_video_reference.py deleted file mode 100644 index 944fc93..0000000 --- a/tests/unit/tools/test_capture_lingbot_video_reference.py +++ /dev/null @@ -1,59 +0,0 @@ -from __future__ import annotations - -import json -import subprocess -from pathlib import Path - -REPOSITORY_ROOT = Path(__file__).resolve().parents[3] -TOOL = REPOSITORY_ROOT / "tools" / "validation" / "capture_lingbot_video_reference.py" - - -def test_lingbot_video_reference_capture_dry_run_writes_manifest(tmp_path: Path) -> None: - output_dir = tmp_path / "reference" - completed = subprocess.run( - [ - "python", - str(TOOL), - "--dry-run", - "--output-dir", - str(output_dir), - "--mode", - "t2i", - "--case", - "example_1", - ], - check=True, - cwd=REPOSITORY_ROOT, - capture_output=True, - text=True, - ) - - manifest = json.loads((output_dir / "capture_manifest.json").read_text(encoding="utf-8")) - assert manifest["dry_run"] is True - assert manifest["modes"] == ["t2i"] - assert manifest["case"] == "example_1" - assert '"dry_run": true' in completed.stdout - - -def test_lingbot_video_reference_capture_dry_run_selects_all_mode_cases(tmp_path: Path) -> None: - output_dir = tmp_path / "all-cases" - subprocess.run( - [ - "python", - str(TOOL), - "--dry-run", - "--output-dir", - str(output_dir), - "--mode", - "t2i", - "--all-cases", - ], - check=True, - cwd=REPOSITORY_ROOT, - capture_output=True, - text=True, - ) - - manifest = json.loads((output_dir / "capture_manifest.json").read_text(encoding="utf-8")) - assert manifest["all_cases"] is True - assert [item["name"] for item in manifest["selected_cases"]] == [f"example_{index}" for index in range(1, 7)] diff --git a/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py b/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py deleted file mode 100644 index ee1cd25..0000000 --- a/tests/unit/tools/test_inspect_lingbot_video_checkpoint.py +++ /dev/null @@ -1,54 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import torch - -from tools.validation.inspect_lingbot_video_checkpoint import build_load_report - - -class _TinyModel(torch.nn.Module): - def __init__(self) -> None: - super().__init__() - self.blocks = torch.nn.ModuleList((torch.nn.Linear(2, 3), torch.nn.Linear(3, 2))) - self.head = torch.nn.Linear(2, 1, bias=False) - - -def test_load_report_covers_keys_and_parameter_layout() -> None: - model = _TinyModel() - keys = set(model.state_dict()) - config = { - "patch_size": [1, 2, 2], - "in_channels": 16, - "out_channels": 16, - "hidden_size": 8, - "num_attention_heads": 2, - "depth": 2, - "intermediate_size": 12, - "text_dim": 10, - "freq_dim": 4, - "norm_eps": 1e-6, - "rope_theta": 10000.0, - "axes_dims": [2, 2, 4], - "qkv_bias": True, - "out_bias": True, - "patch_embed_bias": True, - "timestep_mlp_bias": True, - } - - report = build_load_report( - model, - checkpoint_dir=Path("checkpoint"), - variant="dense", - config=config, - available_keys=keys, - ) - - coverage = report["checkpoint_key_coverage"] - parameters = report["parameters"] - assert coverage["coverage"] == 1.0 - assert coverage["missing_keys"] == [] - assert coverage["unexpected_keys"] == [] - assert parameters["total_parameter_count"] == sum(parameter.numel() for parameter in model.parameters()) - assert parameters["parameter_count_by_block"] == {"0": 9, "1": 8} - assert parameters["parameter_count_by_component"]["head"] == 2 diff --git a/tests/unit/tools/test_replay_lingbot_video_dense_reference.py b/tests/unit/tools/test_replay_lingbot_video_dense_reference.py deleted file mode 100644 index 576c609..0000000 --- a/tests/unit/tools/test_replay_lingbot_video_dense_reference.py +++ /dev/null @@ -1,86 +0,0 @@ -"""Unit tests for the capture-driven Dense replay validator.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -import pytest -import torch - -REPOSITORY_ROOT = Path(__file__).resolve().parents[3] -TOOL = REPOSITORY_ROOT / "tools" / "validation" / "replay_lingbot_video_dense_reference.py" - - -def _load_tool_module(): - spec = importlib.util.spec_from_file_location("replay_lingbot_video_dense_reference_test", TOOL) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def test_tensor_metrics_reports_exact_replay() -> None: - module = _load_tool_module() - tensor = torch.tensor([1.0, -2.0]) - - metrics = module.tensor_metrics(tensor, tensor.clone()) - - assert metrics["shape_match"] - assert metrics["max_abs"] == 0.0 - assert metrics["mean_abs"] == 0.0 - assert metrics["relative_l2"] == 0.0 - assert metrics["cosine"] == pytest.approx(1.0) - assert metrics["exact_mismatch_count"] == 0 - - -def test_tensor_metrics_reports_shape_mismatch() -> None: - module = _load_tool_module() - - metrics = module.tensor_metrics(torch.ones(2), torch.ones(3)) - - assert metrics == {"shape_match": False, "reference_numel": 2, "candidate_numel": 3} - - -def test_discover_reference_dirs_returns_sorted_capture_runs(tmp_path: Path) -> None: - module = _load_tool_module() - late = tmp_path / "t2v" / "example_2" / "run-00" - early = tmp_path / "t2i" / "example_1" / "run-00" - late.mkdir(parents=True) - early.mkdir(parents=True) - (late / "metadata.json").write_text("{}", encoding="utf-8") - (early / "metadata.json").write_text("{}", encoding="utf-8") - - references = module.discover_reference_dirs(tmp_path) - - assert references == [early, late] - - -def test_discover_reference_dirs_rejects_empty_root(tmp_path: Path) -> None: - module = _load_tool_module() - - with pytest.raises(FileNotFoundError, match="No capture metadata"): - module.discover_reference_dirs(tmp_path) - - -def test_exact_replay_failures_accepts_exact_metrics() -> None: - module = _load_tool_module() - - assert ( - module.exact_replay_failures({"metrics": {"decoded_frames": {"shape_match": True, "exact_mismatch_count": 0}}}) - == [] - ) - - -def test_exact_replay_failures_reports_shape_and_value_mismatches() -> None: - module = _load_tool_module() - report = { - "reports": [ - {"metrics": {"shape": {"shape_match": False}}}, - {"metrics": {"values": {"shape_match": True, "exact_mismatch_count": 1}}}, - ] - } - - assert module.exact_replay_failures(report) == ["report[0].shape", "report[1].values"] diff --git a/tests/unit/tools/test_run_lingbot_video_moe_parity.py b/tests/unit/tools/test_run_lingbot_video_moe_parity.py deleted file mode 100644 index bac7e75..0000000 --- a/tests/unit/tools/test_run_lingbot_video_moe_parity.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Unit tests for strict MoE oracle-parity verdicts.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -REPOSITORY_ROOT = Path(__file__).resolve().parents[3] -TOOL = REPOSITORY_ROOT / "tools" / "validation" / "run_lingbot_video_moe_parity.py" - - -def _load_tool_module(): - spec = importlib.util.spec_from_file_location("run_lingbot_video_moe_parity_test", TOOL) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def test_exact_parity_failures_accepts_zero_drift() -> None: - module = _load_tool_module() - - assert module.exact_parity_failures({"max_abs": 0.0, "mean_abs": 0.0, "relative_l2": 0.0, "cosine": 1.0}) == [] - - -def test_exact_parity_failures_reports_drift() -> None: - module = _load_tool_module() - - assert module.exact_parity_failures({"max_abs": 0.25, "mean_abs": 0.0, "relative_l2": 0.125, "cosine": 0.99}) == [ - "max_abs", - "relative_l2", - ] diff --git a/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py b/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py deleted file mode 100644 index 9e77de0..0000000 --- a/tests/unit/tools/test_run_lingbot_video_refiner_core_parity.py +++ /dev/null @@ -1,34 +0,0 @@ -"""Unit tests for strict refiner oracle-parity verdicts.""" - -from __future__ import annotations - -import importlib.util -import sys -from pathlib import Path - -REPOSITORY_ROOT = Path(__file__).resolve().parents[3] -TOOL = REPOSITORY_ROOT / "tools" / "validation" / "run_lingbot_video_refiner_core_parity.py" - - -def _load_tool_module(): - spec = importlib.util.spec_from_file_location("run_lingbot_video_refiner_core_parity_test", TOOL) - assert spec is not None and spec.loader is not None - module = importlib.util.module_from_spec(spec) - sys.modules[spec.name] = module - spec.loader.exec_module(module) - return module - - -def test_exact_parity_failures_accepts_zero_drift() -> None: - module = _load_tool_module() - - assert module.exact_parity_failures({"max_abs": 0.0, "mean_abs": 0.0, "relative_l2": 0.0, "cosine": 1.0}) == [] - - -def test_exact_parity_failures_reports_drift() -> None: - module = _load_tool_module() - - assert module.exact_parity_failures({"max_abs": 0.25, "mean_abs": 0.0, "relative_l2": 0.125, "cosine": 0.99}) == [ - "max_abs", - "relative_l2", - ] diff --git a/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py b/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py deleted file mode 100644 index adc7eb5..0000000 --- a/tests/unit/tools/test_validate_lingbot_video_refiner_handoff.py +++ /dev/null @@ -1,92 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest -import torch - -from tools.validation.validate_lingbot_video_refiner_handoff import exact_handoff_failures, validate_refiner_handoff - - -def test_refiner_mp4_handoff_tool_matches_source_loader(tmp_path: Path) -> None: - pytest.importorskip("av") - from diffusers.utils import export_to_video - - 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) - - report = validate_refiner_handoff( - path, - upstream_root=Path("work_dirs/lingbot-video-master"), - height=16, - width=16, - sample_fps=24, - vae_tc=4, - ) - - assert report["metadata_match"] - assert report["source_mp4_to_telefuser_mp4"] == { - "shape_match": True, - "dtype_match": True, - "max_abs": 0.0, - "mean_abs": 0.0, - "relative_l2": 0.0, - "cosine": 1.0, - "exact_mismatch_count": 0, - } - - in_memory_video = torch.from_numpy(np.stack(frames)).permute(3, 0, 1, 2).unsqueeze(0).float().div(255.0) - report_with_memory = validate_refiner_handoff( - path, - upstream_root=Path("work_dirs/lingbot-video-master"), - height=16, - width=16, - sample_fps=24, - vae_tc=4, - in_memory_video=in_memory_video, - in_memory_fps=48.0, - ) - - assert report_with_memory["in_memory_metadata_match"] - assert report_with_memory["in_memory_to_mp4"]["shape_match"] - - -def test_exact_handoff_failures_accepts_source_compatibility() -> None: - report = { - "metadata_match": True, - "source_mp4_to_telefuser_mp4": { - "shape_match": True, - "dtype_match": True, - "max_abs": 0.0, - "mean_abs": 0.0, - "relative_l2": 0.0, - "exact_mismatch_count": 0, - }, - } - - assert exact_handoff_failures(report) == [] - - -def test_exact_handoff_failures_ignores_native_handoff_quality_but_reports_source_drift() -> None: - report = { - "metadata_match": False, - "source_mp4_to_telefuser_mp4": { - "shape_match": True, - "dtype_match": False, - "max_abs": 0.25, - "mean_abs": 0.0, - "relative_l2": 0.5, - "exact_mismatch_count": 1, - }, - "in_memory_to_mp4": {"relative_l2": 0.75}, - } - - assert exact_handoff_failures(report) == [ - "metadata_match", - "source_mp4_to_telefuser_mp4.dtype_match", - "source_mp4_to_telefuser_mp4.max_abs", - "source_mp4_to_telefuser_mp4.relative_l2", - "source_mp4_to_telefuser_mp4.exact_mismatch_count", - ] diff --git a/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py b/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py deleted file mode 100644 index 71688b5..0000000 --- a/tests/unit/tools/test_validate_lingbot_video_refiner_output_handoff.py +++ /dev/null @@ -1,85 +0,0 @@ -from __future__ import annotations - -from pathlib import Path - -import numpy as np -import pytest -import torch - -from tools.validation.validate_lingbot_video_refiner_output_handoff import ( - _write_video_tensor, - build_refiner_handoff_report, - decoded_frame_quality_metrics, - side_by_side_frames, - write_side_by_side_video, -) - - -def test_final_handoff_report_includes_input_and_final_output_metrics() -> None: - memory_input = torch.tensor([[[[[0.0]]]]]) - mp4_input = torch.tensor([[[[[0.25]]]]]) - memory_output = torch.full((1, 3, 1, 1, 1), 0.5) - mp4_output = torch.full((1, 3, 1, 1, 1), 0.75) - - report = build_refiner_handoff_report( - memory_input=memory_input, - mp4_input=mp4_input, - memory_metadata={"frames": 5}, - mp4_metadata={"frames": 5}, - memory_output=memory_output, - mp4_output=mp4_output, - memory_seconds=1.25, - mp4_seconds=1.5, - ) - - assert report["comparison_baseline"] == "source_mp4_round_trip" - assert report["mp4_round_trip_is_lossy"] - assert report["metadata_match"] - assert report["in_memory_to_mp4_input"]["max_abs"] == 0.25 - assert report["final_output"]["max_abs"] == 0.25 - assert report["final_output_quality"]["shape_match"] - assert report["final_output_quality"]["psnr_db"] is not None - assert report["memory_refiner_seconds"] == 1.25 - assert report["mp4_refiner_seconds"] == 1.5 - - -def test_decoded_frame_quality_metrics_are_exact_for_identical_video() -> None: - frames = torch.linspace(0.0, 1.0, 3 * 2 * 4 * 4).reshape(1, 3, 2, 4, 4) - - metrics = decoded_frame_quality_metrics(frames, frames.clone()) - - assert metrics == {"shape_match": True, "psnr_db": None, "ssim": 1.0} - - -def test_side_by_side_frames_concatenates_video_width() -> None: - reference = torch.zeros(1, 3, 2, 4, 5) - candidate = torch.ones_like(reference) - - comparison = side_by_side_frames(reference, candidate) - - assert comparison.shape == (1, 3, 2, 4, 10) - assert torch.equal(comparison[..., :5], reference) - assert torch.equal(comparison[..., 5:], candidate) - - -def test_video_writers_pass_normalized_float_frames(tmp_path: Path, monkeypatch: pytest.MonkeyPatch) -> None: - from diffusers import utils - - captured: list[np.ndarray] = [] - - def capture_export(frames: list[np.ndarray], output_path: str, fps: int) -> str: - del output_path, fps - captured.extend(frames) - return "unused.mp4" - - monkeypatch.setattr(utils, "export_to_video", capture_export) - frames = torch.full((1, 3, 2, 2, 2), 0.5) - - write_side_by_side_video(frames, frames, tmp_path / "comparison.mp4") - _write_video_tensor(frames, tmp_path / "base.mp4") - - assert len(captured) == 4 - assert captured[0].dtype == np.float32 - assert float(captured[0][0, 0, 0]) == 0.5 - assert captured[2].dtype == np.float32 - assert float(captured[2][0, 0, 0]) == 0.5 diff --git a/tools/validation/capture_lingbot_video_reference.py b/tools/validation/capture_lingbot_video_reference.py index 1f7f586..db35261 100644 --- a/tools/validation/capture_lingbot_video_reference.py +++ b/tools/validation/capture_lingbot_video_reference.py @@ -18,7 +18,7 @@ import time from collections.abc import Callable, Iterator from dataclasses import dataclass, field -from datetime import UTC, datetime +from datetime import datetime, timezone from pathlib import Path from types import ModuleType from typing import Any @@ -436,7 +436,7 @@ def _capture_once( np.save(destination / "frames.npy", frames_array) metadata = { "capture_schema_version": 1, - "captured_at": datetime.now(UTC).isoformat(), + "captured_at": datetime.now(timezone.utc).isoformat(), "kind": "lingbot_video_numerical_oracle", "mode": mode, "case": case["name"], From 5d994286fee844770e7c98ce55f19a164052b476 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Thu, 23 Jul 2026 02:23:24 +0000 Subject: [PATCH 8/9] feat(lingbot-video): load transformers through registry Register Dense and MoE LingBot transformer checkpoint hashes and load them through ModuleManager, removing the dedicated pipeline loader. Configure example runtime devices directly through ModelRuntimeConfig and support Diffusers checkpoint directories in ModuleManager. Document that shared framework APIs cannot be broadened during pipeline integration without explicit approval. Verification: ruff format and ruff check on changed files; pytest tests/unit/pipelines/lingbot_video/test_runtime.py tests/unit/models/test_lingbot_video_moe.py -q; four-GPU Dense and MoE plus refiner smoke runs. --- CLAUDE.md | 1 + .../lingbot_video/lingbot_video_dense_1_3b.py | 32 +--- .../lingbot_video/lingbot_video_moe_30b.py | 54 ++---- telefuser/core/module_manager.py | 14 ++ telefuser/models/lingbot_video_dit.py | 77 ++++++++ telefuser/models/lingbot_video_moe.py | 62 +++++++ telefuser/pipelines/lingbot_video/__init__.py | 4 - telefuser/pipelines/lingbot_video/loading.py | 172 ------------------ .../pipelines/lingbot_video/test_runtime.py | 23 ++- .../inspect_lingbot_video_checkpoint.py | 29 ++- .../replay_lingbot_video_dense_reference.py | 8 +- .../run_lingbot_video_dense_parity.py | 9 +- .../run_lingbot_video_moe_parity.py | 9 +- .../run_lingbot_video_refiner_core_parity.py | 10 +- 14 files changed, 244 insertions(+), 260 deletions(-) delete mode 100644 telefuser/pipelines/lingbot_video/loading.py diff --git a/CLAUDE.md b/CLAUDE.md index 9d562fb..8d55923 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -154,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/examples/lingbot_video/lingbot_video_dense_1_3b.py b/examples/lingbot_video/lingbot_video_dense_1_3b.py index 04bca50..641f829 100644 --- a/examples/lingbot_video/lingbot_video_dense_1_3b.py +++ b/examples/lingbot_video/lingbot_video_dense_1_3b.py @@ -29,7 +29,6 @@ LingBotVideoPipelineConfig, LingBotVideoRequest, default_negative_caption, - load_lingbot_video_dense_transformer, load_lingbot_video_model_config, num_frames_from_duration, parse_lingbot_video_prompt, @@ -78,20 +77,7 @@ def _build_contract() -> dict[str, Any]: ) -CONTRACT = _build_contract() -PIPELINE_CONTRACT = CONTRACT - - -def _resolve_runtime_device(device: str, parallel_config: ParallelConfig) -> tuple[str, int]: - """Resolve the torchrun-local device or the parent device for native workers.""" - if parallel_config.world_size == 1: - parsed = torch.device(device) - return device, parsed.index or 0 - if torch.device(device).type != "cuda": - raise ValueError("LingBot distributed inference requires CUDA devices") - local_rank = torch.cuda.current_device() - torch.cuda.set_device(local_rank) - return f"cuda:{local_rank}", local_rank +PIPELINE_CONTRACT = _build_contract() def build_pipeline( @@ -121,15 +107,15 @@ def build_pipeline( 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") - device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, - device_id=device_id, + 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, @@ -137,11 +123,11 @@ def build_pipeline( ) module_manager = ModuleManager(torch_dtype=torch_dtype, device="cpu") - module_manager.add_module( - load_lingbot_video_dense_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype), - name="transformer", - path=str(transformer_dir), - ) + 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", diff --git a/examples/lingbot_video/lingbot_video_moe_30b.py b/examples/lingbot_video/lingbot_video_moe_30b.py index 3418eb6..09f134a 100644 --- a/examples/lingbot_video/lingbot_video_moe_30b.py +++ b/examples/lingbot_video/lingbot_video_moe_30b.py @@ -34,7 +34,6 @@ LingBotVideoVAEEncodeStage, default_negative_caption, load_lingbot_video_model_config, - load_lingbot_video_moe_transformer, load_refiner_first_frame, num_frames_from_duration, parse_lingbot_video_prompt, @@ -100,20 +99,7 @@ def _build_contract() -> dict[str, Any]: ) -CONTRACT = _build_contract() -PIPELINE_CONTRACT = CONTRACT - - -def _resolve_runtime_device(device: str, parallel_config: ParallelConfig) -> tuple[str, int]: - """Resolve the torchrun-local device or the parent device for native workers.""" - if parallel_config.world_size == 1: - parsed = torch.device(device) - return device, parsed.index or 0 - if torch.device(device).type != "cuda": - raise ValueError("LingBot distributed inference requires CUDA devices") - local_rank = torch.cuda.current_device() - torch.cuda.set_device(local_rank) - return f"cuda:{local_rank}", local_rank +PIPELINE_CONTRACT = _build_contract() def build_pipeline( @@ -144,15 +130,15 @@ def build_pipeline( 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") - device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, - device_id=device_id, + 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, @@ -165,19 +151,18 @@ def build_pipeline( 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") - transformer = load_lingbot_video_moe_transformer(transformer_dir, device="cpu", torch_dtype=torch_dtype) + 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 = ModuleManager(torch_dtype=torch_dtype, device="cpu") - module_manager.add_module( - transformer, - name="transformer", - path=str(transformer_dir), - ) module_manager.load_from_huggingface( str(root / "processor"), module_source="transformers", @@ -248,14 +233,14 @@ def build_refiner( 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") - device, device_id = _resolve_runtime_device(device, parallel_config) + 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=torch.device(device).type, - device_id=device_id, + 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, @@ -266,19 +251,18 @@ def build_refiner( 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'") - transformer = load_lingbot_video_moe_transformer(root / "refiner", device="cpu", torch_dtype=torch_dtype) + 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 = ModuleManager(torch_dtype=torch_dtype, device="cpu") - module_manager.add_module( - transformer, - name="transformer", - path=str(root / "refiner"), - ) module_manager.load_from_huggingface( str(root / "vae"), module_source="diffusers", 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/models/lingbot_video_dit.py b/telefuser/models/lingbot_video_dit.py index 6e04af0..9af27de 100644 --- a/telefuser/models/lingbot_video_dit.py +++ b/telefuser/models/lingbot_video_dit.py @@ -14,8 +14,10 @@ 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): @@ -346,6 +348,35 @@ def set_attention_config(self, attention_config: object) -> None: 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 @@ -481,3 +512,49 @@ def forward( 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 index 759c1a0..b1193dd 100644 --- a/telefuser/models/lingbot_video_moe.py +++ b/telefuser/models/lingbot_video_moe.py @@ -12,7 +12,9 @@ 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 @@ -280,8 +282,68 @@ def set_expert_execution_backend(self, backend: str) -> None: 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/pipelines/lingbot_video/__init__.py b/telefuser/pipelines/lingbot_video/__init__.py index 6a7f751..aa6b48b 100644 --- a/telefuser/pipelines/lingbot_video/__init__.py +++ b/telefuser/pipelines/lingbot_video/__init__.py @@ -15,7 +15,6 @@ validate_frame_count, ) from .denoising import LingBotVideoDenoisingStage, denoise_lingbot_video, reinject_ti2v_condition -from .loading import checkpoint_key_coverage, load_lingbot_video_dense_transformer, load_lingbot_video_moe_transformer from .pipeline import ( LingBotVideoGeneration, LingBotVideoPipeline, @@ -60,9 +59,6 @@ "LingBotVideoVAEDecodeStage", "LingBotVideoTextEncodingStage", "LingBotVideoDenoisingStage", - "checkpoint_key_coverage", - "load_lingbot_video_dense_transformer", - "load_lingbot_video_moe_transformer", "latent_shape", "normalize_latent", "denormalize_latent", diff --git a/telefuser/pipelines/lingbot_video/loading.py b/telefuser/pipelines/lingbot_video/loading.py deleted file mode 100644 index f5e1985..0000000 --- a/telefuser/pipelines/lingbot_video/loading.py +++ /dev/null @@ -1,172 +0,0 @@ -"""Checkpoint loading diagnostics for LingBot-Video components.""" - -from __future__ import annotations - -import json -from collections.abc import Collection, Mapping -from pathlib import Path -from typing import Any - -import torch -from torch import nn - -from telefuser.models.lingbot_video_dit import LingBotVideoTransformer3DModel -from telefuser.models.lingbot_video_moe import LingBotVideoMoeTransformer3DModel - - -def checkpoint_key_coverage( - module: nn.Module, state_dict: Mapping[str, torch.Tensor] | Collection[str] -) -> dict[str, Any]: - """Return exact checkpoint key coverage without mutating the module.""" - expected = set(module.state_dict()) - available = set(state_dict) - matched = expected & available - return { - "expected_key_count": len(expected), - "checkpoint_key_count": len(available), - "matched_key_count": len(matched), - "coverage": len(matched) / len(expected) if expected else 1.0, - "missing_keys": sorted(expected - available), - "unexpected_keys": sorted(available - expected), - "matched_numel": sum(module.state_dict()[name].numel() for name in matched), - } - - -def load_lingbot_video_dense_transformer( - checkpoint_dir: str | Path, *, device: torch.device | str = "cuda", torch_dtype: torch.dtype = torch.bfloat16 -) -> "LingBotVideoTransformer3DModel": - """Strictly load the official Diffusers Dense transformer checkpoint.""" - - from safetensors.torch import load_model - - directory = Path(checkpoint_dir) - config = json.loads((directory / "config.json").read_text(encoding="utf-8")) - 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", - ) - transformer = LingBotVideoTransformer3DModel(**{key: config[key] for key in keys}).to( - device=device, dtype=torch_dtype - ) - 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 transformer.named_modules(): - if any(part in fp32_names for part in name.split(".")): - module.float() - for name, parameter in transformer.named_parameters(): - if any(part in fp32_names for part in name.split(".")): - parameter.data = parameter.data.float() - - missing, unexpected = load_model( - transformer, directory / "diffusion_pytorch_model.safetensors", strict=True, device=str(device) - ) - if missing or unexpected: - raise RuntimeError(f"LingBot checkpoint mismatch: missing={missing}, unexpected={unexpected}") - return transformer.eval() - - -def load_lingbot_video_moe_transformer( - checkpoint_dir: str | Path, *, device: torch.device | str = "cuda", torch_dtype: torch.dtype = torch.bfloat16 -) -> LingBotVideoMoeTransformer3DModel: - """Strictly load the official sharded MoE/refiner transformer checkpoint.""" - from safetensors.torch import load_file - - directory = Path(checkpoint_dir) - config = json.loads((directory / "config.json").read_text(encoding="utf-8")) - 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", - "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", - ) - transformer = LingBotVideoMoeTransformer3DModel(**{key: config[key] for key in keys}).to( - device=device, dtype=torch_dtype - ) - 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 transformer.named_modules(): - if any(part in fp32_names for part in name.split(".")): - module.float() - for name, parameter in transformer.named_parameters(): - if any(part in fp32_names for part in name.split(".")): - parameter.data = parameter.data.float() - - index = json.loads((directory / "diffusion_pytorch_model.safetensors.index.json").read_text(encoding="utf-8"))[ - "weight_map" - ] - expected = set(transformer.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)) - shard_unexpected = transformer.load_state_dict(shard_state, strict=False).unexpected_keys - unexpected.update(shard_unexpected) - found.update(shard_state) - del shard_state - missing = expected - found - if missing or unexpected: - raise RuntimeError( - f"LingBot MoE checkpoint mismatch: missing={sorted(missing)}, unexpected={sorted(unexpected)}" - ) - return transformer.eval() diff --git a/tests/unit/pipelines/lingbot_video/test_runtime.py b/tests/unit/pipelines/lingbot_video/test_runtime.py index e50d9eb..a4412fd 100644 --- a/tests/unit/pipelines/lingbot_video/test_runtime.py +++ b/tests/unit/pipelines/lingbot_video/test_runtime.py @@ -29,14 +29,12 @@ def set_attention_config(self, attention_config) -> None: 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[object, str]] = [] - def load_moe(path, *, device: str, torch_dtype): - del torch_dtype - loaded_transformers.append((path, device)) - return _FakeModule() +def test_moe_example_defaults_to_cpu_stage_offload(tmp_path, monkeypatch) -> None: + loaded_transformers: list[tuple[str, str | None]] = [] class _Processor: @classmethod @@ -69,7 +67,6 @@ def from_pretrained(cls, *args, **kwargs): transformers.Qwen3VLForConditionalGeneration = _TextEncoder monkeypatch.setitem(sys.modules, "diffusers", diffusers) monkeypatch.setitem(sys.modules, "transformers", transformers) - monkeypatch.setattr(moe_example, "load_lingbot_video_moe_transformer", load_moe) monkeypatch.setattr( moe_example, "load_lingbot_video_model_config", @@ -77,6 +74,13 @@ def from_pretrained(cls, *args, **kwargs): ) 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 @@ -91,7 +95,10 @@ def load_from_huggingface(self, module_path, *, module_name=None, module_class=N base = moe_example.build_pipeline(tmp_path, device="cuda") refiner = moe_example.build_refiner(tmp_path, device="cuda") - assert loaded_transformers == [(tmp_path / "transformer", "cpu"), (tmp_path / "refiner", "cpu")] + 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 ( diff --git a/tools/validation/inspect_lingbot_video_checkpoint.py b/tools/validation/inspect_lingbot_video_checkpoint.py index 808ea66..79a856a 100644 --- a/tools/validation/inspect_lingbot_video_checkpoint.py +++ b/tools/validation/inspect_lingbot_video_checkpoint.py @@ -12,11 +12,7 @@ import torch from safetensors import safe_open -from telefuser.pipelines.lingbot_video.loading import ( - checkpoint_key_coverage, - load_lingbot_video_dense_transformer, - load_lingbot_video_moe_transformer, -) +from telefuser.core.module_manager import ModuleManager _DENSE_CONFIG_KEYS = ( "patch_size", @@ -63,6 +59,21 @@ def checkpoint_keys(checkpoint_dir: Path) -> set[str]: 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) @@ -139,8 +150,12 @@ def main() -> None: torch.cuda.memory_allocated(args.device) if torch.cuda.is_available() and "cuda" in args.device else None ) started = time.perf_counter() - loader = load_lingbot_video_dense_transformer if args.variant == "dense" else load_lingbot_video_moe_transformer - model = loader(checkpoint_dir, device=args.device, torch_dtype=_dtype(args.dtype)) + 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, diff --git a/tools/validation/replay_lingbot_video_dense_reference.py b/tools/validation/replay_lingbot_video_dense_reference.py index ecb230b..5f7c033 100644 --- a/tools/validation/replay_lingbot_video_dense_reference.py +++ b/tools/validation/replay_lingbot_video_dense_reference.py @@ -20,7 +20,6 @@ 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.loading import load_lingbot_video_dense_transformer from telefuser.pipelines.lingbot_video.text_encoding import LingBotVideoTextEncodingStage from telefuser.pipelines.lingbot_video.vae import ( LingBotVideoVAEDecodeStage, @@ -200,9 +199,12 @@ def _load_replay_runtime( raise RuntimeError("Dense replay requires diffusers AutoencoderKLWan") from exc device = torch.device("cuda") runtime_config = ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) - transformer = load_lingbot_video_dense_transformer(model_dir / "transformer", device=device) module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) - module_manager.add_module(transformer, name="transformer") + 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 diff --git a/tools/validation/run_lingbot_video_dense_parity.py b/tools/validation/run_lingbot_video_dense_parity.py index 90406b1..5778575 100644 --- a/tools/validation/run_lingbot_video_dense_parity.py +++ b/tools/validation/run_lingbot_video_dense_parity.py @@ -13,7 +13,7 @@ import torch from safetensors.torch import load_model -from telefuser.pipelines.lingbot_video.loading import load_lingbot_video_dense_transformer +from telefuser.core.module_manager import ModuleManager def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: @@ -65,7 +65,12 @@ def main() -> None: 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)) - telefuser = load_lingbot_video_dense_transformer(args.transformer_dir) + 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) diff --git a/tools/validation/run_lingbot_video_moe_parity.py b/tools/validation/run_lingbot_video_moe_parity.py index 44f08e3..f4bd0fd 100644 --- a/tools/validation/run_lingbot_video_moe_parity.py +++ b/tools/validation/run_lingbot_video_moe_parity.py @@ -13,7 +13,7 @@ import torch from safetensors.torch import load_file -from telefuser.pipelines.lingbot_video.loading import load_lingbot_video_moe_transformer +from telefuser.core.module_manager import ModuleManager def _metrics(reference: torch.Tensor, candidate: torch.Tensor) -> dict[str, float]: @@ -83,7 +83,12 @@ def main() -> None: del upstream torch.cuda.empty_cache() - telefuser = load_lingbot_video_moe_transformer(directory, device=device, torch_dtype=torch.bfloat16) + 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) diff --git a/tools/validation/run_lingbot_video_refiner_core_parity.py b/tools/validation/run_lingbot_video_refiner_core_parity.py index af50f4d..d7b2269 100644 --- a/tools/validation/run_lingbot_video_refiner_core_parity.py +++ b/tools/validation/run_lingbot_video_refiner_core_parity.py @@ -12,7 +12,6 @@ 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.loading import load_lingbot_video_moe_transformer from telefuser.pipelines.lingbot_video.refiner import ( LingBotVideoRefinerStage, compute_refiner_sigmas, @@ -129,10 +128,13 @@ def main() -> None: gc.collect() torch.cuda.empty_cache() - local = load_lingbot_video_moe_transformer(DIRECTORY, device=DEVICE, torch_dtype=torch.bfloat16) + 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") - module_manager = ModuleManager(device="cuda", torch_dtype=torch.bfloat16) - module_manager.add_module(local, name="transformer") stage = LingBotVideoRefinerStage( denoising_stage=LingBotVideoDenoisingStage( "refiner", module_manager, ModelRuntimeConfig(device_type="cuda", torch_dtype=torch.bfloat16) From b98971fd8615fde63a1840b7781f4a9908128fc3 Mon Sep 17 00:00:00 2001 From: lzx1413 Date: Thu, 23 Jul 2026 02:33:59 +0000 Subject: [PATCH 9/9] test(lingbot-video): require pipeline contract export Align the LingBot example contract test with the service loader's canonical PIPELINE_CONTRACT export and remove the obsolete CONTRACT alias expectation. Verification: ruff format tests/unit/pipelines/lingbot_video/test_service_contract.py; ruff check tests/unit/pipelines/lingbot_video/test_service_contract.py; pytest tests/unit/pipelines/lingbot_video/test_service_contract.py -q. --- tests/unit/pipelines/lingbot_video/test_service_contract.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/tests/unit/pipelines/lingbot_video/test_service_contract.py b/tests/unit/pipelines/lingbot_video/test_service_contract.py index b24a9cf..4a67157 100644 --- a/tests/unit/pipelines/lingbot_video/test_service_contract.py +++ b/tests/unit/pipelines/lingbot_video/test_service_contract.py @@ -37,9 +37,8 @@ def _load_example(variant: str) -> ModuleType: def test_examples_expose_cli_and_service_entrypoints(variant: str) -> None: module = _load_example(variant) - for name in ("PPL_CONFIG", "CONTRACT", "get_pipeline", "run", "run_with_file"): + for name in ("PPL_CONFIG", "PIPELINE_CONTRACT", "get_pipeline", "run", "run_with_file"): assert hasattr(module, name) - assert module.PIPELINE_CONTRACT is module.CONTRACT assert module.PPL_CONFIG["variant"] == variant prompt_payload = json.loads(module.PPL_CONFIG["prompt"])