Add Echo-WM Base and Flash modular pipelines

[liyaowei-stu]

What does this PR do?

This PR adds native Diffusers support for the Echo-WM Base and Flash checkpoints as modular pipelines for image- and action-conditioned video generation with synchronized audio.

The implementation includes:

• EchoWMTransformer3DModel, including Echo-WM's action conditioning, sigma handling, velocity-to-x0 conversion, and causal cache behavior.
• Separate EchoWMModularPipeline and EchoWMFlashModularPipeline blocksets for the Base and Flash inference workflows.
• Video and audio decoding through the existing LTX-2 VAE and vocoder components.
• A reusable apply_action_overlay post-processing utility for rendering WASD/IJKL controls on generated PIL frames.
• A conversion script for converting the original Echo-WM checkpoints to Diffusers format.
• Documentation and examples for loading the models, running inference, and saving muxed audio-video output.
• Model and modular pipeline tests using the standard Diffusers tester mixins.
• Dedicated tiny Base and Flash checkpoints for pipeline tests.

Model repositories:

• Base: https://huggingface.co/Echo-Team/Echo-WM-Base-Diffusers
• Flash: https://huggingface.co/Echo-Team/Echo-WM-Flash-Diffusers
• Tiny Base: https://huggingface.co/Echo-Team/tiny-echo-wm-base-diffusers
• Tiny Flash: https://huggingface.co/Echo-Team/tiny-echo-wm-flash-diffusers

The PR also fixes an LTX-2 VAE tiled-decoding edge case where a final tile could lie entirely inside the preceding overlap. Such a tile adds no output coverage and may be too small for reflection padding. The updated implementation skips overlap-only tail tiles and preserves the final decoded tile when assembling the output. A regression test covers both portrait and landscape layouts.

No new mandatory dependency is introduced. PyAV and Accelerate remain optional dependencies for media encoding/image CRF processing and automatic CPU offloading respectively.

Examples

Echo-WM Base

import torch

from diffusers import ComponentsManager, ModularPipeline
from diffusers.utils import encode_video, load_image


components_manager = ComponentsManager()
components_manager.enable_auto_cpu_offload(device="cuda")

pipe = ModularPipeline.from_pretrained(
    "Echo-Team/Echo-WM-Base-Diffusers",
    components_manager=components_manager,
)
pipe.load_components(dtype=torch.bfloat16)

image = load_image(
    "https://raw.githubusercontent.com/jd-opensource/JoyAI-Echo/main/"
    "echo_wm/examples/wm_cases/0010/input.png"
)

result = pipe(
    image=image,
    prompt=(
        "Environment: A still teal pool fills a fantasy canyon of pale limestone. "
        "A monumental natural arch frames a distant white palace. "
        "Sounds: Water laps against the shore and birds echo between the cliffs."
    ),
    negative_prompt="blurry, jittery, distorted, text, watermark, user interface",
    action="w-60,a-60,w-60,d-60",
    num_frames=241,
    num_inference_steps=30,
    generator=torch.Generator(device="cuda").manual_seed(34),
    output=["videos", "audio"],
)

encode_video(
    video=result["videos"][0],
    audio=result["audio"][0].float().cpu(),
    fps=24,
    audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
    output_path="echo_wm_base.mp4",
)

Echo-WM Flash

import torch

from diffusers import ComponentsManager, ModularPipeline
from diffusers.utils import encode_video, load_image


components_manager = ComponentsManager()
components_manager.enable_auto_cpu_offload(device="cuda")

pipe = ModularPipeline.from_pretrained(
    "Echo-Team/Echo-WM-Flash-Diffusers",
    components_manager=components_manager,
)
pipe.load_components(dtype=torch.bfloat16)

image = load_image(
    "https://raw.githubusercontent.com/jd-opensource/JoyAI-Echo/main/"
    "echo_wm/examples/wm_causal_cases/0079/input.jpg"
)

result = pipe(
    image=image,
    prompt=(
        "An enchanted crystal cave with massive purple, teal, and pink crystal formations, "
        "bioluminescent fungi, floating light motes, and a crystalline cave monster."
    ),
    action="l-96,l-96,l-96,l-96",
    num_frames=241,
    generator=torch.Generator(device="cuda").manual_seed(42),
    output=["videos", "audio"],
)

encode_video(
    video=result["videos"][0],
    audio=result["audio"][0].float().cpu(),
    fps=24,
    audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
    output_path="echo_wm_flash.mp4",
)

Flash uses its fixed distilled schedule and therefore does not accept negative_prompt or num_inference_steps.

GPU validation and upstream impact

The implementation was validated on an NVIDIA CUDA GPU with the published full-size Base and Flash checkpoints.

The validation covered:

• Loading both checkpoints through ModularPipeline.from_pretrained() and load_components().
• Automatic CPU offloading through ComponentsManager.
• 241-frame Base inference with 30 denoising steps.
• 241-frame Flash inference with its fixed distilled autoregressive schedule.
• Action-conditioned camera trajectories.
• Video and synchronized audio decoding.
• Direct muxing of video and audio into MP4 files with encode_video.
• Optional action-overlay post-processing.

The Echo-WM-specific transformer, denoising, action conditioning, causal cache handling, and post-processing logic are isolated under the Echo-WM model and modular pipeline modules.

This integration does not modify the existing LTX-2 denoiser, vocoder, or ComponentsManager behavior. The only shared runtime change is a focused fix in AutoencoderKLLTX2Video for overlap-only tail tiles during tiled decoding. That fix is model-independent and covered by a dedicated regression test. The remaining changes to existing Diffusers files are limited to class exports, modular pipeline registration, generated dummy objects, documentation registration, and tests.

Before submitting

• Did you use an AI agent (Claude Code, Codex, Cursor, etc.) to help with this PR? If so:
  • Did you read the Coding with AI agents guide?
  • Did you run the self-review skill on the diff?
  • Did you share the final self-review notes in the PR description or a comment?
• Did you read the contributor guideline?
• Did you read our philosophy doc?
• Was this discussed/approved via a GitHub issue or the forum?
• Did you make sure to update the documentation with your changes?
• Did you write any new necessary tests?
• Are you the author (or part of the team) of the model/pipeline?

Validation

• Echo-WM modular pipeline tests: 64 passed, 10 skipped
• Echo-WM loading tests: 18 passed
• make fix-copies
• Ruff lint and formatting checks
• Base and Flash end-to-end inference with muxed video and audio output

AI-assisted development and self-review

This contribution was developed with assistance from Codex.

Final self-review

Verdict: READY

Blocking issues

None found.

Non-blocking notes

1. The tiny test checkpoints currently live under Echo-Team. They can be migrated to hf-internal-testing during review if preferred.
2. The LTX-2 VAE tiled-decoding change modifies shared runtime code, but it is limited to skipping overlap-only tail tiles and retaining the final useful tile. The change is covered by a focused regression test.
3. Echo-WM keeps model-specific denoising, action conditioning, causal cache handling, and post-processing inside its own model and modular pipeline modules. No changes are made to components_manager.py, the LTX-2 vocoder, or the LTX-2 denoiser.

Dead-code review

No likely dead code was identified in the Echo-WM pipeline call paths reviewed for Base and Flash inference.

Who can review?

cc @yiyixuxu @dg845