Repository navigation
Add Echo modular pipeline #14696
New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add Echo modular pipeline #14696
Changes from 3 commits
a256a7b
64f423a
679db3f
9d182ea
90d0c3d
a613668
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
| Original file line number | Diff line number | Diff line change | ||||||||
|---|---|---|---|---|---|---|---|---|---|---|
| @@ -0,0 +1,106 @@ | ||||||||||
| <!-- 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 adds 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 Modular Pipeline so its text encoding, memory conditioning, stochastic DMD denoising, and | ||||||||||
| decoding blocks can be run as a complete workflow or composed independently. | ||||||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||||
|
|
||||||||||
| ## Convert the checkpoint | ||||||||||
|
sayakpaul marked this conversation as resolved.
Outdated
|
||||||||||
|
|
||||||||||
| Convert the BF16 Echo release checkpoint before loading it. The converter reuses the Gemma text encoder and | ||||||||||
| tokenizer directly from `google/gemma-3-12b-it`. | ||||||||||
|
|
||||||||||
| ```bash | ||||||||||
| python scripts/convert_echo_to_diffusers.py \ | ||||||||||
| --checkpoint /path/to/echo15_full_dmd \ | ||||||||||
| --output-path /path/to/Echo-Diffusers \ | ||||||||||
| --repo-id Echo-Team/JoyAI-Echo | ||||||||||
| ``` | ||||||||||
|
|
||||||||||
| The Gemma repository is gated, so users must accept its license and authenticate with Hugging Face before loading the | ||||||||||
| pipeline. Pass a different `--base-model` only when the compatible Gemma model and tokenizer are stored together at | ||||||||||
| that repository or path root. `--repo-id` records portable Hub references for the converted Echo components; without | ||||||||||
| it, the index targets the local output path. | ||||||||||
|
|
||||||||||
| ## Inference | ||||||||||
|
|
||||||||||
| 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 | ||||||||||
|
|
||||||||||
|
|
||||||||||
| model_path = "/path/to/Echo-Diffusers" | ||||||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. We should be able to load from the official checkpoint. Please prefer doing so. None of our official pipelines do this.
Contributor
Author
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Done in a613668. The example now uses
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Did you see if the pipeline runs and produces expected output? |
||||||||||
| manager = ComponentsManager() | ||||||||||
| pipe = ModularPipeline.from_pretrained(model_path, 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 | ||||||||||
|
Member
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
Suggested change
|
||||||||||
| 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 and reuse the condition encoder's packed token outputs by running that block separately. | ||||||||||
|
|
||||||||||
| ## EchoModularPipeline | ||||||||||
|
|
||||||||||
| [[autodoc]] EchoModularPipeline | ||||||||||
|
|
||||||||||
| ## EchoBlocks | ||||||||||
|
|
||||||||||
| [[autodoc]] EchoBlocks | ||||||||||
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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, | ||
| ) |
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -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) |
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.