AI/diffusers
1
0
Fork 0
mirror of https://github.com/huggingface/diffusers.git synced 2026-10-23 09:09:01 +03:00
Code Activity

[modular] Echo team joy future academy jd echo by @Gelercatty (#14941)

Browse Source 
* Add Echo modular pipeline

* Address Echo modular pipeline review feedback

* Inline Echo input list normalization

* Align Echo latent boundaries and batch expansion

* Polish Echo documentation per review

* Load official Echo checkpoint and separate block tests

* consolidate testing.

* formatting

---------

Co-authored-by: Yanwen Ma <gelercat117@gmail.com>
This commit is contained in:
Sayak Paul
2026-10-05 14:04:27 +05:30
committed by GitHub
parent 9aababd755
commit cff9dafbbf
18 changed files with 2953 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
docs/source/en/_toctree.yml
View File

@@ -575,6 +575,8 @@
title: DreamLite
- local: api/pipelines/easyanimate
title: EasyAnimate
- local: api/pipelines/echo
title: Echo
- local: api/pipelines/ernie_image
title: ERNIE-Image
- local: api/pipelines/flux

98
docs/source/en/api/pipelines/echo.md Normal file
View File

@@ -0,0 +1,98 @@
<!-- Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License. -->
# Echo
[Echo](https://github.com/jd-opensource/JoyAI-Echo) is a long-video generation model. It supports an optional clean
first frame, ordered image/audio memory slots, and a stochastic few-step Distribution Matching Distillation (DMD)
sampler.
The pipeline generates synchronized video and audio.
Echo is implemented as a [`ModularPipeline`] so its text encoding, memory conditioning, stochastic DMD denoising, and
decoding blocks can be run as a complete workflow or composed independently.
## Inference
Load the official [jdopensource/JoyAI-Echo](https://huggingface.co/jdopensource/JoyAI-Echo) checkpoint directly.
It uses the text encoder and tokenizer from [google/gemma-3-12b-it](https://huggingface.co/google/gemma-3-12b-it).
The Gemma repository is gated, so accept its license and authenticate with Hugging Face before loading the pipeline.
The released model uses 241 frames in its long-video example. The video RoPE coordinates remain at the training rate
of 24 fps, independently of the output container rate.
```py
import torch
import torchaudio
from PIL import Image
from diffusers import ComponentsManager, ModularPipeline
from diffusers.utils import encode_video
manager = ComponentsManager()
pipe = ModularPipeline.from_pretrained("jdopensource/JoyAI-Echo", components_manager=manager)
pipe.load_components(dtype={"default": torch.bfloat16, "audio_vae": torch.float32})
manager.enable_auto_cpu_offload(device="cuda")
pipe.vae.enable_tiling()
first_frame = Image.open("first_frame.png").convert("RGB")
memory_images = [Image.open(path).convert("RGB") for path in ["memory_0.png", "memory_1.png"]]
memory_audio_with_rates = [torchaudio.load(path) for path in ["memory_0.wav", "memory_1.wav"]]
memory_audio = [waveform for waveform, _ in memory_audio_with_rates]
memory_audio_rates = [sample_rate for _, sample_rate in memory_audio_with_rates]
output = pipe(
prompt="A cinematic dialogue scene in a quiet cafe.",
image=first_frame,
memory_images=memory_images,
memory_audio_waveforms=memory_audio,
memory_audio_sample_rates=memory_audio_rates,
width=1280,
height=736,
num_frames=241,
frame_rate=25.0,
model_frame_rate=24.0,
generator=torch.Generator(device="cuda").manual_seed(42),
output_type="np",
output=["videos", "audio"],
)
encode_video(
output["videos"][0],
fps=25,
audio=output["audio"][0].float().cpu(),
audio_sample_rate=pipe.vocoder.config.output_sampling_rate,
output_path="echo.mp4",
)
```
The default DMD sigma schedule is the released eight-step schedule. It predicts `x0` at every step and re-noises with
fresh Gaussian noise at the next sigma, so a seeded `torch.Generator` controls both the initial noise and all
intermediate re-noising.
Raw audio-memory encoding requires `torchaudio`. For reference parity, keep `audio_vae` in FP32 as shown above.
Modular workflows can cache the VAE encoder's normalized, unpacked tensors. Video latents have shape
`(batch, channels, frames, height, width)` and audio latents have shape `(batch, channels, time, mel_bins)`.
The core `denoise` block packs those tensors, expands conditioning for `num_videos_per_prompt`, and unpacks
denoised outputs back to the same VAE form. Pass initial `latents` and `audio_latents` in that unpacked form too.
Decoders take normalized VAE tensors and denormalize immediately before
decoding. `output=["latents", "audio_latents"]` returns normalized VAE tensors. With `output_type="latent"` and
`output=["videos", "audio"]`, you get the denormalized VAE tensors without decoding.
## EchoModularPipeline
[[autodoc]] EchoModularPipeline
## EchoBlocks
[[autodoc]] EchoBlocks

179
scripts/convert_echo_to_diffusers.py Normal file
View File

@@ -0,0 +1,179 @@
#!/usr/bin/env python3
"""Convert the Echo BF16 release checkpoint into a Diffusers Modular Pipeline repository."""
import argparse
import gc
import json
from pathlib import Path
import torch
from convert_ltx2_to_diffusers import (
convert_ltx2_audio_vae,
convert_ltx2_connectors,
convert_ltx2_transformer,
convert_ltx2_video_vae,
convert_ltx2_vocoder,
get_model_state_dict_from_combined_ckpt,
)
from safetensors.torch import load_file
from diffusers import __version__
def resolve_checkpoint(path: Path) -> Path:
path = path.expanduser().resolve()
if path.is_file():
return path
manifest_path = path / "checkpoint.json"
if not manifest_path.is_file():
raise FileNotFoundError(f"Checkpoint manifest not found: {manifest_path}")
manifest = json.loads(manifest_path.read_text(encoding="utf-8"))
precision = str(manifest.get("precision", "")).lower()
if precision not in {"bf16", "bfloat16"}:
raise ValueError(
"Echo conversion supports the BF16 release only. The FP8 and FP4 releases use custom packed "
f"kernels and cannot be represented by this converter, got precision={precision!r}."
)
model_filename = manifest.get("files", {}).get("model")
if not model_filename:
raise ValueError(f"Checkpoint manifest has no `files.model`: {manifest_path}")
return path / model_filename
def save_component(model, output_path: Path, name: str, max_shard_size: str) -> None:
model.to(torch.bfloat16).save_pretrained(
output_path / name,
safe_serialization=True,
max_shard_size=max_shard_size,
)
def component_entry(library: str, class_name: str, repo: str, subfolder: str | None):
return [
None,
None,
{
"type_hint": [library, class_name],
"pretrained_model_name_or_path": repo,
"subfolder": subfolder,
"variant": None,
"revision": None,
},
]
def write_modular_index(output_path: Path, base_model: str, repo_id: str | None) -> None:
echo_repo = repo_id or str(output_path)
index = {
"_class_name": "EchoModularPipeline",
"_diffusers_version": __version__,
"_blocks_class_name": "EchoBlocks",
"text_encoder": component_entry("transformers", "Gemma3ForConditionalGeneration", base_model, None),
"tokenizer": component_entry("transformers", "GemmaTokenizerFast", base_model, None),
"connectors": component_entry("ltx2", "LTX2TextConnectors", echo_repo, "connectors"),
"vae": component_entry("diffusers", "AutoencoderKLLTX2Video", echo_repo, "vae"),
"audio_vae": component_entry("diffusers", "AutoencoderKLLTX2Audio", echo_repo, "audio_vae"),
"transformer": component_entry("diffusers", "LTX2VideoTransformer3DModel", echo_repo, "transformer"),
"vocoder": component_entry("ltx2", "LTX2VocoderWithBWE", echo_repo, "vocoder"),
}
(output_path / "modular_model_index.json").write_text(json.dumps(index, indent=2) + "\n", encoding="utf-8")
def convert(
checkpoint: Path,
output_path: Path,
base_model: str,
repo_id: str | None,
max_shard_size: str,
) -> None:
source = resolve_checkpoint(checkpoint)
output_path = output_path.expanduser().resolve()
output_path.mkdir(parents=True, exist_ok=True)
print(f"Loading Echo checkpoint from {source}", flush=True)
combined = load_file(str(source), device="cpu")
component_states = {
"dit": get_model_state_dict_from_combined_ckpt(combined, "model.diffusion_model"),
"vae": get_model_state_dict_from_combined_ckpt(combined, "vae"),
"audio_vae": get_model_state_dict_from_combined_ckpt(combined, "audio_vae"),
"vocoder": get_model_state_dict_from_combined_ckpt(combined, "vocoder"),
}
del combined
missing = [name for name, state in component_states.items() if not state]
if missing:
raise ValueError(f"Echo checkpoint is missing components: {missing}")
tensor_counts = {}
dit_state = component_states.pop("dit")
transformer = convert_ltx2_transformer(dict(dit_state), version="2.3")
tensor_counts["transformer"] = len(transformer.state_dict())
save_component(transformer, output_path, "transformer", max_shard_size)
del transformer
gc.collect()
connectors = convert_ltx2_connectors(dict(dit_state), version="2.3")
tensor_counts["connectors"] = len(connectors.state_dict())
save_component(connectors, output_path, "connectors", max_shard_size)
del connectors, dit_state
gc.collect()
vae = convert_ltx2_video_vae(component_states.pop("vae"), version="2.3", timestep_conditioning=False)
tensor_counts["vae"] = len(vae.state_dict())
save_component(vae, output_path, "vae", max_shard_size)
del vae
gc.collect()
audio_vae = convert_ltx2_audio_vae(component_states.pop("audio_vae"), version="2.3")
tensor_counts["audio_vae"] = len(audio_vae.state_dict())
save_component(audio_vae, output_path, "audio_vae", max_shard_size)
del audio_vae
gc.collect()
vocoder = convert_ltx2_vocoder(component_states.pop("vocoder"), version="2.3")
tensor_counts["vocoder"] = len(vocoder.state_dict())
save_component(vocoder, output_path, "vocoder", max_shard_size)
del vocoder
gc.collect()
write_modular_index(output_path, base_model=base_model, repo_id=repo_id)
report = {
"schema": "echo.diffusers.conversion.v1",
"source": str(source),
"precision": "bfloat16",
"tensor_counts": tensor_counts,
"base_model": base_model,
"model_repo": repo_id,
"vocoder_output_sample_rate": 48000,
}
(output_path / "conversion_report.json").write_text(json.dumps(report, indent=2) + "\n", encoding="utf-8")
print(f"Echo conversion complete: {output_path}", flush=True)
def parse_args() -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument("--checkpoint", type=Path, required=True, help="Echo BF16 checkpoint or release folder.")
parser.add_argument("--output-path", type=Path, required=True, help="Destination Diffusers model directory.")
parser.add_argument(
"--base-model",
default="google/gemma-3-12b-it",
help="Gemma model repo/path containing the text encoder and tokenizer at its root.",
)
parser.add_argument(
"--repo-id",
default=None,
help="Future Hub repo id for the converted Echo components. Defaults to the local output path.",
)
parser.add_argument("--max-shard-size", default="5GB")
return parser.parse_args()
if __name__ == "__main__":
args = parse_args()
convert(
checkpoint=args.checkpoint,
output_path=args.output_path,
base_model=args.base_model,
repo_id=args.repo_id,
max_shard_size=args.max_shard_size,
)

4
src/diffusers/__init__.py
View File

@@ -521,6 +521,8 @@ else:
"Cosmos3DistilledModularPipeline",
"Cosmos3OmniBlocks",
"Cosmos3OmniModularPipeline",
"EchoBlocks",
"EchoModularPipeline",
"ErnieImageAutoBlocks",
"ErnieImageModularPipeline",
"Flux2AutoBlocks",
@@ -1380,6 +1382,8 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
Cosmos3DistilledModularPipeline,
Cosmos3OmniBlocks,
Cosmos3OmniModularPipeline,
EchoBlocks,
EchoModularPipeline,
ErnieImageAutoBlocks,
ErnieImageModularPipeline,
Flux2AutoBlocks,

5
src/diffusers/modular_pipelines/__init__.py
View File

@@ -115,6 +115,10 @@ else:
"ErnieImageAutoBlocks",
"ErnieImageModularPipeline",
]
_import_structure["echo"] = [
"EchoBlocks",
"EchoModularPipeline",
]
_import_structure["hunyuan_video1_5"] = [
"HunyuanVideo15AutoBlocks",
"HunyuanVideo15ModularPipeline",
@@ -158,6 +162,7 @@ if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
Cosmos3OmniBlocks,
Cosmos3OmniModularPipeline,
)
from .echo import EchoBlocks, EchoModularPipeline
from .ernie_image import ErnieImageAutoBlocks, ErnieImageModularPipeline
from .flux import FluxAutoBlocks, FluxKontextAutoBlocks, FluxKontextModularPipeline, FluxModularPipeline
from .flux2 import (

47
src/diffusers/modular_pipelines/echo/__init__.py Normal file
View File

@@ -0,0 +1,47 @@
from typing import TYPE_CHECKING
from ...utils import (
DIFFUSERS_SLOW_IMPORT,
OptionalDependencyNotAvailable,
_LazyModule,
get_objects_from_module,
is_torch_available,
is_transformers_available,
)
_dummy_objects = {}
_import_structure = {}
try:
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from ...utils import dummy_torch_and_transformers_objects # noqa F403
_dummy_objects.update(get_objects_from_module(dummy_torch_and_transformers_objects))
else:
_import_structure["modular_blocks_echo"] = ["EchoBlocks"]
_import_structure["modular_pipeline"] = ["EchoModularPipeline"]
if TYPE_CHECKING or DIFFUSERS_SLOW_IMPORT:
try:
if not (is_transformers_available() and is_torch_available()):
raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
from ...utils.dummy_torch_and_transformers_objects import * # noqa F403
else:
from .modular_blocks_echo import EchoBlocks
from .modular_pipeline import EchoModularPipeline
else:
import sys
sys.modules[__name__] = _LazyModule(
__name__,
globals()["__file__"],
_import_structure,
module_spec=__spec__,
)
for name, value in _dummy_objects.items():
setattr(sys.modules[__name__], name, value)

501
src/diffusers/modular_pipelines/echo/before_denoise.py Normal file
View File

@@ -0,0 +1,501 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import torch
from ...models import LTX2VideoTransformer3DModel
from ...utils.torch_utils import randn_tensor
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
# Copied from diffusers.modular_pipelines.ltx2.before_denoise._pack_latents
def _pack_latents(latents: torch.Tensor, patch_size: int = 1, patch_size_t: int = 1) -> torch.Tensor:
batch_size, num_channels, num_frames, height, width = latents.shape
post_patch_num_frames = num_frames // patch_size_t
post_patch_height = height // patch_size
post_patch_width = width // patch_size
latents = latents.reshape(
batch_size,
-1,
post_patch_num_frames,
patch_size_t,
post_patch_height,
patch_size,
post_patch_width,
patch_size,
)
latents = latents.permute(0, 2, 4, 6, 1, 3, 5, 7).flatten(4, 7).flatten(1, 3)
return latents
def _pack_audio_latents(latents: torch.Tensor) -> torch.Tensor:
# Echo packs the full mel axis into each audio token: [B, C, L, M] -> [B, L, C * M].
return latents.transpose(1, 2).flatten(2, 3)
def _video_memory_coords(
transformer,
slots: int,
latent_height: int,
latent_width: int,
frame_rate: float,
position_offset: float,
slot_stride: float,
device: torch.device,
) -> torch.Tensor:
base = transformer.rope.prepare_video_coords(1, 1, latent_height, latent_width, device, fps=frame_rate)
coords = []
for slot_index in range(slots):
slot_coords = base.clone()
center = position_offset + slot_index * slot_stride
midpoint = (slot_coords[:, 0, :1, 0] + slot_coords[:, 0, :1, 1]) * 0.5
slot_coords[:, 0] += (center - midpoint).view(1, 1, 1)
coords.append(slot_coords)
return torch.cat(coords, dim=2)
def _audio_memory_coords(
transformer,
lengths: list[int],
position_offset: float,
slot_stride: float,
device: torch.device,
) -> torch.Tensor:
coords = []
for slot_index, length in enumerate(lengths):
slot_coords = transformer.audio_rope.prepare_audio_coords(1, length, device)
center = position_offset + slot_index * slot_stride
midpoint = (slot_coords[:, 0, :1, 0] + slot_coords[:, 0, -1:, 1]) * 0.5
slot_coords[:, 0] += (center - midpoint).view(1, 1, 1)
coords.append(slot_coords)
return torch.cat(coords, dim=2)
class EchoPrepareConditioningStep(ModularPipelineBlocks):
"""Pack VAE latents and build the transformer-specific memory coordinates."""
model_name = "echo"
@property
def description(self) -> str:
return "Packs Echo VAE latents and assigns each memory slot its transformer RoPE coordinates."
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"first_frame_latents",
type_hint=torch.Tensor,
required=False,
description="Normalized VAE latents for the optional clean first frame.",
),
InputParam(
"memory_video_latents",
type_hint=list,
required=False,
description="Normalized VAE latents for each ordered memory image.",
),
InputParam(
"memory_audio_latents",
type_hint=list,
required=False,
description="Normalized audio VAE latents of shape (B, C, L, M) for each memory slot.",
),
InputParam(
"model_frame_rate",
type_hint=float,
default=24.0,
description="Training-time frame rate used for Echo video RoPE coordinates.",
),
InputParam(
"memory_position_offset",
type_hint=float,
default=500.0,
description="Temporal center assigned to the first memory slot.",
),
InputParam(
"memory_position_slot_stride",
type_hint=float,
default=50.0,
description="Temporal distance between consecutive memory-slot centers.",
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("first_frame_tokens", type_hint=torch.Tensor, description="Packed first-frame tokens."),
OutputParam("memory_video_tokens", type_hint=torch.Tensor, description="Packed image-memory tokens."),
OutputParam(
"memory_video_coords", type_hint=torch.Tensor, description="RoPE coordinates for image-memory tokens."
),
OutputParam("memory_audio_tokens", type_hint=torch.Tensor, description="Packed audio-memory tokens."),
OutputParam(
"memory_audio_coords", type_hint=torch.Tensor, description="RoPE coordinates for audio-memory tokens."
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
transformer = components.transformer
device = components._execution_device
dtype = transformer.dtype
patch_size = components.transformer_spatial_patch_size
patch_size_t = components.transformer_temporal_patch_size
first_frame_tokens = None
if block_state.first_frame_latents is not None:
first_frame_tokens = _pack_latents(block_state.first_frame_latents, patch_size, patch_size_t)
memory_video_latents = block_state.memory_video_latents or []
video_slots = [_pack_latents(latents, patch_size, patch_size_t) for latents in memory_video_latents]
memory_video_tokens = torch.cat(video_slots, dim=1) if video_slots else None
memory_video_coords = None
if video_slots:
latent_height, latent_width = memory_video_latents[0].shape[-2:]
memory_video_coords = _video_memory_coords(
transformer,
len(video_slots),
latent_height,
latent_width,
block_state.model_frame_rate,
block_state.memory_position_offset,
block_state.memory_position_slot_stride,
device,
)
memory_audio_latents = block_state.memory_audio_latents or []
template = next((value for value in memory_audio_latents if value is not None), None)
memory_audio_tokens = None
memory_audio_coords = None
if template is not None:
aligned_audio = [
_pack_audio_latents(value if value is not None else torch.zeros_like(template))
for value in memory_audio_latents
]
lengths = [value.shape[1] for value in aligned_audio]
memory_audio_tokens = torch.cat(aligned_audio, dim=1)
memory_audio_coords = _audio_memory_coords(
transformer,
lengths,
block_state.memory_position_offset,
block_state.memory_position_slot_stride,
device,
)
block_state.first_frame_tokens = (
None if first_frame_tokens is None else first_frame_tokens.to(device=device, dtype=dtype)
)
block_state.memory_video_tokens = (
None if memory_video_tokens is None else memory_video_tokens.to(device=device, dtype=dtype)
)
block_state.memory_video_coords = memory_video_coords
block_state.memory_audio_tokens = (
None if memory_audio_tokens is None else memory_audio_tokens.to(device=device, dtype=dtype)
)
block_state.memory_audio_coords = memory_audio_coords
self.set_block_state(state, block_state)
return components, state
class EchoInputsStep(ModularPipelineBlocks):
"""Expand per-prompt text and shared or per-prompt memory conditioning once before denoising."""
model_name = "echo"
@property
def description(self) -> str:
return "Expands text, first-frame, and memory conditioning to the requested number of videos per prompt."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam("num_videos_per_prompt", type_hint=int, default=1, description="Number of videos per prompt."),
InputParam(
"connector_prompt_embeds",
type_hint=torch.Tensor,
required=True,
description="Per-prompt video-branch text conditioning.",
),
InputParam(
"connector_audio_prompt_embeds",
type_hint=torch.Tensor,
required=True,
description="Per-prompt audio-branch text conditioning.",
),
InputParam(
"connector_attention_mask",
type_hint=torch.Tensor,
required=True,
description="Per-prompt binary text attention mask.",
),
InputParam("first_frame_tokens", type_hint=torch.Tensor, description="Packed first-frame tokens."),
InputParam("memory_video_tokens", type_hint=torch.Tensor, description="Packed image-memory tokens."),
InputParam("memory_audio_tokens", type_hint=torch.Tensor, description="Packed audio-memory tokens."),
InputParam("memory_video_coords", type_hint=torch.Tensor, description="Image-memory RoPE coordinates."),
InputParam("memory_audio_coords", type_hint=torch.Tensor, description="Audio-memory RoPE coordinates."),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(param.name, type_hint=torch.Tensor, description=f"Batch-expanded {param.description}")
for param in self.inputs
if param.name != "num_videos_per_prompt"
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
prompt_batch_size = block_state.connector_prompt_embeds.shape[0]
num_videos_per_prompt = block_state.num_videos_per_prompt
if num_videos_per_prompt < 1:
raise ValueError("`num_videos_per_prompt` must be a positive integer.")
for param in self.intermediate_outputs:
value = getattr(block_state, param.name)
if value is None:
continue
if value.shape[0] == 1:
value = value.repeat_interleave(prompt_batch_size, dim=0)
elif value.shape[0] != prompt_batch_size:
raise ValueError(
f"`{param.name}` must have batch size 1 or {prompt_batch_size}, got {value.shape[0]}."
)
setattr(block_state, param.name, value.repeat_interleave(num_videos_per_prompt, dim=0))
self.set_block_state(state, block_state)
return components, state
class EchoPrepareLatentsStep(ModularPipelineBlocks):
"""Prepare the target video/audio sequences and their fixed Echo coordinates."""
model_name = "echo"
@property
def description(self) -> str:
return (
"Creates packed video and audio noise, inserts the optional clean first frame, and prepares fixed target "
"RoPE coordinates for Echo's stochastic DMD loop."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("height", default=512),
InputParam.template("width", default=704),
InputParam(
"num_frames",
type_hint=int,
default=241,
description="Number of generated pixel frames; must be `1 + k * vae_temporal_compression_ratio`.",
),
InputParam("frame_rate", type_hint=float, default=25.0, description="Frame rate of the generated video."),
InputParam(
"model_frame_rate",
type_hint=float,
default=24.0,
description="Training-time frame rate used for video RoPE coordinates.",
),
InputParam(
"latents",
type_hint=torch.Tensor,
description="Optional initial video noise in VAE form (B, C, F, H, W).",
),
InputParam(
"audio_latents",
type_hint=torch.Tensor,
default=None,
description="Optional initial audio noise in VAE form (B, C, L, M).",
),
InputParam.template("generator"),
InputParam(
"connector_prompt_embeds",
type_hint=torch.Tensor,
required=True,
description="Positive video-branch text conditioning.",
),
InputParam("first_frame_tokens", type_hint=torch.Tensor, required=False),
InputParam("memory_video_tokens", type_hint=torch.Tensor, required=False),
InputParam("memory_video_coords", type_hint=torch.Tensor, required=False),
InputParam("memory_audio_tokens", type_hint=torch.Tensor, required=False),
InputParam("memory_audio_coords", type_hint=torch.Tensor, required=False),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam("latents", type_hint=torch.Tensor, description="Packed initial noisy video latents."),
OutputParam("audio_latents", type_hint=torch.Tensor, description="Packed initial noisy audio latents."),
OutputParam(
"video_coords", type_hint=torch.Tensor, description="RoPE coordinates for target video tokens."
),
OutputParam(
"audio_coords", type_hint=torch.Tensor, description="RoPE coordinates for target audio tokens."
),
OutputParam("latent_num_frames", type_hint=int, description="Number of target video latent frames."),
OutputParam("latent_height", type_hint=int, description="Target video latent height."),
OutputParam("latent_width", type_hint=int, description="Target video latent width."),
OutputParam("audio_num_frames", type_hint=int, description="Number of target audio latent frames."),
OutputParam(
"audio_latent_mel_bins", type_hint=int, description="Number of mel bins in audio VAE latents."
),
OutputParam(
"first_frame_token_count", type_hint=int, description="Number of clean first-frame video tokens."
),
OutputParam(
"memory_video_token_count", type_hint=int, description="Number of prepended image-memory tokens."
),
OutputParam(
"memory_audio_token_count", type_hint=int, description="Number of prepended audio-memory tokens."
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = components._execution_device
dtype = components.transformer.dtype
batch_size = block_state.connector_prompt_embeds.shape[0]
if block_state.num_frames < 1 or (block_state.num_frames - 1) % components.vae_temporal_compression_ratio:
raise ValueError(
"`num_frames` must be `1 + k * vae_temporal_compression_ratio` for Echo, got "
f"{block_state.num_frames}."
)
if block_state.height % components.vae_spatial_compression_ratio or block_state.width % (
components.vae_spatial_compression_ratio
):
raise ValueError(
"`height` and `width` must be divisible by the video VAE spatial compression ratio "
f"({components.vae_spatial_compression_ratio})."
)
latent_num_frames = 1 + (block_state.num_frames - 1) // components.vae_temporal_compression_ratio
latent_height = block_state.height // components.vae_spatial_compression_ratio
latent_width = block_state.width // components.vae_spatial_compression_ratio
video_token_count = (
(latent_num_frames // components.transformer_temporal_patch_size)
* (latent_height // components.transformer_spatial_patch_size)
* (latent_width // components.transformer_spatial_patch_size)
)
if block_state.latents is None:
block_state.latents = randn_tensor(
(batch_size, video_token_count, components.transformer.config.in_channels),
generator=block_state.generator,
device=device,
dtype=dtype,
)
else:
latents = block_state.latents.to(device=device, dtype=dtype)
if latents.ndim != 5:
raise ValueError("Echo `latents` must be in VAE form (B, C, F, H, W).")
latents = _pack_latents(
latents,
components.transformer_spatial_patch_size,
components.transformer_temporal_patch_size,
)
if latents.shape != (batch_size, video_token_count, components.transformer.config.in_channels):
raise ValueError(
"Unexpected Echo video latent shape: expected "
f"{(batch_size, video_token_count, components.transformer.config.in_channels)}, got "
f"{tuple(latents.shape)}."
)
block_state.latents = latents
duration = block_state.num_frames / block_state.frame_rate
audio_frames_per_second = (
components.audio_sampling_rate
/ components.audio_hop_length
/ components.audio_vae_temporal_compression_ratio
)
audio_num_frames = round(duration * audio_frames_per_second)
if block_state.audio_latents is None:
block_state.audio_latents = randn_tensor(
(batch_size, audio_num_frames, components.transformer.config.audio_in_channels),
generator=block_state.generator,
device=device,
dtype=dtype,
)
else:
audio_latents = block_state.audio_latents.to(device=device, dtype=dtype)
if audio_latents.ndim != 4:
raise ValueError("Echo `audio_latents` must be in VAE form (B, C, L, M).")
audio_latents = _pack_audio_latents(audio_latents)
if audio_latents.shape != (
batch_size,
audio_num_frames,
components.transformer.config.audio_in_channels,
):
raise ValueError(
"Unexpected Echo audio latent shape: expected "
f"{(batch_size, audio_num_frames, components.transformer.config.audio_in_channels)}, got "
f"{tuple(audio_latents.shape)}."
)
block_state.audio_latents = audio_latents
first_frame_tokens = block_state.first_frame_tokens
first_frame_token_count = 0
if first_frame_tokens is not None:
first_frame_token_count = first_frame_tokens.shape[1]
expected_first_frame_tokens = (latent_height // components.transformer_spatial_patch_size) * (
latent_width // components.transformer_spatial_patch_size
)
if first_frame_token_count != expected_first_frame_tokens:
raise ValueError(
"The encoded first frame has the wrong token count: expected "
f"{expected_first_frame_tokens}, got {first_frame_token_count}."
)
block_state.latents = block_state.latents.clone()
block_state.latents[:, :first_frame_token_count] = first_frame_tokens
block_state.video_coords = components.transformer.rope.prepare_video_coords(
batch_size,
latent_num_frames,
latent_height,
latent_width,
device,
fps=block_state.model_frame_rate,
)
block_state.audio_coords = components.transformer.audio_rope.prepare_audio_coords(
batch_size, audio_num_frames, device
)
block_state.latent_num_frames = latent_num_frames
block_state.latent_height = latent_height
block_state.latent_width = latent_width
block_state.audio_num_frames = audio_num_frames
block_state.audio_latent_mel_bins = components.audio_latent_mel_bins
block_state.first_frame_token_count = first_frame_token_count
block_state.memory_video_token_count = (
0 if block_state.memory_video_tokens is None else block_state.memory_video_tokens.shape[1]
)
block_state.memory_audio_token_count = (
0 if block_state.memory_audio_tokens is None else block_state.memory_audio_tokens.shape[1]
)
self.set_block_state(state, block_state)
return components, state

95
src/diffusers/modular_pipelines/echo/constants.py Normal file
View File

@@ -0,0 +1,95 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# Per-channel latent statistics used by the Echo checkpoint. These values provide pipeline property fallbacks when
# the corresponding autoencoder is not loaded.
# fmt: off
ECHO_LATENTS_MEAN = [
-0.077637, -0.062012, 0.073242, 0.089355, 0.002945, 0.002457, 0.013428, -0.054199,
0.022461, 0.037354, -0.025757, 0.055176, -0.050293, 0.017212, 0.009949, 0.042236,
-0.032715, 0.055908, -0.048584, -0.005981, -0.008545, 0.008911, -0.006714, 0.010742,
-0.040039, 0.030762, -0.003494, 0.019775, -0.058105, -0.030029, -0.049316, -0.040771,
0.011536, -0.009338, 0.029175, 0.004547, -0.013245, 0.029541, -0.019897, 0.026123,
-0.064453, 0.018555, -0.100098, -0.032715, 0.05249, 0.125, -0.010498, 0.031494,
0.166016, -0.044434, -0.016357, 0.012451, 0.01239, -0.034424, -0.016724, 0.041992,
-0.003189, -0.005249, -0.040527, -0.018066, -0.058105, 0.020874, 0.034668, 0.002426,
-0.09375, -0.050537, 0.074219, -0.041016, -0.077148, 0.0019, -0.097168, -0.004639,
-0.005859, 0.006744, 0.040771, -0.003418, 0.062012, 0.083008, 0.092773, -0.003708,
0.031494, -0.048828, -0.02063, -0.005127, -0.054932, -0.068359, 0.000687, -0.032471,
-0.018677, 0.010193, 0.005707, -0.042236, 0.109375, 0.013367, -0.041504, -0.009399,
-0.050293, 0.019043, 0.040527, 0.036377, -0.005127, 0.009338, -0.00206, 0.006805,
0.023315, -0.008728, 0.010986, -0.015625, -0.154297, -0.003021, -0.012573, -0.030273,
0.017944, 0.005249, 0.039551, -0.096191, -0.031494, 0.009033, -0.002014, -0.039795,
-0.047119, -0.003555, -0.016479, -0.09082, 0.007996, -0.064941, 0.142578, 0.023315,
]
ECHO_LATENTS_STD = [
0.161133, 0.12793, 0.168945, 0.189453, 0.135742, 0.12207, 0.150391, 0.193359,
0.099121, 0.120117, 0.124023, 0.117676, 0.160156, 0.244141, 0.143555, 0.136719,
0.113281, 0.128906, 0.222656, 0.150391, 0.120605, 0.092773, 0.296875, 0.12793,
0.141602, 0.179688, 0.144531, 0.125, 0.144531, 0.096191, 0.191406, 0.114258,
0.178711, 0.131836, 0.103516, 0.099609, 0.129883, 0.166992, 0.105469, 0.103027,
0.126953, 0.103027, 0.205078, 0.239258, 0.167969, 0.1875, 0.104492, 0.253906,
0.170898, 0.129883, 0.158203, 0.099121, 0.117188, 0.147461, 0.129883, 0.209961,
0.141602, 0.124512, 0.15332, 0.132812, 0.116699, 0.140625, 0.220703, 0.108887,
0.351562, 0.115723, 0.15332, 0.114258, 0.217773, 0.121582, 0.144531, 0.167969,
0.198242, 0.134766, 0.138672, 0.120605, 0.118652, 0.200195, 0.106934, 0.126953,
0.116211, 0.144531, 0.253906, 0.116699, 0.181641, 0.141602, 0.146484, 0.126953,
0.107422, 0.121094, 0.106445, 0.21582, 0.433594, 0.105957, 0.186523, 0.130859,
0.112793, 0.121094, 0.111816, 0.132812, 0.161133, 0.108398, 0.15332, 0.12793,
0.121094, 0.104492, 0.107422, 0.09668, 0.349609, 0.103027, 0.108887, 0.208984,
0.150391, 0.228516, 0.112793, 0.15918, 0.15625, 0.110352, 0.10498, 0.130859,
0.095215, 0.112793, 0.137695, 0.166016, 0.150391, 0.125977, 0.279297, 0.107422,
]
ECHO_AUDIO_LATENTS_MEAN = [
1.539062, 2.84375, 2.359375, 1.992188, 1.539062, 1.296875, 1.085938, 0.871094,
0.628906, 0.404297, 0.197266, -0.102051, -0.431641, -0.734375, -1.054688, -0.652344,
0.4375, -0.679688, -0.429688, -0.296875, -0.112793, 0.035645, 0.145508, 0.211914,
0.294922, 0.388672, 0.462891, 0.574219, 0.660156, 0.894531, 0.490234, 0.707031,
0.753906, 0.412109, 0.365234, 0.149414, -0.007996, -0.122559, -0.193359, -0.310547,
-0.431641, -0.486328, -0.644531, -0.738281, -0.878906, -1.070312, -0.941406, -0.976562,
0.064453, -0.373047, -0.138672, -0.257812, -0.384766, -0.498047, -0.5625, -0.59375,
-0.691406, -0.765625, -0.785156, -0.878906, -0.898438, -0.976562, -1.085938, -0.507812,
1.617188, 1.671875, 1.789062, 1.742188, 1.710938, 1.6875, 1.664062, 1.609375,
1.546875, 1.546875, 1.507812, 1.429688, 1.390625, 1.203125, 1.3125, 0.773438,
-0.242188, -1.960938, -1.328125, -1.203125, -0.992188, -0.796875, -0.703125, -0.574219,
-0.466797, -0.363281, -0.227539, -0.101074, 0.016602, 0.231445, 0.375, 0.15625,
0.960938, 0.523438, 0.589844, 0.3125, 0.164062, 0.067871, 0.006317, -0.093262,
-0.174805, -0.224609, -0.316406, -0.376953, -0.474609, -0.703125, -0.824219, -1.484375,
0.21875, 0.111816, 0.289062, 0.00705, -0.234375, -0.408203, -0.527344, -0.671875,
-0.789062, -0.949219, -1.070312, -1.257812, -1.390625, -1.585938, -1.671875, -1.804688,
]
ECHO_AUDIO_LATENTS_STD = [
1.992188, 2.125, 2.078125, 2.046875, 2.046875, 2.03125, 2.015625, 1.976562,
1.984375, 1.96875, 1.96875, 1.945312, 1.929688, 1.90625, 1.945312, 1.726562,
1.09375, 1.265625, 1.242188, 1.203125, 1.171875, 1.148438, 1.117188, 1.085938,
1.085938, 1.054688, 1.023438, 1.007812, 0.980469, 0.964844, 0.90625, 0.882812,
1.3125, 1.328125, 1.289062, 1.234375, 1.21875, 1.179688, 1.179688, 1.117188,
1.0625, 1.039062, 1.046875, 1.015625, 0.996094, 0.980469, 1.007812, 1.015625,
1.429688, 1.210938, 1.148438, 1.101562, 1.078125, 1.039062, 0.988281, 0.980469,
0.949219, 0.902344, 0.902344, 0.875, 0.867188, 0.847656, 0.820312, 0.914062,
0.769531, 1.109375, 1.109375, 1.09375, 1.078125, 1.070312, 1.039062, 1.007812,
0.960938, 0.914062, 0.871094, 0.824219, 0.792969, 0.777344, 0.773438, 0.753906,
1.320312, 1.382812, 1.296875, 1.296875, 1.257812, 1.242188, 1.234375, 1.203125,
1.179688, 1.171875, 1.179688, 1.125, 1.125, 1.101562, 1.109375, 1.171875,
1.46875, 1.28125, 1.242188, 1.210938, 1.210938, 1.1875, 1.140625, 1.132812,
1.109375, 1.09375, 1.085938, 1.0625, 1.054688, 1.0625, 1.101562, 1.351562,
1.203125, 1.5625, 1.460938, 1.445312, 1.398438, 1.390625, 1.359375, 1.328125,
1.3125, 1.289062, 1.296875, 1.242188, 1.234375, 1.21875, 1.226562, 1.054688,
]
# fmt: on

194
src/diffusers/modular_pipelines/echo/decoders.py Normal file
View File

@@ -0,0 +1,194 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from typing import Any
import torch
from ...configuration_utils import FrozenDict
from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video
from ...pipelines.ltx2.vocoder import LTX2Vocoder
from ...utils.torch_utils import randn_tensor
from ...video_processor import VideoProcessor
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
# Copied from diffusers.modular_pipelines.ltx2.decoders._denormalize_latents
def _denormalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
# Denormalize video latents across the channel dimension [B, C, F, H, W].
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = latents * latents_std / scaling_factor + latents_mean
return latents
def _denormalize_audio_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor
) -> torch.Tensor:
# Mirror the encoder's channel/mel statistics on unpacked VAE latents.
latents_mean = latents_mean.view(1, latents.shape[1], 1, latents.shape[3]).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, latents.shape[1], 1, latents.shape[3]).to(latents.device, latents.dtype)
return (latents * latents_std) + latents_mean
class EchoVaeDecoderStep(ModularPipelineBlocks):
model_name = "echo"
@property
def description(self) -> str:
return "Denormalizes and decodes Echo video VAE latents into videos or returns denormalized latents."
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLLTX2Video),
ComponentSpec(
"video_processor",
VideoProcessor,
config=FrozenDict({"vae_scale_factor": 32}),
default_creation_method="from_config",
),
]
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"latents",
type_hint=torch.Tensor,
required=True,
description="Normalized video VAE latents of shape (B, C, F, H, W).",
),
InputParam.template("output_type", default="pil"),
InputParam("decode_timestep", default=0.0, description="Timestep used to decode the final latents."),
InputParam(
"decode_noise_scale",
default=None,
description="Noise interpolation factor applied at the decode timestep.",
),
InputParam.template("generator"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam.template("videos")]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
vae = components.vae
latents = block_state.latents
if block_state.output_type == "latent":
block_state.videos = _denormalize_latents(
latents,
components.latents_mean,
components.latents_std,
components.vae_scaling_factor,
)
self.set_block_state(state, block_state)
return components, state
if not vae.config.timestep_conditioning:
timestep = None
else:
batch_size = latents.shape[0]
decode_timestep = block_state.decode_timestep
decode_noise_scale = block_state.decode_noise_scale
noise = randn_tensor(
latents.shape,
generator=block_state.generator,
device=latents.device,
dtype=latents.dtype,
)
if not isinstance(decode_timestep, list):
decode_timestep = [decode_timestep] * batch_size
if decode_noise_scale is None:
decode_noise_scale = decode_timestep
elif not isinstance(decode_noise_scale, list):
decode_noise_scale = [decode_noise_scale] * batch_size
timestep = torch.tensor(decode_timestep, device=latents.device, dtype=latents.dtype)
decode_noise_scale = torch.tensor(
decode_noise_scale,
device=latents.device,
dtype=latents.dtype,
)[:, None, None, None, None]
latents = (1 - decode_noise_scale) * latents + decode_noise_scale * noise
latents = _denormalize_latents(
latents,
components.latents_mean,
components.latents_std,
components.vae_scaling_factor,
).to(vae.dtype)
video = vae.decode(latents, timestep, return_dict=False)[0]
block_state.videos = components.video_processor.postprocess_video(video, output_type=block_state.output_type)
self.set_block_state(state, block_state)
return components, state
class EchoAudioDecoderStep(ModularPipelineBlocks):
"""Decode Echo audio latents while preserving per-component mixed precision."""
model_name = "echo"
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("audio_vae", AutoencoderKLLTX2Audio),
ComponentSpec("vocoder", LTX2Vocoder),
]
@property
def inputs(self) -> list[tuple[str, Any]]:
return [
InputParam(
"audio_latents",
type_hint=torch.Tensor,
required=True,
description="Normalized audio VAE latents of shape (B, C, L, M).",
),
InputParam.template("output_type", default="pil"),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform.")]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
audio_vae = components.audio_vae
audio_latents = _denormalize_audio_latents(
block_state.audio_latents,
components.audio_latents_mean,
components.audio_latents_std,
)
if block_state.output_type == "latent":
block_state.audio = audio_latents
else:
audio_latents = audio_latents.to(audio_vae.dtype)
generated_mel_spectrograms = audio_vae.decode(audio_latents, return_dict=False)[0]
block_state.audio = components.vocoder(generated_mel_spectrograms.to(components.vocoder.dtype))
self.set_block_state(state, block_state)
return components, state

524
src/diffusers/modular_pipelines/echo/denoise.py Normal file
View File

@@ -0,0 +1,524 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import torch
from ...models import LTX2VideoTransformer3DModel
from ...utils.torch_utils import randn_tensor
from ..modular_pipeline import (
BlockState,
LoopSequentialPipelineBlocks,
ModularPipelineBlocks,
PipelineState,
)
from ..modular_pipeline_utils import ComponentSpec, InputParam, InsertableDict, OutputParam
from .modular_pipeline import EchoModularPipeline
DEFAULT_ECHO_SIGMAS = (1.0, 0.99375, 0.9875, 0.98125, 0.975, 0.909375, 0.725, 0.421875, 0.0)
class EchoLoopBeforeDenoiser(ModularPipelineBlocks):
model_name = "echo"
@property
def description(self) -> str:
return "Builds Echo's per-token video/audio timesteps and prepends clean memory tokens for one DMD step."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("latents", required=True),
InputParam(
"audio_latents", type_hint=torch.Tensor, required=True, description="Packed noisy target audio tokens."
),
InputParam(
"first_frame_token_count",
type_hint=int,
required=True,
description="Number of clean first-frame tokens.",
),
InputParam(
"memory_video_tokens",
type_hint=torch.Tensor,
required=False,
description="Packed clean image-memory tokens.",
),
InputParam(
"memory_video_coords",
type_hint=torch.Tensor,
required=False,
description="RoPE coordinates for image-memory tokens.",
),
InputParam(
"memory_audio_tokens",
type_hint=torch.Tensor,
required=False,
description="Packed clean audio-memory tokens.",
),
InputParam(
"memory_audio_coords",
type_hint=torch.Tensor,
required=False,
description="RoPE coordinates for audio-memory tokens.",
),
InputParam(
"video_coords",
type_hint=torch.Tensor,
required=True,
description="RoPE coordinates for target video tokens.",
),
InputParam(
"audio_coords",
type_hint=torch.Tensor,
required=True,
description="RoPE coordinates for target audio tokens.",
),
]
@torch.no_grad()
def __call__(
self, components, block_state: BlockState, i: int, sigma: float
) -> tuple[EchoModularPipeline, BlockState]:
batch_size, video_token_count = block_state.latents.shape[:2]
audio_token_count = block_state.audio_latents.shape[1]
scaled_sigma = float(sigma) * float(components.transformer.config.timestep_scale_multiplier)
video_timestep = torch.full(
(batch_size, video_token_count),
scaled_sigma,
device=block_state.latents.device,
dtype=torch.float32,
)
if block_state.first_frame_token_count:
video_timestep[:, : block_state.first_frame_token_count] = 0
audio_timestep = torch.full(
(batch_size, audio_token_count),
scaled_sigma,
device=block_state.audio_latents.device,
dtype=torch.float32,
)
block_state.latent_model_input = block_state.latents
block_state.model_video_timestep = video_timestep
block_state.model_video_coords = block_state.video_coords
if block_state.memory_video_tokens is not None:
memory_timestep = torch.zeros(
(batch_size, block_state.memory_video_tokens.shape[1]),
device=video_timestep.device,
dtype=video_timestep.dtype,
)
block_state.latent_model_input = torch.cat([block_state.memory_video_tokens, block_state.latents], dim=1)
block_state.model_video_timestep = torch.cat([memory_timestep, video_timestep], dim=1)
block_state.model_video_coords = torch.cat(
[block_state.memory_video_coords, block_state.video_coords], dim=2
)
block_state.audio_latent_model_input = block_state.audio_latents
block_state.model_audio_timestep = audio_timestep
block_state.model_audio_coords = block_state.audio_coords
if block_state.memory_audio_tokens is not None:
memory_audio_timestep = torch.zeros(
(batch_size, block_state.memory_audio_tokens.shape[1]),
device=audio_timestep.device,
dtype=audio_timestep.dtype,
)
block_state.audio_latent_model_input = torch.cat(
[block_state.memory_audio_tokens, block_state.audio_latents], dim=1
)
block_state.model_audio_timestep = torch.cat([memory_audio_timestep, audio_timestep], dim=1)
block_state.model_audio_coords = torch.cat(
[block_state.memory_audio_coords, block_state.audio_coords], dim=2
)
block_state.video_timestep = video_timestep
block_state.audio_timestep = audio_timestep
return components, block_state
class EchoLoopDenoiser(ModularPipelineBlocks):
model_name = "echo"
@property
def description(self) -> str:
return (
"Runs the Echo transformer once without CFG, using clean ref/memory token timesteps and cross-modal "
"timestep exchange, then converts velocity predictions to x0."
)
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"connector_prompt_embeds",
type_hint=torch.Tensor,
required=True,
description="Positive video-branch text conditioning.",
),
InputParam(
"connector_audio_prompt_embeds",
type_hint=torch.Tensor,
required=True,
description="Positive audio-branch text conditioning.",
),
InputParam(
"connector_attention_mask",
type_hint=torch.Tensor,
required=True,
description="Binary attention mask for text conditioning.",
),
InputParam.template("latents", required=True),
InputParam(
"audio_latents", type_hint=torch.Tensor, required=True, description="Packed noisy target audio tokens."
),
InputParam(
"video_coords",
type_hint=torch.Tensor,
required=True,
description="RoPE coordinates for target video tokens.",
),
InputParam(
"audio_coords",
type_hint=torch.Tensor,
required=True,
description="RoPE coordinates for target audio tokens.",
),
InputParam(
"latent_num_frames", type_hint=int, required=True, description="Number of target video latent frames."
),
InputParam("latent_height", type_hint=int, required=True, description="Target video latent height."),
InputParam("latent_width", type_hint=int, required=True, description="Target video latent width."),
InputParam(
"audio_num_frames", type_hint=int, required=True, description="Number of target audio latent frames."
),
InputParam(
"memory_video_token_count",
type_hint=int,
required=True,
description="Number of prepended image-memory tokens.",
),
InputParam(
"memory_audio_token_count",
type_hint=int,
required=True,
description="Number of prepended audio-memory tokens.",
),
InputParam("frame_rate", type_hint=float, default=25.0, description="Frame rate of the generated video."),
InputParam.template("attention_kwargs"),
]
@torch.no_grad()
def __call__(
self, components, block_state: BlockState, i: int, sigma: float
) -> tuple[EchoModularPipeline, BlockState]:
transformer_dtype = components.transformer.dtype
video_context = block_state.connector_prompt_embeds.to(transformer_dtype)
audio_context = block_state.connector_audio_prompt_embeds.to(transformer_dtype)
context_mask = block_state.connector_attention_mask
velocity_video, velocity_audio = components.transformer(
hidden_states=block_state.latent_model_input.to(transformer_dtype),
audio_hidden_states=block_state.audio_latent_model_input.to(transformer_dtype),
encoder_hidden_states=video_context,
audio_encoder_hidden_states=audio_context,
timestep=block_state.model_video_timestep,
audio_timestep=block_state.model_audio_timestep,
# Echo deliberately uses the first target token as its global video sigma. When a clean first frame is
# present this is zero, while the audio branch keeps the current DMD sigma. `use_cross_timestep=True`
# exchanges these values in the cross-modal modulation blocks, matching the released wrapper.
sigma=block_state.video_timestep[:, 0],
audio_sigma=block_state.audio_timestep[:, 0],
encoder_attention_mask=context_mask,
audio_encoder_attention_mask=context_mask,
num_frames=block_state.latent_num_frames,
height=block_state.latent_height,
width=block_state.latent_width,
fps=block_state.frame_rate,
audio_num_frames=block_state.audio_num_frames,
video_coords=block_state.model_video_coords,
audio_coords=block_state.model_audio_coords,
isolate_modalities=False,
spatio_temporal_guidance_blocks=None,
perturbation_mask=None,
use_cross_timestep=True,
attention_kwargs=block_state.attention_kwargs,
return_dict=False,
)
velocity_video = velocity_video[:, block_state.memory_video_token_count :]
velocity_audio = velocity_audio[:, block_state.memory_audio_token_count :]
timestep_scale = float(components.transformer.config.timestep_scale_multiplier)
block_state.predicted_video_x0 = (
block_state.latents.float()
- velocity_video.float() * block_state.video_timestep.unsqueeze(-1) / timestep_scale
).to(transformer_dtype)
block_state.predicted_audio_x0 = (
block_state.audio_latents.float()
- velocity_audio.float() * block_state.audio_timestep.unsqueeze(-1) / timestep_scale
).to(transformer_dtype)
if block_state.first_frame_tokens is not None:
block_state.predicted_video_x0[:, : block_state.first_frame_token_count] = block_state.first_frame_tokens
return components, block_state
class EchoLoopAfterDenoiser(ModularPipelineBlocks):
model_name = "echo"
@property
def description(self) -> str:
return "Re-noises each predicted x0 with fresh Gaussian noise at the next Echo DMD sigma."
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("latents", required=True),
InputParam(
"audio_latents", type_hint=torch.Tensor, required=True, description="Packed target audio tokens."
),
InputParam(
"sigmas",
type_hint=list | tuple,
required=True,
description="DMD sigma schedule including the terminal zero.",
),
InputParam.template("generator"),
InputParam(
"first_frame_tokens",
type_hint=torch.Tensor,
required=False,
description="Packed clean first-frame tokens.",
),
InputParam(
"first_frame_token_count",
type_hint=int,
required=True,
description="Number of clean first-frame tokens.",
),
]
@torch.no_grad()
def __call__(
self, components, block_state: BlockState, i: int, sigma: float
) -> tuple[EchoModularPipeline, BlockState]:
next_sigma = float(block_state.sigmas[i + 1])
if next_sigma > 0:
video_noise = randn_tensor(
block_state.predicted_video_x0.shape,
generator=block_state.generator,
device=block_state.predicted_video_x0.device,
dtype=block_state.predicted_video_x0.dtype,
)
audio_noise = randn_tensor(
block_state.predicted_audio_x0.shape,
generator=block_state.generator,
device=block_state.predicted_audio_x0.device,
dtype=block_state.predicted_audio_x0.dtype,
)
block_state.latents = block_state.predicted_video_x0 * (1.0 - next_sigma) + video_noise * next_sigma
block_state.audio_latents = block_state.predicted_audio_x0 * (1.0 - next_sigma) + audio_noise * next_sigma
if block_state.first_frame_tokens is not None:
block_state.latents[:, : block_state.first_frame_token_count] = block_state.first_frame_tokens
else:
block_state.latents = block_state.predicted_video_x0
block_state.audio_latents = block_state.predicted_audio_x0
return components, block_state
EchoDenoiseLoopBlocks = InsertableDict(
[
("before_denoiser", EchoLoopBeforeDenoiser()),
("denoiser", EchoLoopDenoiser()),
("after_denoiser", EchoLoopAfterDenoiser()),
]
)
# auto_docstring
class EchoDenoiseLoopStep(LoopSequentialPipelineBlocks):
"""
Iteratively predicts clean video/audio latents and re-noises them with fresh Gaussian noise according to Echo's DMD
sigma schedule.
Components:
transformer (`LTX2VideoTransformer3DModel`)
Inputs:
sigmas (`list | tuple`):
DMD sigma schedule, including the terminal zero.
latents (`Tensor`):
Pre-generated noisy latents for image generation.
audio_latents (`Tensor`):
Packed noisy target audio tokens.
first_frame_token_count (`int`):
Number of clean first-frame tokens.
memory_video_tokens (`Tensor`, *optional*):
Packed clean image-memory tokens.
memory_video_coords (`Tensor`, *optional*):
RoPE coordinates for image-memory tokens.
memory_audio_tokens (`Tensor`, *optional*):
Packed clean audio-memory tokens.
memory_audio_coords (`Tensor`, *optional*):
RoPE coordinates for audio-memory tokens.
video_coords (`Tensor`):
RoPE coordinates for target video tokens.
audio_coords (`Tensor`):
RoPE coordinates for target audio tokens.
connector_prompt_embeds (`Tensor`):
Positive video-branch text conditioning.
connector_audio_prompt_embeds (`Tensor`):
Positive audio-branch text conditioning.
connector_attention_mask (`Tensor`):
Binary attention mask for text conditioning.
latent_num_frames (`int`):
Number of target video latent frames.
latent_height (`int`):
Target video latent height.
latent_width (`int`):
Target video latent width.
audio_num_frames (`int`):
Number of target audio latent frames.
memory_video_token_count (`int`):
Number of prepended image-memory tokens.
memory_audio_token_count (`int`):
Number of prepended audio-memory tokens.
frame_rate (`float`, *optional*, defaults to 25.0):
Frame rate of the generated video.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
audio_latents (`Tensor`):
Packed target audio tokens.
sigmas (`list | tuple`):
DMD sigma schedule including the terminal zero.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
first_frame_tokens (`Tensor`, *optional*):
Packed clean first-frame tokens.
"""
model_name = "echo"
block_classes = EchoDenoiseLoopBlocks.values()
block_names = EchoDenoiseLoopBlocks.keys()
@property
def description(self) -> str:
return (
"Iteratively predicts clean video/audio latents and re-noises them with fresh Gaussian noise according "
"to Echo's DMD sigma schedule."
)
@property
def loop_expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]
@property
def loop_inputs(self) -> list[InputParam]:
return [
InputParam(
"sigmas",
type_hint=list | tuple,
default=DEFAULT_ECHO_SIGMAS,
description="DMD sigma schedule, including the terminal zero.",
)
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
sigmas = [float(value) for value in block_state.sigmas]
if len(sigmas) < 2 or sigmas[-1] != 0.0:
raise ValueError("Echo `sigmas` must contain at least two values and end at 0.")
if any(left < right for left, right in zip(sigmas, sigmas[1:])):
raise ValueError("Echo `sigmas` must be monotonically non-increasing.")
block_state.sigmas = sigmas
with self.progress_bar(total=len(sigmas) - 1) as progress_bar:
for i, sigma in enumerate(sigmas[:-1]):
components, block_state = self.loop_step(components, block_state, i=i, sigma=sigma)
progress_bar.update()
self.set_block_state(state, block_state)
return components, state
# Copied from diffusers.modular_pipelines.ltx2.decoders._unpack_latents
def _unpack_latents(
latents: torch.Tensor, num_frames: int, height: int, width: int, patch_size: int = 1, patch_size_t: int = 1
) -> torch.Tensor:
# Packed video latents of shape [B, S, D] are unpacked into a video tensor of shape [B, C, F, H, W].
batch_size = latents.size(0)
latents = latents.reshape(batch_size, num_frames, height, width, -1, patch_size_t, patch_size, patch_size)
latents = latents.permute(0, 4, 1, 5, 2, 6, 3, 7).flatten(6, 7).flatten(4, 5).flatten(2, 3)
return latents
class EchoUnpackLatentsStep(ModularPipelineBlocks):
"""Return normalized video and audio VAE tensors at the core-denoise boundary."""
model_name = "echo"
@property
def description(self) -> str:
return "Unpacks denoised video and audio tokens while leaving VAE denormalization to the decoders."
@property
def expected_components(self) -> list[ComponentSpec]:
return [ComponentSpec("transformer", LTX2VideoTransformer3DModel)]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("latents", required=True),
InputParam("audio_latents", type_hint=torch.Tensor, required=True, description="Denoised audio tokens."),
InputParam("latent_num_frames", type_hint=int, required=True, description="Video latent frame count."),
InputParam("latent_height", type_hint=int, required=True, description="Video latent height."),
InputParam("latent_width", type_hint=int, required=True, description="Video latent width."),
InputParam(
"audio_latent_mel_bins", type_hint=int, required=True, description="Audio VAE latent mel-bin count."
),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="Normalized video VAE latents (B, C, F, H, W)."
),
OutputParam(
"audio_latents", type_hint=torch.Tensor, description="Normalized audio VAE latents (B, C, L, M)."
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
patch_size = components.transformer_spatial_patch_size
patch_size_t = components.transformer_temporal_patch_size
block_state.latents = _unpack_latents(
block_state.latents,
block_state.latent_num_frames // patch_size_t,
block_state.latent_height // patch_size,
block_state.latent_width // patch_size,
patch_size,
patch_size_t,
)
block_state.audio_latents = block_state.audio_latents.unflatten(
2, (-1, block_state.audio_latent_mel_bins)
).transpose(1, 2)
self.set_block_state(state, block_state)
return components, state

499
src/diffusers/modular_pipelines/echo/encoders.py Normal file
View File

@@ -0,0 +1,499 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import annotations
import math
import numpy as np
import PIL.Image
import torch
import torch.nn.functional as F
from transformers import PreTrainedModel, PreTrainedTokenizerBase
from ...models import AutoencoderKLLTX2Audio, AutoencoderKLLTX2Video
from ...pipelines.ltx2.connectors import LTX2TextConnectors
from ...utils import logging
from ..modular_pipeline import ModularPipelineBlocks, PipelineState
from ..modular_pipeline_utils import ComponentSpec, InputParam, OutputParam
logger = logging.get_logger(__name__)
MAX_MEMORY_SLOTS = 7
MAX_MEMORY_AUDIO_DURATION_SECONDS = 9.62
def _validate_memory_slot_count(slot_count: int) -> None:
if slot_count > MAX_MEMORY_SLOTS:
raise ValueError(f"Echo accepts at most {MAX_MEMORY_SLOTS} memory slots, but received {slot_count}.")
def _get_prompt_embeds(
tokenizer,
text_encoder,
prompt: str | list[str],
max_sequence_length: int,
device: torch.device,
dtype: torch.dtype,
):
prompt = [prompt] if isinstance(prompt, str) else prompt
tokenizer.padding_side = "left"
if tokenizer.pad_token is None:
tokenizer.pad_token = tokenizer.eos_token
prompt = [text.strip() for text in prompt]
text_inputs = tokenizer(
prompt,
padding="max_length",
max_length=max_sequence_length,
truncation=True,
add_special_tokens=True,
return_tensors="pt",
)
text_input_ids = text_inputs.input_ids.to(device)
prompt_attention_mask = text_inputs.attention_mask.to(device)
text_encoder_outputs = text_encoder(
input_ids=text_input_ids,
attention_mask=prompt_attention_mask,
output_hidden_states=True,
)
text_encoder_hidden_states = torch.stack(text_encoder_outputs.hidden_states, dim=-1)
prompt_embeds = text_encoder_hidden_states.flatten(2, 3).to(dtype=dtype)
return prompt_embeds, prompt_attention_mask
class EchoTextEncoderStep(ModularPipelineBlocks):
"""Encode only the positive prompt used by the guidance-free Echo DMD checkpoint."""
model_name = "echo"
@property
def description(self) -> str:
return "Encodes the positive Echo prompt into packed per-layer Gemma hidden states."
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("text_encoder", PreTrainedModel),
ComponentSpec("tokenizer", PreTrainedTokenizerBase),
]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam.template("prompt", required=True),
InputParam.template("max_sequence_length", default=1024),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"prompt_embeds",
type_hint=torch.Tensor,
description="Packed per-layer Gemma hidden states for the prompt.",
),
OutputParam(
"prompt_attention_mask",
type_hint=torch.Tensor,
description="Binary attention mask for `prompt_embeds`.",
),
]
@staticmethod
def check_inputs(block_state):
if block_state.prompt is not None and not isinstance(block_state.prompt, (str, list)):
raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(block_state.prompt)}")
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
self.check_inputs(block_state)
prompt = [block_state.prompt] if isinstance(block_state.prompt, str) else block_state.prompt
block_state.prompt_embeds, block_state.prompt_attention_mask = _get_prompt_embeds(
components.tokenizer,
components.text_encoder,
prompt,
block_state.max_sequence_length,
components._execution_device,
components.text_encoder.dtype,
)
self.set_block_state(state, block_state)
return components, state
class EchoTextConnectorStep(ModularPipelineBlocks):
"""Project positive Gemma embeddings into Echo's video and audio context spaces."""
model_name = "echo"
@property
def description(self) -> str:
return "Adapts positive Gemma embeddings for the Echo video and audio transformer branches."
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("connectors", LTX2TextConnectors),
ComponentSpec("tokenizer", PreTrainedTokenizerBase),
]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam("prompt_embeds", type_hint=torch.Tensor, required=True),
InputParam("prompt_attention_mask", type_hint=torch.Tensor, required=True),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"connector_prompt_embeds",
type_hint=torch.Tensor,
description="Video-branch positive text conditioning.",
),
OutputParam(
"connector_audio_prompt_embeds",
type_hint=torch.Tensor,
description="Audio-branch positive text conditioning.",
),
OutputParam(
"connector_attention_mask",
type_hint=torch.Tensor,
description="Binary attention mask for the positive text conditioning.",
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
(
block_state.connector_prompt_embeds,
block_state.connector_audio_prompt_embeds,
block_state.connector_attention_mask,
) = components.connectors(
block_state.prompt_embeds,
block_state.prompt_attention_mask,
padding_side=components.tokenizer.padding_side,
)
self.set_block_state(state, block_state)
return components, state
# Copied from diffusers.modular_pipelines.ltx2.before_denoise._normalize_latents
def _normalize_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor, scaling_factor: float = 1.0
) -> torch.Tensor:
latents_mean = latents_mean.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, -1, 1, 1, 1).to(latents.device, latents.dtype)
latents = (latents - latents_mean) * scaling_factor / latents_std
return latents
def _normalize_audio_latents(
latents: torch.Tensor, latents_mean: torch.Tensor, latents_std: torch.Tensor
) -> torch.Tensor:
# Audio statistics are stored in flattened (channel, mel-bin) order. Broadcast them over the time axis
# without packing the VAE tensor, keeping the encoder independent of transformer tokenization.
latents_mean = latents_mean.view(1, latents.shape[1], 1, latents.shape[3]).to(latents.device, latents.dtype)
latents_std = latents_std.view(1, latents.shape[1], 1, latents.shape[3]).to(latents.device, latents.dtype)
return (latents - latents_mean) / latents_std
def _prepare_image(image, height: int, width: int) -> torch.Tensor:
if isinstance(image, PIL.Image.Image):
image = image.convert("RGB")
if image.size != (width, height):
image = image.resize((width, height), PIL.Image.Resampling.BICUBIC)
image = torch.from_numpy(np.asarray(image, dtype=np.float32).copy()).permute(2, 0, 1) / 127.5 - 1.0
return image.unsqueeze(0)
image = torch.as_tensor(image).detach().float()
if image.ndim == 3:
image = image.unsqueeze(0)
if image.ndim != 4 or image.shape[0] != 1 or image.shape[1] != 3:
raise ValueError(
"Each Echo image must be a PIL image or a tensor with shape (3, height, width) or "
f"(1, 3, height, width), but got {tuple(image.shape)}."
)
if image.shape[-2:] != (height, width):
image = F.interpolate(image, size=(height, width), mode="bicubic", align_corners=False)
if image.amin() >= 0:
image = image / 127.5 - 1.0 if image.amax() > 1.0 else image * 2.0 - 1.0
return image
def _normalize_waveform(waveform: torch.Tensor) -> torch.Tensor:
waveform = torch.as_tensor(waveform).detach().float()
while waveform.ndim > 2 and waveform.shape[0] == 1:
waveform = waveform.squeeze(0)
if waveform.ndim == 1:
waveform = waveform.unsqueeze(0)
if waveform.ndim != 2 or waveform.shape[-1] < 1:
raise ValueError(
"Each Echo memory waveform must have shape (samples,) or (channels, samples), "
f"but got {tuple(waveform.shape)}."
)
if waveform.shape[0] == 1:
waveform = waveform.repeat(2, 1)
elif waveform.shape[0] > 2:
waveform = waveform[:2]
return waveform.contiguous()
def _encode_image(
vae,
latents_mean: torch.Tensor,
latents_std: torch.Tensor,
scaling_factor: float,
image,
height: int,
width: int,
device: torch.device,
) -> torch.Tensor:
pixels = _prepare_image(image, height, width).unsqueeze(2).to(device=device, dtype=vae.dtype)
latents = vae.encode(pixels).latent_dist.mode().float()
return _normalize_latents(latents, latents_mean, latents_std, scaling_factor)
def _encode_audio(
audio_vae,
audio_latents_mean: torch.Tensor,
audio_latents_std: torch.Tensor,
waveform: torch.Tensor,
sample_rate: int,
device: torch.device,
) -> torch.Tensor:
try:
import torchaudio
except ImportError as error:
raise ImportError(
"Encoding raw `memory_audio_waveforms` requires torchaudio. Install torchaudio before running the "
"Echo VAE encoder."
) from error
if sample_rate <= 0:
raise ValueError(f"`memory_audio_sample_rates` values must be positive, but got {sample_rate}.")
waveform = _normalize_waveform(waveform).to(device=device, dtype=torch.float32)
target_rate = int(audio_vae.config.sample_rate)
n_fft = 1024
min_input_samples = math.ceil(n_fft * sample_rate / target_rate)
if waveform.shape[-1] < min_input_samples:
waveform = F.pad(waveform, (0, min_input_samples - waveform.shape[-1]))
if sample_rate != target_rate:
waveform = torchaudio.functional.resample(waveform, sample_rate, target_rate)
if waveform.shape[-1] < n_fft:
waveform = F.pad(waveform, (0, n_fft - waveform.shape[-1]))
hop_length = int(audio_vae.config.mel_hop_length)
mel_transform = torchaudio.transforms.MelSpectrogram(
sample_rate=target_rate,
n_fft=n_fft,
win_length=n_fft,
hop_length=hop_length,
f_min=0.0,
f_max=target_rate / 2.0,
n_mels=int(audio_vae.config.mel_bins),
window_fn=torch.hann_window,
center=True,
pad_mode="reflect",
power=1.0,
mel_scale="slaney",
norm="slaney",
).to(device)
mel = torch.log(torch.clamp(mel_transform(waveform), min=1e-5)).permute(0, 2, 1).unsqueeze(0)
max_samples = round(MAX_MEMORY_AUDIO_DURATION_SECONDS * target_rate)
if waveform.shape[-1] > max_samples:
max_mel_steps = max_samples // hop_length + 1
num_time_steps = mel.shape[2]
max_start = num_time_steps - max_mel_steps
scan_stride = max(1, max_mel_steps // 4)
candidate_starts = list(range(0, max_start + 1, scan_stride))
if candidate_starts[-1] != max_start:
candidate_starts.append(max_start)
response = mel.float().exp().sum(dim=(0, 1, 3))
cumulative_response = torch.cat([response.new_zeros(1), response.cumsum(dim=0)])
starts = torch.tensor(candidate_starts, device=mel.device, dtype=torch.long)
scores = cumulative_response[starts + max_mel_steps] - cumulative_response[starts]
start_mel_step = candidate_starts[int(scores.argmax().item())]
start_sample = min(start_mel_step * hop_length, waveform.shape[-1] - max_samples)
waveform = waveform[..., start_sample : start_sample + max_samples]
mel = torch.log(torch.clamp(mel_transform(waveform), min=1e-5)).permute(0, 2, 1).unsqueeze(0)
latents = audio_vae.encode(mel.to(audio_vae.dtype)).latent_dist.mode()
return _normalize_audio_latents(latents, audio_latents_mean, audio_latents_std).float()
class EchoVaeEncoderStep(ModularPipelineBlocks):
"""Encode Echo's clean first frame and ordered image/audio memory slots."""
model_name = "echo"
@property
def description(self) -> str:
return "Encodes the optional first frame and image/audio memory slots into normalized VAE latents."
@property
def expected_components(self) -> list[ComponentSpec]:
return [
ComponentSpec("vae", AutoencoderKLLTX2Video),
ComponentSpec("audio_vae", AutoencoderKLLTX2Audio),
]
@property
def inputs(self) -> list[InputParam]:
return [
InputParam(
"image",
type_hint=PIL.Image.Image | torch.Tensor,
default=None,
description="Optional single first frame used as a clean reference condition.",
),
InputParam(
"memory_images",
type_hint=list,
default=None,
description="Ordered reference images, one per Echo memory slot.",
),
InputParam(
"memory_audio_waveforms",
type_hint=list,
default=None,
description=(
"Ordered memory waveforms as `(channels, samples)` tensors. Inputs longer than 9.62 seconds are "
"cropped to their highest-response window. Use `None` for a silent slot."
),
),
InputParam(
"memory_audio_sample_rates",
type_hint=int | list,
default=None,
description="Sampling rate shared by all memory waveforms, or one rate per slot.",
),
InputParam.template("height", default=512),
InputParam.template("width", default=704),
]
@property
def intermediate_outputs(self) -> list[OutputParam]:
return [
OutputParam(
"first_frame_latents", type_hint=torch.Tensor, description="Normalized first-frame VAE latents."
),
OutputParam(
"memory_video_latents", type_hint=list, description="Normalized VAE latents for each memory image."
),
OutputParam(
"memory_audio_latents",
type_hint=list,
description="Normalized audio VAE latents of shape (B, C, L, M) for each memory slot.",
),
]
@torch.no_grad()
def __call__(self, components, state: PipelineState) -> PipelineState:
block_state = self.get_block_state(state)
device = components._execution_device
if block_state.image is not None:
first_frame_latents = _encode_image(
components.vae,
components.latents_mean,
components.latents_std,
components.vae_scaling_factor,
block_state.image,
block_state.height,
block_state.width,
device,
)
else:
first_frame_latents = None
memory_images = block_state.memory_images
if memory_images is None:
memory_images = []
elif not isinstance(memory_images, (list, tuple)):
memory_images = [memory_images]
else:
memory_images = list(memory_images)
slot_count = len(memory_images)
_validate_memory_slot_count(slot_count)
memory_video_latents = [
_encode_image(
components.vae,
components.latents_mean,
components.latents_std,
components.vae_scaling_factor,
image,
block_state.height,
block_state.width,
device,
)
for image in memory_images
]
raw_audio = block_state.memory_audio_waveforms
if raw_audio is None:
raw_audio = []
elif not isinstance(raw_audio, (list, tuple)):
raw_audio = [raw_audio]
else:
raw_audio = list(raw_audio)
if raw_audio and len(raw_audio) != slot_count:
raise ValueError("`memory_audio_waveforms` must have one entry per `memory_images` slot.")
memory_audio_latents: list[torch.Tensor | None] = []
if raw_audio:
rates = block_state.memory_audio_sample_rates
if rates is None:
rates = int(components.audio_vae.config.sample_rate)
rates = [rates] * len(raw_audio) if isinstance(rates, int) else list(rates)
if len(rates) != len(raw_audio):
raise ValueError("`memory_audio_sample_rates` must be an int or have one entry per waveform.")
if components.audio_vae.dtype != torch.float32:
logger.warning_once(
"Echo was trained with FP32 audio-memory encoding. For parity, load `audio_vae` with "
"`dtype={'default': torch.bfloat16, 'audio_vae': torch.float32}`."
)
for waveform, sample_rate in zip(raw_audio, rates):
memory_audio_latents.append(
None
if waveform is None
else _encode_audio(
components.audio_vae,
components.audio_latents_mean,
components.audio_latents_std,
waveform,
int(sample_rate),
device,
)
)
block_state.first_frame_latents = first_frame_latents
block_state.memory_video_latents = memory_video_latents or None
block_state.memory_audio_latents = memory_audio_latents or None
self.set_block_state(state, block_state)
return components, state

300
src/diffusers/modular_pipelines/echo/modular_blocks_echo.py Normal file
View File

@@ -0,0 +1,300 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from ..modular_pipeline import SequentialPipelineBlocks
from ..modular_pipeline_utils import InsertableDict, OutputParam
from .before_denoise import EchoInputsStep, EchoPrepareConditioningStep, EchoPrepareLatentsStep
from .decoders import EchoAudioDecoderStep, EchoVaeDecoderStep
from .denoise import EchoDenoiseLoopStep, EchoUnpackLatentsStep
from .encoders import EchoTextConnectorStep, EchoTextEncoderStep, EchoVaeEncoderStep
EchoTextEncoderBlocks = InsertableDict(
[
("encode", EchoTextEncoderStep()),
("connect", EchoTextConnectorStep()),
]
)
# auto_docstring
class EchoTextConditioningStep(SequentialPipelineBlocks):
"""
Positive-only text conditioning for the guidance-free Echo DMD checkpoint.
Components:
text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`)
Inputs:
prompt (`str`):
The prompt or prompts to guide image generation.
max_sequence_length (`int`, *optional*, defaults to 1024):
Maximum sequence length for prompt encoding.
Outputs:
prompt_embeds (`Tensor`):
Packed per-layer Gemma hidden states for the prompt.
prompt_attention_mask (`Tensor`):
Binary attention mask for `prompt_embeds`.
connector_prompt_embeds (`Tensor`):
Video-branch positive text conditioning.
connector_audio_prompt_embeds (`Tensor`):
Audio-branch positive text conditioning.
connector_attention_mask (`Tensor`):
Binary attention mask for the positive text conditioning.
"""
model_name = "echo"
block_classes = EchoTextEncoderBlocks.values()
block_names = EchoTextEncoderBlocks.keys()
@property
def description(self) -> str:
return "Positive-only text conditioning for the guidance-free Echo DMD checkpoint."
EchoDenoiseBlocks = InsertableDict(
[
("prepare_conditioning", EchoPrepareConditioningStep()),
("input", EchoInputsStep()),
("prepare_latents", EchoPrepareLatentsStep()),
("denoise_loop", EchoDenoiseLoopStep()),
("after_denoise", EchoUnpackLatentsStep()),
]
)
# auto_docstring
class EchoDenoiseStep(SequentialPipelineBlocks):
"""
Prepare Echo conditioning and target latents, then run stochastic DMD denoising.
Components:
transformer (`LTX2VideoTransformer3DModel`)
Inputs:
first_frame_latents (`Tensor`, *optional*):
Normalized VAE latents for the optional clean first frame.
memory_video_latents (`list`, *optional*):
Normalized VAE latents for each ordered memory image.
memory_audio_latents (`list`, *optional*):
Normalized audio VAE latents of shape (B, C, L, M) for each memory slot.
model_frame_rate (`float`, *optional*, defaults to 24.0):
Training-time frame rate used for Echo video RoPE coordinates.
memory_position_offset (`float`, *optional*, defaults to 500.0):
Temporal center assigned to the first memory slot.
memory_position_slot_stride (`float`, *optional*, defaults to 50.0):
Temporal distance between consecutive memory-slot centers.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos per prompt.
connector_prompt_embeds (`Tensor`):
Per-prompt video-branch text conditioning.
connector_audio_prompt_embeds (`Tensor`):
Per-prompt audio-branch text conditioning.
connector_attention_mask (`Tensor`):
Per-prompt binary text attention mask.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to 704):
The width in pixels of the generated image.
num_frames (`int`, *optional*, defaults to 241):
Number of generated pixel frames; must be `1 + k * vae_temporal_compression_ratio`.
frame_rate (`float`, *optional*, defaults to 25.0):
Frame rate of the generated video.
latents (`Tensor`, *optional*):
Optional initial video noise in VAE form (B, C, F, H, W).
audio_latents (`Tensor`, *optional*):
Optional initial audio noise in VAE form (B, C, L, M).
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
sigmas (`list | tuple`):
DMD sigma schedule, including the terminal zero.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
Outputs:
latents (`Tensor`):
Normalized video VAE latents (B, C, F, H, W).
audio_latents (`Tensor`):
Normalized audio VAE latents (B, C, L, M).
"""
model_name = "echo"
block_classes = EchoDenoiseBlocks.values()
block_names = EchoDenoiseBlocks.keys()
@property
def description(self) -> str:
return "Prepare Echo conditioning and target latents, then run stochastic DMD denoising."
@property
def outputs(self) -> list[OutputParam]:
return [
OutputParam(
"latents", type_hint=torch.Tensor, description="Normalized video VAE latents (B, C, F, H, W)."
),
OutputParam(
"audio_latents", type_hint=torch.Tensor, description="Normalized audio VAE latents (B, C, L, M)."
),
]
EchoDecoderBlocks = InsertableDict(
[
("video_decode", EchoVaeDecoderStep()),
("audio_decode", EchoAudioDecoderStep()),
]
)
# auto_docstring
class EchoDecoderStep(SequentialPipelineBlocks):
"""
Decode Echo video and audio outputs with mixed-precision-safe audio vocoding.
Components:
vae (`AutoencoderKLLTX2Video`) video_processor (`VideoProcessor`) audio_vae (`AutoencoderKLLTX2Audio`)
vocoder (`LTX2Vocoder`)
Inputs:
latents (`Tensor`):
Normalized video VAE latents of shape (B, C, F, H, W).
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
decode_timestep (`None`, *optional*, defaults to 0.0):
Timestep used to decode the final latents.
decode_noise_scale (`None`, *optional*):
Noise interpolation factor applied at the decode timestep.
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
audio_latents (`Tensor`):
Normalized audio VAE latents of shape (B, C, L, M).
Outputs:
videos (`list`):
The generated videos.
audio (`Tensor`):
The generated audio waveform.
"""
model_name = "echo"
block_classes = EchoDecoderBlocks.values()
block_names = EchoDecoderBlocks.keys()
@property
def description(self) -> str:
return "Decode Echo video and audio outputs with mixed-precision-safe audio vocoding."
@property
def outputs(self) -> list[OutputParam]:
return [
OutputParam.template("videos"),
OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."),
]
EchoPipelineBlocks = InsertableDict(
[
("text_encoder", EchoTextConditioningStep()),
("vae_encoder", EchoVaeEncoderStep()),
("denoise", EchoDenoiseStep()),
("decode", EchoDecoderStep()),
]
)
# auto_docstring
class EchoBlocks(SequentialPipelineBlocks):
"""
Echo reference-to-video generation with clean first-frame conditioning, ordered image/audio memory slots, and
stochastic DMD denoising.
Components:
text_encoder (`PreTrainedModel`) tokenizer (`PreTrainedTokenizerBase`) connectors (`LTX2TextConnectors`) vae
(`AutoencoderKLLTX2Video`) audio_vae (`AutoencoderKLLTX2Audio`) transformer (`LTX2VideoTransformer3DModel`)
video_processor (`VideoProcessor`) vocoder (`LTX2Vocoder`)
Inputs:
prompt (`str`):
The prompt or prompts to guide image generation.
max_sequence_length (`int`, *optional*, defaults to 1024):
Maximum sequence length for prompt encoding.
image (`Image | Tensor`, *optional*):
Optional single first frame used as a clean reference condition.
memory_images (`list`, *optional*):
Ordered reference images, one per Echo memory slot.
memory_audio_waveforms (`list`, *optional*):
Ordered memory waveforms as `(channels, samples)` tensors. Inputs longer than 9.62 seconds are cropped to
their highest-response window. Use `None` for a silent slot.
memory_audio_sample_rates (`int | list`, *optional*):
Sampling rate shared by all memory waveforms, or one rate per slot.
height (`int`, *optional*, defaults to 512):
The height in pixels of the generated image.
width (`int`, *optional*, defaults to 704):
The width in pixels of the generated image.
model_frame_rate (`float`, *optional*, defaults to 24.0):
Training-time frame rate used for Echo video RoPE coordinates.
memory_position_offset (`float`, *optional*, defaults to 500.0):
Temporal center assigned to the first memory slot.
memory_position_slot_stride (`float`, *optional*, defaults to 50.0):
Temporal distance between consecutive memory-slot centers.
num_videos_per_prompt (`int`, *optional*, defaults to 1):
Number of videos per prompt.
num_frames (`int`, *optional*, defaults to 241):
Number of generated pixel frames; must be `1 + k * vae_temporal_compression_ratio`.
frame_rate (`float`, *optional*, defaults to 25.0):
Frame rate of the generated video.
latents (`Tensor`, *optional*):
Optional initial video noise in VAE form (B, C, F, H, W).
audio_latents (`Tensor`, *optional*):
Optional initial audio noise in VAE form (B, C, L, M).
generator (`Generator`, *optional*):
Torch generator for deterministic generation.
sigmas (`list | tuple`):
DMD sigma schedule, including the terminal zero.
attention_kwargs (`dict`, *optional*):
Additional kwargs for attention processors.
output_type (`str`, *optional*, defaults to pil):
Output format: 'pil', 'np', 'pt'.
decode_timestep (`None`, *optional*, defaults to 0.0):
Timestep used to decode the final latents.
decode_noise_scale (`None`, *optional*):
Noise interpolation factor applied at the decode timestep.
Outputs:
videos (`list`):
The generated videos.
audio (`Tensor`):
The generated audio waveform.
"""
model_name = "echo"
block_classes = EchoPipelineBlocks.values()
block_names = EchoPipelineBlocks.keys()
@property
def description(self) -> str:
return (
"Echo reference-to-video generation with clean first-frame conditioning, ordered image/audio memory "
"slots, and stochastic DMD denoising."
)
@property
def outputs(self) -> list[OutputParam]:
return [
OutputParam.template("videos"),
OutputParam("audio", type_hint=torch.Tensor, description="The generated audio waveform."),
]

108
src/diffusers/modular_pipelines/echo/modular_pipeline.py Normal file
View File

@@ -0,0 +1,108 @@
# Copyright 2026 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from ...loaders import LTX2LoraLoaderMixin as EchoLoraLoaderMixin
from ..modular_pipeline import ModularPipeline
from .constants import (
ECHO_AUDIO_LATENTS_MEAN,
ECHO_AUDIO_LATENTS_STD,
ECHO_LATENTS_MEAN,
ECHO_LATENTS_STD,
)
class EchoModularPipeline(ModularPipeline, EchoLoraLoaderMixin):
"""A ModularPipeline for the Echo long-video DMD checkpoint."""
default_blocks_name = "EchoBlocks"
@property
def vae_spatial_compression_ratio(self):
if getattr(self, "vae", None) is not None:
return self.vae.spatial_compression_ratio
return 32
@property
def vae_temporal_compression_ratio(self):
if getattr(self, "vae", None) is not None:
return self.vae.temporal_compression_ratio
return 8
@property
def vae_scaling_factor(self):
if getattr(self, "vae", None) is not None:
return self.vae.config.scaling_factor
return 1.0
@property
def latents_mean(self):
if getattr(self, "vae", None) is not None:
return self.vae.latents_mean
return torch.tensor(ECHO_LATENTS_MEAN)
@property
def latents_std(self):
if getattr(self, "vae", None) is not None:
return self.vae.latents_std
return torch.tensor(ECHO_LATENTS_STD)
@property
def transformer_spatial_patch_size(self):
if getattr(self, "transformer", None) is not None:
return self.transformer.config.patch_size
return 1
@property
def transformer_temporal_patch_size(self):
if getattr(self, "transformer", None) is not None:
return self.transformer.config.patch_size_t
return 1
@property
def audio_latent_mel_bins(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.config.mel_bins // self.audio_vae.mel_compression_ratio
return 16
@property
def audio_vae_temporal_compression_ratio(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.temporal_compression_ratio
return 4
@property
def audio_latents_mean(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.latents_mean
return torch.tensor(ECHO_AUDIO_LATENTS_MEAN)
@property
def audio_latents_std(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.latents_std
return torch.tensor(ECHO_AUDIO_LATENTS_STD)
@property
def audio_sampling_rate(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.config.sample_rate
return 16000
@property
def audio_hop_length(self):
if getattr(self, "audio_vae", None) is not None:
return self.audio_vae.config.mel_hop_length
return 160

1
src/diffusers/modular_pipelines/modular_pipeline.py
View File

@@ -196,6 +196,7 @@ MODULAR_PIPELINE_MAPPING = OrderedDict(
("helios", _create_default_map_fn("HeliosModularPipeline")),
("helios-pyramid", _helios_pyramid_map_fn),
("hunyuan-video-1.5", _create_default_map_fn("HunyuanVideo15ModularPipeline")),
("echo", _create_default_map_fn("EchoModularPipeline")),
("ltx", _create_default_map_fn("LTXModularPipeline")),
("ltx2", _create_default_map_fn("LTX2ModularPipeline")),
("ltx2.5", _create_default_map_fn("LTX25ModularPipeline")),

30
src/diffusers/utils/dummy_torch_and_transformers_objects.py
View File

@@ -92,6 +92,36 @@ class Cosmos3OmniModularPipeline(metaclass=DummyObject):
requires_backends(cls, ["torch", "transformers"])
class EchoBlocks(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class EchoModularPipeline(metaclass=DummyObject):
_backends = ["torch", "transformers"]
def __init__(self, *args, **kwargs):
requires_backends(self, ["torch", "transformers"])
@classmethod
def from_config(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
@classmethod
def from_pretrained(cls, *args, **kwargs):
requires_backends(cls, ["torch", "transformers"])
class ErnieImageAutoBlocks(metaclass=DummyObject):
_backends = ["torch", "transformers"]

0
tests/modular_pipelines/echo/__init__.py Normal file
View File

302
tests/modular_pipelines/echo/test_modular_pipeline_echo.py Normal file
View File

@@ -0,0 +1,302 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from copy import deepcopy
import pytest
import torch
from diffusers.modular_pipelines.echo.before_denoise import EchoInputsStep, EchoPrepareConditioningStep
from diffusers.modular_pipelines.echo.decoders import EchoAudioDecoderStep
from diffusers.modular_pipelines.echo.encoders import EchoVaeEncoderStep
from diffusers.modular_pipelines.echo.modular_blocks_echo import EchoDecoderStep
from ..testing_utils import (
BaseModularPipelineOutputMixin,
ModularLoadingTesterMixin,
ModularMemoryTesterMixin,
ModularPipelineTesterMixin,
ModularWorkflowTesterMixin,
)
from .testing_utils import EchoModularPipelineTesterConfig
class TestEchoModularPipelineFast(EchoModularPipelineTesterConfig, ModularPipelineTesterMixin):
def test_multi_prompt_multi_video(self):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
inputs["prompt"] = ["a robot dancing", "a cat walking"]
inputs["num_videos_per_prompt"] = 2
output = pipe(**inputs, output=["videos", "audio"])
assert output["videos"].shape == (4, 5, 3, 32, 32)
assert output["audio"].shape[0] == 4
assert torch.isfinite(output["videos"]).all()
assert torch.isfinite(output["audio"]).all()
def test_denoise_returns_initial_vae_form_at_zero_sigma(self):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
inputs.pop("image")
inputs.pop("memory_images")
inputs["sigmas"] = [0.0, 0.0]
inputs["latents"] = torch.randn(1, 4, 3, 16, 16)
inputs["audio_latents"] = torch.randn(1, 2, 5, 2)
output = pipe(**inputs, output=["latents", "audio_latents"])
torch.testing.assert_close(output["latents"], inputs["latents"], rtol=0, atol=0)
torch.testing.assert_close(output["audio_latents"], inputs["audio_latents"], rtol=0, atol=0)
def test_inference_batch_single_identical(self):
super().test_inference_batch_single_identical(expected_max_diff=1e-3)
def test_video_and_audio_outputs(self):
pipe = self.get_pipeline()
output = pipe(**self.get_dummy_inputs(), output=["videos", "audio"])
assert output["videos"].shape == (1, 5, 3, 32, 32)
assert output["audio"].shape[0] == 1
assert output["audio"].shape[1] == pipe.vocoder.config.out_channels
assert torch.isfinite(output["videos"]).all()
assert torch.isfinite(output["audio"]).all()
def test_seeded_dmd_renoising_is_deterministic(self):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs(seed=1)
inputs["output_type"] = "latent"
first = pipe(**inputs, output=["videos", "audio"])
inputs = self.get_dummy_inputs(seed=1)
inputs["output_type"] = "latent"
second = pipe(**inputs, output=["videos", "audio"])
assert torch.equal(first["videos"], second["videos"])
assert torch.equal(first["audio"], second["audio"])
def test_audio_decoder_supports_mixed_precision(self):
pipe = self.get_pipeline()
pipe.vocoder.to(dtype=torch.float64)
output = pipe(**self.get_dummy_inputs(), output=["videos", "audio"])
assert output["audio"].dtype == torch.float64
assert torch.isfinite(output["audio"]).all()
def test_user_latents_are_not_modified(self):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
inputs["output_type"] = "latent"
latent_num_frames = 1 + (inputs["num_frames"] - 1) // pipe.vae_temporal_compression_ratio
latent_height = inputs["height"] // pipe.vae_spatial_compression_ratio
latent_width = inputs["width"] // pipe.vae_spatial_compression_ratio
latents = torch.randn(1, pipe.vae.config.latent_channels, latent_num_frames, latent_height, latent_width)
original = latents.clone()
inputs["latents"] = latents
pipe(**inputs, output=["videos", "audio"])
assert torch.equal(latents, original)
@pytest.mark.parametrize("waveform_lengths", [(3200, None), (3200, 4800), (1, 512)])
def test_raw_audio_memory(self, waveform_lengths):
pytest.importorskip("torchaudio")
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
inputs["memory_audio_waveforms"] = [
None if length is None else torch.rand(2, length) for length in waveform_lengths
]
inputs["memory_audio_sample_rates"] = 16000
inputs["output_type"] = "latent"
output = pipe(**inputs, output=["videos", "audio"])
assert torch.isfinite(output["videos"]).all()
assert torch.isfinite(output["audio"]).all()
@pytest.mark.parametrize(
("sigmas", "message"),
[
([1.0], "at least two values"),
([1.0, 0.5], "end at 0"),
([0.5, 1.0, 0.0], "monotonically non-increasing"),
],
)
def test_invalid_dmd_schedules(self, sigmas, message):
pipe = self.get_pipeline()
inputs = self.get_dummy_inputs()
inputs["sigmas"] = sigmas
with pytest.raises(ValueError, match=message):
pipe(**inputs)
class TestEchoModularBlocks(EchoModularPipelineTesterConfig, BaseModularPipelineOutputMixin):
def test_vae_encoder_rejects_more_than_seven_memory_slots(self):
pipe = EchoVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
with pytest.raises(ValueError, match="at most 7 memory slots"):
pipe(memory_images=[torch.zeros(1, 3, 32, 32) for _ in range(8)], height=32, width=32)
def test_vae_encoder_runs_without_transformer(self):
pytest.importorskip("torchaudio")
pipe = EchoVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
output = pipe(
image=torch.rand(1, 3, 32, 32),
memory_images=[torch.rand(1, 3, 32, 32)],
memory_audio_waveforms=[torch.rand(2, 3200)],
memory_audio_sample_rates=16000,
height=32,
width=32,
output=["first_frame_latents", "memory_video_latents", "memory_audio_latents"],
)
assert set(pipe.components) == {"vae", "audio_vae"}
assert output["first_frame_latents"].shape == (1, 4, 1, 16, 16)
assert len(output["memory_video_latents"]) == 1
assert output["memory_audio_latents"][0].shape == (1, 2, 6, 2)
@pytest.mark.parametrize("dtype", [torch.float32, torch.bfloat16])
def test_audio_encoder_decoder_use_channel_mel_statistics(self, dtype):
pytest.importorskip("torchaudio")
encoder = EchoVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
encoder.load_components(dtype={"default": torch.float32, "audio_vae": dtype})
identity_vae = deepcopy(encoder.audio_vae)
identity_vae.latents_mean.zero_()
identity_vae.latents_std.fill_(1)
encoder.update_components(audio_vae=identity_vae)
inputs = {
"memory_images": [torch.rand(1, 3, 32, 32, generator=self.get_generator(0))],
"memory_audio_waveforms": [torch.rand(2, 3200, generator=self.get_generator(1))],
"memory_audio_sample_rates": 16000,
"height": 32,
"width": 32,
}
raw_latents = encoder(**inputs, output="memory_audio_latents")[0].to(dtype)
normalized_vae = deepcopy(identity_vae)
mean = torch.linspace(-0.5, 0.5, raw_latents.shape[1] * raw_latents.shape[3]).to(dtype)
std = torch.linspace(0.5, 1.5, mean.numel()).to(dtype)
normalized_vae.latents_mean.copy_(mean)
normalized_vae.latents_std.copy_(std)
encoder.update_components(audio_vae=normalized_vae)
normalized = encoder(**inputs, output="memory_audio_latents")[0]
# Statistics are stored in flattened (channel, mel-bin) order, independent of the time axis.
expected = (raw_latents.transpose(1, 2).flatten(2) - mean) / std
torch.testing.assert_close(normalized.transpose(1, 2).flatten(2), expected.float(), rtol=0, atol=0)
decoder = EchoAudioDecoderStep().init_pipeline(self.pretrained_model_name_or_path)
decoder.load_components(dtype=torch.float32)
decoder.update_components(audio_vae=normalized_vae)
decoded = decoder(audio_latents=normalized.to(dtype), output_type="latent", output="audio")
expected_decoded = (
(expected * std + mean).unflatten(2, (raw_latents.shape[1], raw_latents.shape[3])).transpose(1, 2)
)
torch.testing.assert_close(decoded, expected_decoded, rtol=0, atol=0)
def test_conditioning_uses_latent_geometry(self):
pipe = EchoPrepareConditioningStep().init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
video_latents = [torch.rand(1, 4, 1, 6, 10) for _ in range(3)]
audio_latents = [torch.rand(1, 2, 5, 2), None, torch.rand(1, 2, 7, 2)]
output = pipe(
memory_video_latents=video_latents,
memory_audio_latents=audio_latents,
output=["memory_video_tokens", "memory_video_coords", "memory_audio_tokens", "memory_audio_coords"],
)
assert set(pipe.components) == {"transformer"}
assert output["memory_video_tokens"].shape == (1, 180, 4)
assert output["memory_video_coords"].shape == (1, 3, 180, 2)
assert output["memory_audio_tokens"].shape == (1, 17, 4)
assert output["memory_audio_coords"].shape == (1, 1, 17, 2)
torch.testing.assert_close(output["memory_audio_tokens"][:, :5], audio_latents[0].transpose(1, 2).flatten(2))
assert torch.count_nonzero(output["memory_audio_tokens"][:, 5:10]) == 0
torch.testing.assert_close(output["memory_audio_tokens"][:, 10:], audio_latents[2].transpose(1, 2).flatten(2))
@pytest.mark.parametrize("num_videos", [1, 2])
@pytest.mark.parametrize("condition_batch", [1, 2])
def test_inputs_expand_cached_conditioning(self, num_videos, condition_batch):
pipe = EchoInputsStep().init_pipeline()
pipe.load_components()
inputs = {
"connector_prompt_embeds": torch.randn(2, 4, 8),
"connector_audio_prompt_embeds": torch.randn(2, 4, 4),
"connector_attention_mask": torch.tensor([[1, 1, 0, 0], [1, 1, 1, 0]]),
"first_frame_tokens": torch.randn(condition_batch, 6, 4),
"memory_video_tokens": torch.randn(condition_batch, 12, 4),
"memory_audio_tokens": torch.randn(condition_batch, 5, 4),
"memory_video_coords": torch.randn(1, 3, 12, 2),
"memory_audio_coords": torch.randn(1, 1, 5, 2),
}
originals = {name: value.clone() for name, value in inputs.items()}
# Run the same cached tensors repeatedly with different requested multiplicities.
for count in (num_videos, 1, 2):
output = pipe(**inputs, num_videos_per_prompt=count, output=list(inputs))
for name, value in inputs.items():
expected = value.repeat_interleave(2 if value.shape[0] == 1 else 1, dim=0)
torch.testing.assert_close(output[name], expected.repeat_interleave(count, dim=0))
torch.testing.assert_close(value, originals[name], rtol=0, atol=0)
def test_decoder_runs_with_only_vae_form_latents(self):
pipe = self.get_pipeline()
output = pipe(**self.get_dummy_inputs(), output=["latents", "audio_latents", "videos", "audio"])
assert output["latents"].shape == (1, 4, 3, 16, 16)
assert output["audio_latents"].shape == (1, 2, 5, 2)
decoder = EchoDecoderStep().init_pipeline(self.pretrained_model_name_or_path)
decoder.load_components(dtype=torch.float32)
decoded = decoder(
latents=output["latents"],
audio_latents=output["audio_latents"],
output_type="pt",
output=["videos", "audio"],
)
assert "transformer" not in decoder.components
torch.testing.assert_close(decoded["videos"], output["videos"], rtol=0, atol=0)
torch.testing.assert_close(decoded["audio"], output["audio"], rtol=0, atol=0)
def test_audio_memory_is_cropped_to_highest_response_window(self):
pytest.importorskip("torchaudio")
pipe = EchoVaeEncoderStep().init_pipeline(self.pretrained_model_name_or_path)
pipe.load_components(dtype=torch.float32)
sample_rate = pipe.audio_vae.config.sample_rate
waveform = torch.zeros(2, 12 * sample_rate)
waveform[:, -sample_rate:] = torch.randn(2, sample_rate, generator=self.get_generator(0))
inputs = {
"memory_images": [torch.rand(1, 3, 32, 32, generator=self.get_generator(1))],
"memory_audio_sample_rates": sample_rate,
"height": 32,
"width": 32,
}
output = pipe(**inputs, memory_audio_waveforms=[waveform], output="memory_audio_latents")[0]
# Only the last second contains sound, so the highest-response 9.62-second window ends at the tail.
cropped_waveform = waveform[:, -round(9.62 * sample_rate) :]
expected = pipe(**inputs, memory_audio_waveforms=[cropped_waveform], output="memory_audio_latents")[0]
mel_steps = cropped_waveform.shape[-1] // pipe.audio_vae.config.mel_hop_length + 1
assert output.shape[2] == (mel_steps - 1) // pipe.audio_vae.temporal_compression_ratio + 1
torch.testing.assert_close(output, expected, rtol=0, atol=0)
class TestEchoModularPipelineLoading(EchoModularPipelineTesterConfig, ModularLoadingTesterMixin):
pass
class TestEchoModularPipelineWorkflow(EchoModularPipelineTesterConfig, ModularWorkflowTesterMixin):
pass
class TestEchoModularPipelineMemory(EchoModularPipelineTesterConfig, ModularMemoryTesterMixin):
pass

64
tests/modular_pipelines/echo/testing_utils.py Normal file
View File

@@ -0,0 +1,64 @@
# coding=utf-8
# Copyright 2026 HuggingFace Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
from diffusers.modular_pipelines import EchoBlocks, EchoModularPipeline
from ..testing_utils import BaseModularPipelineTesterConfig
class EchoModularPipelineTesterConfig(BaseModularPipelineTesterConfig):
pipeline_class = EchoModularPipeline
pipeline_blocks_class = EchoBlocks
pretrained_model_name_or_path = "Echo-Team/tiny-echo-modular-pipe"
params = frozenset(
[
"prompt",
"image",
"memory_images",
"memory_audio_waveforms",
"height",
"width",
"num_frames",
"frame_rate",
"model_frame_rate",
"sigmas",
]
)
batch_params = frozenset(["prompt"])
optional_params = frozenset(["num_videos_per_prompt", "latents", "audio_latents", "output_type"])
not_params = frozenset(["negative_prompt", "guidance_scale", "num_inference_steps"])
expected_workflow_blocks = {}
output_name = "videos"
def get_dummy_inputs(self, seed=0):
generator = self.get_generator(seed)
return {
"prompt": "a robot dancing",
"image": torch.rand((1, 3, 32, 32), generator=generator),
"memory_images": [
torch.rand((1, 3, 32, 32), generator=generator),
torch.rand((1, 3, 32, 32), generator=generator),
],
"generator": self.get_generator(seed),
"sigmas": [1.0, 0.0],
"height": 32,
"width": 32,
"num_frames": 5,
"frame_rate": 25.0,
"max_sequence_length": 16,
"output_type": "pt",
}