Motivation.

JoyAI-Echo (https://github.com/jd-opensource/JoyAI-Echo) is the recently open-sourced JD model for minute-scale multi-shot joint text-to-video + audio generation. It is built on top of LTX-2.3 but introduces three capabilities that are not directly supported by the existing LTX23Pipeline family in vLLM-Omni:

1. Joint A/V generation — a single transformer forward produces video and audio latents that are then decoded by separate VAEs (no cascaded "video then dub" pipeline).
2. Multi-shot long-video generation — a PairedAudioVideoMemoryBank conditions later shots on the audio + video latents of earlier shots so that the same scene/character/voice can persist for ~minute-scale clips.
3. DMD few-step distilled inference — 8-step (9-sigma) schedule out of the box, ~75 s per 5-second clip on a single B200 (bf16,480x832).

JoyAI-Echo is published under the LTX-2 Community License (the same one as upstream LTX-2.3 weights; vLLM-Omni already supports inference under that license for LTX23Pipeline). The Hub repo is jdopensource/JoyAI-Echo.

The reasons we cannot simply subclass LTX23Pipeline are summarised in the Proposed Change below.

Proposed Change.

1. Overview

This RFC proposes adding JoyAI-Echo as a new diffusion pipeline, JoyAIEchoPipeline, registered under vllm_omni/diffusion/models/joyai_echo/. It is the first vLLM-Omni pipeline to natively support paired text-to-video + audio with multi-shot memory conditioning, packaged on top of vLLM-Omni's existing diffusion engine, scheduler, and CPU-offload infrastructure.

The work is split into 3 sequential PRs so each one is reviewable on its own:

PR	Scope	Status
PR1	Single-shot T2V+Audio + ltx_core port + L4 functionality test	Ready (this RFC's first PR, verified on B200)
PR2	Multi-shot PairedAudioVideoMemoryBank + warmup re-enable	Planned
PR3	Online serving + recipe + numerical-fidelity test + Attention(role) migration	Planned

Related context:

• Reference implementation: https://github.com/jd-opensource/JoyAI-Echo
• License: LTX-2 Community License (vLLM-Omni already runs LTX23Pipeline weights under the same license).
• Closest neighbours in-tree: LTX23Pipeline, Wan2_2S2VPipeline (audio output, dummy-warmup-skipped pattern), BagelPipeline (eager ModelLedger weight loading, weights_sources=[]).

2. Scope & Objectives

Goals

#	Goal	Validation
G1	Single-shot text → joint video + audio generation on real Hub weights	L4 functionality test at 480x832x121 frames, 8 DMD steps; asserts video tensor (T,H,W,3) uint8 and 24 kHz audio waveform
G2	Bit-for-bit-faithful port of the upstream ltx_core subtree (transformer / video_vae / audio_vae / Gemma-3 text_encoder / loader) — no algorithmic re-design	bf16 same-seed PSNR ≥ 30 dB vs JoyAI-Echo/inference.py (PR3 deliverable)
G3	Compatibility with transformers >= 5 Gemma3 (the upstream codebase was on 4.57)	Existing fixes in the PR1 branch: vision_tower.vision_model.* SDOps flatten + per-layer-type Gemma3RotaryEmbedding buffer init
G4	Multi-shot long-video generation (1–5 shots, prefix-continuous memory)	PR2 deliverable: dedicated multi-shot smoke test
G5	OpenAI-compatible online serving	PR3 deliverable: examples/online_serving/joyai_echo/{run_server.sh, openai_chat_client.py}
G6	Honour vLLM-Omni's --diffusion-attention-config.per_role.* on every attention site	PR3 deliverable: migrate to Attention(role="self"/"cross"/"joyai_echo.audio_to_video")
G7	CPU-offload buckets via SupportsComponentDiscovery	Already in PR1

Non-Goals

#	Out of scope (and why)
NG1	LoRA fuse path. Upstream supports it via apply_loras; we ship the code but do not expose it through od_config. Re-enable as a follow-up once vLLM-Omni's LoRA manager is plumbed through.
NG2	TensorRT-LLM FP8 quantisation kernels (ltx_core/quantization/fp8_*.py). Imported but never invoked; vLLM-Omni already has a generic FP8 path that PR3+ can adopt.
NG3	Image-to-video. The upstream model has an I2V config but the released checkpoint is T2V-only.
NG4	torch.compile of the transformer block. _repeated_blocks is not declared, so compile is skipped (warned at startup). PR3 can add this once attention sites are wired to Attention(role).
NG5	Tensor / sequence / CFG parallelism. Single-GPU only on PR1. Multi-GPU support gated on the attention migration (PR3).
NG6	New kernels. We reuse vLLM-Omni's existing rms_norm / fused_add_rms_norm IR ops; nothing new lands in vllm_omni/kernel/.

3. Design

3.1 Architecture

Top-level component layout (PR1 single-shot path):

+-------------------------------------------------------------------------+
|                         vLLM-Omni Diffusion Engine                      |
|                                                                         |
|   OmniDiffusionRequest  -->  pre_process_func (none)  -->  pipeline     |
|                                                                         |
|   +-------------------------------------------------------------+       |
|   |  JoyAIEchoPipeline (vllm_omni/diffusion/models/joyai_echo/) |       |
|   |                                                             |       |
|   |    __init__ : ModelLedger.load() — eager safetensors load   |       |
|   |               (Bagel-style; weights_sources = [])           |       |
|   |                                                             |       |
|   |    forward(req):                                            |       |
|   |      1. encode_prompt   (Gemma3 + EmbeddingsProcessor)      |       |
|   |      2. compute_latent_shapes (video_vae + audio_vae)       |       |
|   |      3. BidirectionalAVInferencePipeline.generate           |       |
|   |           -> 8-step DMD few-step denoise                    |       |
|   |           -> LTX2DiffusionWrapper.forward                   |       |
|   |      4. video_vae.decode  + audio_vae.decode                |       |
|   |      5. return DiffusionOutput((video_uint8, audio))        |       |
|   |                                                             |       |
|   |   SupportsComponentDiscovery:                               |       |
|   |     dit      = [transformer]                                |       |
|   |     encoder  = [text_encoder]                               |       |
|   |     vae      = [video_vae, audio_vae]                       |       |
|   +-------------------------------------------------------------+       |
|                                                                         |
|   post_process_func: tuple -> {"video", "audio", "audio_sample_rate"}   |
+-------------------------------------------------------------------------+
                                |
                                v
   OmniRequestOutput(images=[video_uint8], multimodal_output={"audio":..})

PR2 layers a PairedAudioVideoMemoryBank between steps (3) and (4) above that threads the previous shot's video+audio latents into the next shot's transformer call as memory_video / paired_audio_memory arguments.

PR3 swaps the attention sites inside LTX2DiffusionWrapper's transformer blocks for vllm_omni.diffusion.attention.layer.Attention(role=...) so that the engine's --diffusion-attention-config.per_role.* flags reach each site.

3.2 API & Interface Changes

New public surface (PR1)

Path	Symbol	Purpose
vllm_omni/diffusion/models/joyai_echo/__init__.py	JoyAIEchoPipeline, get_joyai_echo_post_process_func	Pipeline + post-process factory
vllm_omni/diffusion/models/joyai_echo/pipeline_joyai_echo.py	JoyAIEchoPipeline(nn.Module, SupportsComponentDiscovery, ProgressBarMixin)	Single-shot forward + CPU-offload buckets
examples/offline_inference/joyai_echo/end2end.py	CLI entry point	Offline reference command
tests/e2e/offline_inference/test_joyai_echo_expansion.py	test_joyai_echo_single_shot_t2v_audio	L4 functionality test

Existing files modified

File	Change
vllm_omni/diffusion/registry.py	Two entries: _DIFFUSION_MODELS["JoyAIEchoPipeline"] = ("joyai_echo", "pipeline_joyai_echo", "JoyAIEchoPipeline") and _DIFFUSION_POST_PROCESS_FUNCS["JoyAIEchoPipeline"] = "get_joyai_echo_post_process_func"
docs/models/supported_models.md	New row for JoyAIEchoPipeline (T2V+Audio, single-shot, repo jdopensource/JoyAI-Echo)

Invariants & compatibility

• Public engine API unchanged. No new fields in OmniDiffusionRequest, OmniDiffusionSamplingParams, OmniRequestOutput, or the DiffusionPipeline protocol. Output flows through the same images=[video_uint8] + multimodal_output={audio, audio_sample_rate} contract that Wan2_2S2VPipeline already uses.
• Registry-only opt-in. No code path is reachable unless the user passes model_class_name="JoyAIEchoPipeline" (or autodetect resolves the same value from model_index.json).
• load_weights() returns None. Intentional — opts out of the strict weights_not_loaded audit in DiffusersPipelineLoader.load_weights because the state dict is loaded eagerly via ModelLedger and is not visible to AutoWeightsLoader. This is the same idiom used by Bagel.

3.3 Key Technical Decisions

D1. Parallel top-level directory rather than ltx2/joyai_echo/ subclass

• Decision — land under vllm_omni/diffusion/models/joyai_echo/ as a parallel implementation, not as a subclass of LTX23Pipeline.
• Alternatives — (a) subclass LTX23Pipeline and override the multi-shot bits; (b) re-use the existing ltx2/ transformer with a thin memory_inject mixin.
• Rationale — four divergence points each by themselves justify a parallel tree:
  1. Checkpoint format: JoyAI-Echo ships a single ComfyUI-style JoyAI-Echo-release.safetensors; ltx2/ expects diffusers-format component subfolders. The SD-loader rewrite alone is non-trivial.
  2. Transformer interface: JoyAI-Echo takes paired Modality objects (video+audio) plus memory_video / paired_audio_memory. The ltx2/ transformer signature has none of those.
  3. Text encoder chain: in-house EmbeddingsProcessor plus two Embeddings1DConnectors on top of Gemma-3-12B, vs single-headed diffusers chain in ltx2/.
  4. Multi-shot: no equivalent of PairedAudioVideoMemoryBank in the existing tree; bolting it on would warp ltx2/ for a model it was not designed for.
• Trade-off — the port adds ~16k LOC and duplicates ~3k LOC's worth of LTX-2.3 building blocks. We accept that in exchange for keeping ltx2/ stable for its current users; PR3 is the right time to look at code de-duplication once the attention migration lands.

D2. Bagel-style eager weight load via ModelLedger

• Decision — __init__ runs ModelLedger.load() (driven by module_ops + SDOps) and load_weights() is a no-op returning None.
• Alternatives — (a) implement weights_sources as a list of DiffusersPipelineLoader.ComponentSource per submodule; (b) subclass DiffusersAdapterPipeline.
• Rationale — JoyAI-Echo's checkpoint is a single ComfyUI-style monolith plus 5 Gemma shards; DiffusersPipelineLoader expects per-component subfolders. ModelLedger is the codepath the upstream reference already uses (Builder + SDOps + ModuleOps), so a faithful port stays mechanically equivalent.
• Risk — bypasses the strict weights_not_loaded audit, so a mismatched checkpoint would surface as a NaN at inference instead of an explicit error. Mitigation: _return_model warns on residual meta parameters; the L4 test catches mismatches end-to-end.

D3. dummy_run_num_frames = 0

• Decision — skip the diffusion-engine warmup call.
• Alternatives — (a) keep the upstream default warmup at 9 frames / 512x512; (b) ship a mini-warmup with a paired memory_audio of the right shape.
• Rationale — the per-token AdaLN modulation requires video.timesteps to broadcast against the live video token grid. At the warmup canvas (512x512x9 frames) the audio-derived per-token timestep length (1840) does not equal the video token count (512), tripping the AdaLN multiply. The same root cause is hit by Wan2_2S2VPipeline, which also sets dummy_run_num_frames = 0. PR2 re-enables a proper paired warmup once the multi-shot shape inference is in place.
• Trade-off — no GPU memory profile run on PR1, so OOM probes are deferred until the first real request. Acceptable for the single-shot smoke path; not acceptable for online serving (covered in PR3).

D4. Migrate attention sites in PR3, not PR1

• Decision — PR1 keeps the original LTX2DiffusionWrapper.forward attention path; the migration to Attention(role=...) lands in PR3.
• Alternatives — do the migration in PR1 to satisfy the diffusion-model contributing guide on day one.
• Rationale — PR1 is already 16k LOC of port work; introducing the attention rewrite in the same PR makes regression triage harder and blocks a clean numerical-fidelity comparison. PR3 owns the attention rewrite together with the same-seed PSNR test that proves the rewrite is faithful.
• Cost on the user side — until PR3, JoyAI-Echo cannot honour --diffusion-attention-config.per_role overrides. Documented in the PR1 description as a known limitation.

D5. Wrapper consumes live latent shape

• Decision — LTX2DiffusionWrapper._compute_timesteps_for_tokens and _unflatten_video_latent derive tokens_per_frame and H,W from noisy_image_or_video.shape[-2:], not from self.video_frame_seqlen / self.latent_height / self.latent_width.
• Alternatives — instantiate a fresh wrapper per request.
• Rationale — the upstream wrapper is constructed once with default video_height=736 / video_width=1280, but vLLM-Omni drives it with arbitrary per-request resolutions. Making the per-token grid live fixes the AdaLN shape mismatch at any resolution; the alternative would re-allocate the transformer module tree on every request.

3.4 Dependencies & Risks

Dep	Version pin	Notes
transformers	already >= 5 in vllm-omni	We had to backport the Gemma3 vision_tower / RotaryEmbedding compatibility shims (D3 Compatibility fixes section).
torch	follows vllm 0.22 (currently 2.11+cu130 on B200; tested with 2.12+cu128 earlier)	LTX VAE custom ops use only stable Tensor APIs; no torch-version-specific kernels.
safetensors	>= 0.7	Multi-shard load via safe_open.
aenum, diffusers, torchaudio, torchvision	already in vllm-omni's lock	Used by VAE encoder/decoder and write_video helper.
av (PyAV)	already used by vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes	Used by the offline example to mux mp4.

Risks:

Risk	Likelihood	Mitigation
Upstream transformers continues to refactor Gemma3 internals	medium	All compatibility shims are guarded by hasattr(...) / layer_types is not None so the older path stays viable; CI matrix will catch a drift.
16k LOC PR is too large for a single review	high	RFC-level offer: split PR1 into "[Misc][Vendor] port ltx_core" + "[Model] add JoyAIEchoPipeline" if a reviewer asks.
Bagel-style eager load conflates init with weight I/O	low	Same idiom shipped in BagelPipeline. PR3 can revisit if the team prefers the weights_sources path once DiffusersPipelineLoader supports monolithic safetensors.
LTX-2 Community License compatibility with Apache-2.0 vLLM-Omni	low	Same license already runs in-tree for LTX23Pipeline weights. SPDX headers on every new file: Apache-2.0 + Adapted from JoyAI-Echo (LTX-2 Community License).

3.5 Performance

PR1 baseline measured on a single NVIDIA B200 (CUDA 13.0, Driver 580.82.07, torch 2.11+cu130, vllm 0.22.0, vllm-omni @ HEAD; bf16):

Setting	Value
resolution	480 × 832
frames	121 (≈ 4.84 s @ 25 fps)
audio	9.62 s @ 24 kHz
DMD steps	8
init time	≈ 28 s
diffusion time	≈ 75 s
peak GPU mem	39 GB
output mp4	1.04 MB

Expected PR3 wins from the attention migration:

• --diffusion-attention-config.per_role.self.backend FLASH_ATTN_3 (or SAGE_ATTN) on the dual-stream block — Wan2.2 saw ~1.4× on B200.
• Optional _repeated_blocks declaration once the block boundary is re-discoverable through Attention(role) — unlocks regional torch.compile.

Known bottlenecks today (PR1):

• Two sequential _move(...) between transformer and VAE-decode stages add ~1–2 s of D2H/H2D traffic on the 12B transformer; PR3 can fold this into vLLM-Omni's standard CPU-offload scheduler.
• No torch.compile (NG4); each DMD step pays full eager overhead.

4. Correctness & Testing plans

Definition of "correct"

Level	Invariant
Per-component	LTX2DiffusionWrapper.forward(latent[B,F,C,H,W], …) returns a tensor whose T = F*H*W and C=128; AdaLN multiply broadcasts cleanly at any (H,W) where H*W matches the live latent.
Per-pipeline	JoyAIEchoPipeline.forward(req) returns DiffusionOutput(output=(video_uint8 [T,H,W,3], audio_waveform)) with T == num_frames, H,W == request H,W, audio sample rate 24000 Hz.
Numerical (PR3)	bf16, same seed, same prompt: PSNR ≥ 30 dB on the first 16 frames vs python /JoyAI-Echo/inference.py.

Test pyramid

Level	File	Hardware	Status
L4 functionality	tests/e2e/offline_inference/test_joyai_echo_expansion.py	B200 cuda	PR1
L1 unit (port-level)	tests/diffusion/models/joyai_echo/test_sdops.py (rename mappings) + test_model_ledger.py (build under meta + populate buffers)	CPU	PR3
L2 smoke	tests/e2e/offline_inference/test_joyai_echo_base.py (smaller canvas, 1 step)	L4 cuda	PR3
L4 numerical fidelity	tests/e2e/accuracy/test_joyai_echo_psnr.py	B200 cuda	PR3

PR1 also ships:

• examples/offline_inference/joyai_echo/end2end.py — manually exercised at the resolutions in §3.5; mp4 verified with ffprobe.
• pre-commit run is clean on every file touched.

5. Open Questions & Discussions

#	Question
Q1	Is a 16k-LOC PR1 acceptable, or should it be split into "[Misc][Vendor] port ltx_core" + "[Model] add JoyAIEchoPipeline"? Most of the LOC is the verbatim ltx_core port.
Q2	Is the load_weights() -> None opt-out from weights_not_loaded audit the preferred idiom for "ModelLedger handled it eagerly", or should we evolve DiffusersPipelineLoader to recognise a pipeline.weights_loaded_eagerly = True marker?
Q3	dummy_run_num_frames = 0 mirrors Wan2_2S2VPipeline but means we lose the warmup OOM probe. Is that acceptable on PR1 given that PR2 re-enables it, or should PR1 ship a paired-tensor mini-warmup?
Q4	The Attention(role) migration needs at least one model-specific role (joyai_echo.audio_to_video for the cross-modality attention). Should the role string be namespaced under the model dir name or under a more generic family name like ltx2.audio_to_video?
Q5	For the multi-shot PairedAudioVideoMemoryBank in PR2, should the per-shot prompt list be carried via OmniDiffusionRequest.prompts as a Sequence[OmniTextPrompt] (current shape) or via an extension field on OmniDiffusionSamplingParams.extra_args?

6. References

• Upstream: https://github.com/jd-opensource/JoyAI-Echo
• Upstream model card / weights: https://huggingface.co/jdopensource/JoyAI-Echo
• LTX-2 Community License: shipped in /home/.../JoyAI-Echo/checkpoints/JoyAI-Echo/LICENSE (and via THIRD_PARTY_NOTICES.md in this PR).
• In-tree neighbours studied for design alignment:
  • vllm_omni/diffusion/models/ltx2/ — diffusers-format LTX-2.3.
  • vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_s2v.py — dummy_run_num_frames = 0 precedent and audio-output handling.
  • vllm_omni/diffusion/models/bagel/ — eager ModelLedger-style load
    • weights_sources = [].
• Contributing docs followed:
  • https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/model/adding_diffusion_model/
  • https://docs.vllm.ai/projects/vllm-omni/en/latest/contributing/
• Branch under review: https://github.com/jefferrrrrrrrrrry/vllm-omni/tree/joyai-echo-pr1 (commit ca085f9b, 110 files / 16,256 LOC, signed off).

Feedback Period.

One week from posting (or up to PR1 review start, whichever is sooner).

CC List.

@claude @Yikun @russellb @qibaoyuan — happy to take review pointers from anyone familiar with LTX23Pipeline / Wan2_2S2VPipeline / Bagel's eager-load pattern.

Any Other Things.

• PR1 is already commit-ready on https://github.com/jefferrrrrrrrrrry/vllm-omni/tree/joyai-echo-pr1 (110 files / ~16.2k LOC). Verified end-to-end on a single B200(CUDA 13.0): 121 frames @ 480x832 + 9.62 s of 24 kHz audio →1.04 MB mp4 in ~75 s, 39 GB peak.
• The PR carries explicit "known limitations" pointing to PR2/PR3 (no Attention(role) wiring yet, dummy warmup skipped at small canvases, etc.) so reviewers can scope cleanly.
• Happy to split PR1 further if 16k LOC is too large to land in onego — most of the LOC is the verbatim ltx_core port. Suggestion:could land ltx_core/ (port-only) in a "[Misc][Vendor]" PR first and then JoyAIEchoPipeline (registry + pipeline + example + test) in a smaller "[Model]" PR.

Before submitting a new issue...

• Make sure you already searched for relevant issues, and asked the chatbot living at the bottom right corner of the documentation page, which can answer lots of frequently asked questions.