minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[diffusion] model: support JoyEcho multi-shot A/V generation support (#27420)

Browse Source 
Co-authored-by: niehen6174 <niehen6174@users.noreply.github.com>
Co-authored-by: 1639206518@qq.com <niehen6174>
This commit is contained in:
WenhaoZhang
2026-06-26 15:46:31 +08:00
committed by GitHub
co-authored by niehen6174 1639206518@qq.com
parent dd56a9f069
commit aeb4e98108
20 changed files with 2265 additions and 24 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/sglang/multimodal_gen/configs/models/dits/joy_echo.py
+22
View File
@@ -0,0 +1,22 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from sglang.multimodal_gen.configs.models.dits.ltx_2 import (
LTX2ArchConfig,
LTX2Config,
)
@dataclass
class JoyEchoArchConfig(LTX2ArchConfig):
"""JoyEcho DiT architecture config (LTX-2.3 AV base)."""
caption_proj_before_connector: bool = True
cross_attention_adaln: bool = True
apply_gated_attention: bool = True
@dataclass
class JoyEchoConfig(LTX2Config):
arch_config: JoyEchoArchConfig = field(default_factory=JoyEchoArchConfig)
prefix: str = "JoyEcho"
python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py
+2
View File
@@ -55,6 +55,8 @@ class LTXVideoVAEArchConfig(VAEArchConfig):
# Native LTX variant metadata.
ltx_variant: str = "ltx_2"
condition_encoder_subdir: str = ""
video_encoder_variant: str = "ltx_2"
video_encoder_config: dict[str, Any] = field(default_factory=dict)
video_decoder_variant: str = "ltx_2"
video_decoder_config: dict[str, Any] = field(default_factory=dict)
python/sglang/multimodal_gen/configs/pipeline_configs/joy_echo.py
+60
View File
@@ -0,0 +1,60 @@
# SPDX-License-Identifier: Apache-2.0
from dataclasses import dataclass, field
from typing import Optional
from sglang.multimodal_gen.configs.models.dits.joy_echo import JoyEchoConfig
from sglang.multimodal_gen.configs.models.vaes.ltx_video import LTXVideoVAEConfig
from sglang.multimodal_gen.configs.pipeline_configs.base import ModelTaskType
from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import LTX2PipelineConfig
def _default_joy_echo_vae_config() -> LTXVideoVAEConfig:
vae_config = LTXVideoVAEConfig()
vae_config.arch_config.ltx_variant = "ltx_2_3"
return vae_config
JOY_ECHO_DEFAULT_SIGMAS: tuple[float, ...] = (
1.0,
0.99375,
0.9875,
0.98125,
0.975,
0.909375,
0.725,
0.421875,
0.0,
)
@dataclass
class JoyEchoPipelineConfig(LTX2PipelineConfig):
"""Pipeline configuration for JoyEcho long-video generation."""
task_type: ModelTaskType = ModelTaskType.T2V
dit_config: JoyEchoConfig = field(default_factory=JoyEchoConfig)
vae_config: LTXVideoVAEConfig = field(default_factory=_default_joy_echo_vae_config)
monolithic_checkpoint: Optional[str] = None
gemma_model_path: str = "google/gemma-3-12b-it"
default_sigmas: tuple[float, ...] = field(
default_factory=lambda: JOY_ECHO_DEFAULT_SIGMAS
)
enable_memory_bank: bool = True
memory_max_size: int = 7
memory_num_fix_frames: int = 3
memory_position_mode: str = "reference"
audio_window_size: int = 96
audio_mel_bins: int = 128
audio_mel_hop_length: int = 160
audio_n_fft: int = 1024
audio_downsample_factor: int = 4
audio_window_selection_mode: str = "max_response"
memory_video_clip_num_frames: int = 9
video_memory_frame_selection_mode: str = "center"
late_layer_ratio: float = 0.7
python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py
+2
View File
@@ -136,6 +136,8 @@ def sync_ltx23_runtime_vae_markers(
for key in (
"ltx_variant",
"condition_encoder_subdir",
"video_encoder_variant",
"video_encoder_config",
"video_decoder_variant",
"video_decoder_config",
):
python/sglang/multimodal_gen/configs/sample/joy_echo.py
+34
View File
@@ -0,0 +1,34 @@
# SPDX-License-Identifier: Apache-2.0
import dataclasses
from dataclasses import field
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
JOY_ECHO_DEFAULT_SIGMAS,
)
from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams
@dataclasses.dataclass
class JoyEchoSamplingParams(LTX2SamplingParams):
"""Sampling parameters for JoyEcho DMD inference."""
seed: int = 12345
generator_device: str = "cuda"
height: int = 480
width: int = 832
num_frames: int = 121
fps: int = 25
guidance_scale: float = 1.0
num_inference_steps: int = 8
sigmas: tuple[float, ...] = field(default_factory=lambda: JOY_ECHO_DEFAULT_SIGMAS)
negative_prompt: str | None = None
video_cfg_scale: float = 1.0
audio_cfg_scale: float = 1.0
enable_memory_bank: bool = True
reset_memory_bank: bool = True
python/sglang/multimodal_gen/registry.py
+15
View File
@@ -60,6 +60,9 @@ from sglang.multimodal_gen.configs.pipeline_configs.hunyuan3d import (
from sglang.multimodal_gen.configs.pipeline_configs.ideogram import (
Ideogram4PipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
JoyEchoPipelineConfig,
)
from sglang.multimodal_gen.configs.pipeline_configs.joy_image import (
JoyImageEditPipelineConfig,
)
@@ -114,6 +117,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import (
)
from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams
from sglang.multimodal_gen.configs.sample.ideogram import Ideogram4SamplingParams
from sglang.multimodal_gen.configs.sample.joy_echo import JoyEchoSamplingParams
from sglang.multimodal_gen.configs.sample.joy_image import (
JoyImageEditSamplingParams,
)
@@ -1060,6 +1064,17 @@ def _register_configs():
lambda hf_id: "joyai-image-edit" in hf_id.lower(),
],
)
register_configs(
sampling_param_cls=JoyEchoSamplingParams,
pipeline_config_cls=JoyEchoPipelineConfig,
hf_model_paths=[
"jdopensource/JoyAI-Echo",
],
model_detectors=[
lambda hf_id: ("joy-echo" in hf_id.lower() or "joyai-echo" in hf_id.lower())
and "image-edit" not in hf_id.lower(),
],
)
# Ideogram 4
register_configs(
python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py
+49 -7
View File
@@ -663,6 +663,7 @@ class LTX2Attention(nn.Module):
all_perturbed: bool = False,
skip_sequence_parallel_override: bool = False,
gather_context_kv_for_sp: bool = False,
context_replicated_prefix_len: int = 0,
) -> torch.Tensor:
gate_input = x
context_ = x if context is None else context
@@ -703,13 +704,38 @@ class LTX2Attention(nn.Module):
k = k.view(*k.shape[:-1], self.local_heads, self.dim_head)
if gather_context_kv_for_sp:
k_full = sequence_model_parallel_all_gather(k.contiguous(), dim=1)
v_full = sequence_model_parallel_all_gather(v.contiguous(), dim=1)
gathered_mask = None
if mask is not None:
gathered_mask = sequence_model_parallel_all_gather(
mask.contiguous(), dim=1
# Replicated prefix (e.g. JoyEcho memory) is identical on every rank; only gather the sharded suffix.
if context_replicated_prefix_len > 0:
prefix = int(context_replicated_prefix_len)
k_prefix, k_suffix = k[:, :prefix], k[:, prefix:]
v_prefix, v_suffix = v[:, :prefix], v[:, prefix:]
k_full = torch.cat(
[
k_prefix,
sequence_model_parallel_all_gather(
k_suffix.contiguous(), dim=1
),
],
dim=1,
)
v_full = torch.cat(
[
v_prefix,
sequence_model_parallel_all_gather(
v_suffix.contiguous(), dim=1
),
],
dim=1,
)
gathered_mask = mask
else:
k_full = sequence_model_parallel_all_gather(k.contiguous(), dim=1)
v_full = sequence_model_parallel_all_gather(v.contiguous(), dim=1)
gathered_mask = None
if mask is not None:
gathered_mask = sequence_model_parallel_all_gather(
mask.contiguous(), dim=1
)
if self.use_local_attention:
out = self.attn(q, k_full, v_full, attn_mask=gathered_mask)
else:
@@ -1009,6 +1035,7 @@ class LTX2TransformerBlock(nn.Module):
a2v_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
v2a_cross_attn_perturbation_mask: Optional[torch.Tensor] = None,
audio_replicated_for_sp: bool = False,
video_memory_prefix_len: int = 0,
) -> tuple[torch.Tensor, torch.Tensor]:
batch_size = hidden_states.size(0)
@@ -1027,6 +1054,7 @@ class LTX2TransformerBlock(nn.Module):
perturbation_mask=video_self_attn_perturbation_mask,
all_perturbed=skip_video_self_attn,
gather_context_kv_for_sp=audio_replicated_for_sp,
context_replicated_prefix_len=video_memory_prefix_len,
)
hidden_states = hidden_states + attn_hidden_states * vgate_msa
@@ -1213,6 +1241,7 @@ class LTX2TransformerBlock(nn.Module):
k_pe=ca_video_rotary_emb,
mask=v2a_cross_attention_mask,
gather_context_kv_for_sp=audio_replicated_for_sp,
context_replicated_prefix_len=video_memory_prefix_len,
)
if v2a_cross_attn_perturbation_mask is not None:
v2a_attn_hidden_states = (
@@ -1639,6 +1668,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
disable_a2v_cross_attn: bool = False,
disable_v2a_cross_attn: bool = False,
audio_replicated_for_sp: bool = False,
video_memory_prefix_len: int = 0,
late_layer_ratio: float = 1.0,
late_audio_self_attention_mask: Optional[torch.Tensor] = None,
**kwargs,
) -> tuple[torch.Tensor | None, torch.Tensor | None]:
@@ -1873,6 +1905,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
audio_hidden_states,
)
)
late_layer_start = int(len(self.transformer_blocks) * float(late_layer_ratio))
for block in self.transformer_blocks:
block_idx = getattr(block, "idx", -1)
video_self_attn_perturbation_mask = None
@@ -1883,6 +1916,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
skip_audio_self_attn = block_idx in skip_audio_self_attn_blocks
skip_a2v_cross_attn = disable_a2v_cross_attn
skip_v2a_cross_attn = disable_v2a_cross_attn
block_audio_self_attention_mask = audio_self_attention_mask
if (
block_idx >= late_layer_start
and late_audio_self_attention_mask is not None
):
block_audio_self_attention_mask = late_audio_self_attention_mask
elif block_idx >= late_layer_start and late_layer_ratio < 1.0:
block_audio_self_attention_mask = None
if perturbation_configs is not None:
if not skip_video_self_attn:
assert video_self_attn_perturbation_states is not None
@@ -1923,7 +1964,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
encoder_attention_mask=encoder_attention_mask,
audio_encoder_attention_mask=audio_encoder_attention_mask,
video_self_attention_mask=video_self_attention_mask,
audio_self_attention_mask=audio_self_attention_mask,
audio_self_attention_mask=block_audio_self_attention_mask,
a2v_cross_attention_mask=a2v_cross_attention_mask,
v2a_cross_attention_mask=v2a_cross_attention_mask,
skip_video_self_attn=skip_video_self_attn,
@@ -1935,6 +1976,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
a2v_cross_attn_perturbation_mask=a2v_cross_attn_perturbation_mask,
v2a_cross_attn_perturbation_mask=v2a_cross_attn_perturbation_mask,
audio_replicated_for_sp=audio_replicated_for_sp,
video_memory_prefix_len=video_memory_prefix_len,
)
# 6. Output layers
python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py
+50 -16
View File
@@ -1647,26 +1647,45 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE):
config.arch_config, "timestep_conditioning", False
)
use_ltx23_video_decoder = (
str(getattr(config.arch_config, "video_decoder_variant", "ltx_2"))
== "ltx_2_3"
str(config.arch_config.video_decoder_variant) == "ltx_2_3"
)
use_ltx23_condition_encoder = (
str(config.arch_config.video_encoder_variant) == "ltx_2_3_condition"
)
self._use_ltx23_condition_encoder = use_ltx23_condition_encoder
decoder_causal = config.arch_config.decoder_causal
decoder_spatial_padding_mode = config.arch_config.decoder_spatial_padding_mode
self.encoder = LTX2VideoEncoder3d(
in_channels,
latent_channels,
block_out_channels,
down_block_types,
spatio_temporal_scaling,
layers_per_block,
downsample_type,
patch_size,
patch_size_t,
resnet_norm_eps,
encoder_causal,
encoder_spatial_padding_mode,
)
if use_ltx23_condition_encoder:
from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import (
LTX23VideoConditionEncoder,
)
video_encoder_config = dict(
config.arch_config.video_encoder_config
or config.arch_config.video_decoder_config
)
if not video_encoder_config:
raise ValueError(
"LTX-2.3 condition video encoder requires video_encoder_config "
"or video_decoder_config."
)
self.encoder = LTX23VideoConditionEncoder(video_encoder_config)
else:
self.encoder = LTX2VideoEncoder3d(
in_channels,
latent_channels,
block_out_channels,
down_block_types,
spatio_temporal_scaling,
layers_per_block,
downsample_type,
patch_size,
patch_size_t,
resnet_norm_eps,
encoder_causal,
encoder_spatial_padding_mode,
)
if use_ltx23_video_decoder:
video_decoder_config = dict(config.arch_config.video_decoder_config)
@@ -1806,6 +1825,9 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE):
)
def _encode(self, x: torch.Tensor, causal: Optional[bool] = None) -> torch.Tensor:
if self._use_ltx23_condition_encoder:
return self.encoder(x)
batch_size, num_channels, num_frames, height, width = x.shape
if self.use_framewise_decoding and num_frames > self.tile_sample_min_num_frames:
@@ -1835,6 +1857,18 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE):
The latent representations of the encoded videos. If `return_dict` is True, a
[`~models.autoencoder_kl.AutoencoderKLOutput`] is returned, otherwise a plain `tuple` is returned.
"""
if self._use_ltx23_condition_encoder:
if self.use_slicing and x.shape[0] > 1:
encoded_slices = [
self._encode(x_slice, causal=causal) for x_slice in x.split(1)
]
h = torch.cat(encoded_slices)
else:
h = self._encode(x, causal=causal)
if not return_dict:
return (h,)
return DecoderOutput(sample=h)
if self.use_slicing and x.shape[0] > 1:
encoded_slices = [
self._encode(x_slice, causal=causal) for x_slice in x.split(1)
python/sglang/multimodal_gen/runtime/pipelines/joy_echo_pipeline.py
+98
View File
@@ -0,0 +1,98 @@
# SPDX-License-Identifier: Apache-2.0
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
JoyEchoPipelineConfig,
)
from sglang.multimodal_gen.runtime.pipelines.ltx_2_pipeline import (
_add_ltx2_front_stages,
_BaseLTX2Pipeline,
prepare_ltx2_mu,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.image_encoding import (
LTX2ImageEncodingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.joy_echo import (
JoyEchoAVDecodingStage,
JoyEchoDMDDenoisingStage,
JoyEchoMemoryBankFetchStage,
JoyEchoMultishotSetupStage,
JoyEchoSigmaPreparationStage,
PairedAudioVideoMemoryBank,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2 import (
LTX2AVLatentPreparationStage,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
class JoyEchoPipeline(_BaseLTX2Pipeline):
pipeline_name = "JoyEchoPipeline"
is_video_pipeline = True
def __init__(self, *args, **kwargs):
self._memory_bank: PairedAudioVideoMemoryBank | None = None
self.multishot_index: int = 0
self._multishot_session_id: str | None = None
super().__init__(*args, **kwargs)
def _get_or_create_memory_bank(
self, config: JoyEchoPipelineConfig
) -> PairedAudioVideoMemoryBank:
if self._memory_bank is None:
self._memory_bank = PairedAudioVideoMemoryBank(
max_size=int(config.memory_max_size),
num_fix_frames=int(config.memory_num_fix_frames),
)
return self._memory_bank
def reset_memory_bank(self) -> None:
if self._memory_bank is not None:
self._memory_bank.memory.clear()
def create_pipeline_stages(self, server_args: ServerArgs):
config = server_args.pipeline_config
if not isinstance(config, JoyEchoPipelineConfig):
raise TypeError(
f"JoyEchoPipeline requires JoyEchoPipelineConfig, got {type(config)}"
)
memory_bank = self._get_or_create_memory_bank(config)
self.add_stage(JoyEchoMultishotSetupStage(pipeline=self))
_add_ltx2_front_stages(self)
self.add_stage(JoyEchoSigmaPreparationStage())
self.add_standard_timestep_preparation_stage(
prepare_extra_kwargs=[prepare_ltx2_mu]
)
self.add_stages(
[
LTX2AVLatentPreparationStage(
scheduler=self.get_module("scheduler"),
transformer=self.get_module("transformer"),
audio_vae=self.get_module("audio_vae"),
),
LTX2ImageEncodingStage(
vae=self.get_module("vae"),
),
JoyEchoMemoryBankFetchStage(
memory_bank=memory_bank,
vae=self.get_module("vae"),
),
JoyEchoDMDDenoisingStage(
transformer=self.get_module("transformer"),
scheduler=self.get_module("scheduler"),
vae=self.get_module("vae"),
audio_vae=self.get_module("audio_vae"),
sampler_name="euler",
pipeline=self,
),
JoyEchoAVDecodingStage(
vae=self.get_module("vae"),
audio_vae=self.get_module("audio_vae"),
vocoder=self.get_module("vocoder"),
memory_bank=memory_bank,
pipeline=self,
),
]
)
EntryClass = JoyEchoPipeline
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/__init__.py
+24
View File
@@ -0,0 +1,24 @@
# SPDX-License-Identifier: Apache-2.0
"""JoyEcho-specific pipeline stages."""
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.joy_echo.denoising import (
JoyEchoDMDDenoisingStage,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.joy_echo.memory import (
JoyEchoAVDecodingStage,
JoyEchoMemoryBankFetchStage,
PairedAudioVideoMemoryBank,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.joy_echo.setup import (
JoyEchoMultishotSetupStage,
JoyEchoSigmaPreparationStage,
)
__all__ = [
"JoyEchoAVDecodingStage",
"JoyEchoDMDDenoisingStage",
"JoyEchoMemoryBankFetchStage",
"JoyEchoMultishotSetupStage",
"JoyEchoSigmaPreparationStage",
"PairedAudioVideoMemoryBank",
]
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/denoising.py
+575
View File
@@ -0,0 +1,575 @@
# SPDX-License-Identifier: Apache-2.0
import dataclasses
import torch
from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import (
JoyEchoPipelineConfig,
)
from sglang.multimodal_gen.runtime.distributed import get_sp_world_size
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.joy_echo.memory import (
build_memory_audio_rope_coords,
build_memory_self_attention_block_mask,
build_memory_video_rope_coords,
build_paired_memory_cross_mask,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising import (
DenoisingStepState,
LTX2DenoisingContext,
LTX2ModelInputs,
)
from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.denoising_av import (
LTX2AVDenoisingStage,
)
from sglang.multimodal_gen.runtime.server_args import ServerArgs
class JoyEchoDMDDenoisingStage(LTX2AVDenoisingStage):
"""JoyEcho DMD denoising with optional memory prefix and late-layer masks."""
def _prepare_denoising_loop(
self,
batch: Req,
server_args: ServerArgs,
) -> LTX2DenoisingContext:
ctx = super()._prepare_denoising_loop(batch, server_args)
if get_sp_world_size() <= 1:
return ctx
# JoyEcho DMD: shard video on time, replicate full audio on every rank.
# Dual-sharding audio+video adds heavy a2v/v2a all_gather per layer; shot 1
# already uses this layout when memory is injected.
ctx.replicate_audio_for_sp = True
batch.ltx23_audio_replicated_for_sp = True
batch.did_sp_shard_audio_latents = False
return ctx
@staticmethod
def _zero_sp_shard_padding(
latents: torch.Tensor,
*,
valid_token_count: int | None,
) -> torch.Tensor:
if valid_token_count is None or int(valid_token_count) >= int(latents.shape[1]):
return latents
latents = latents.clone()
latents[:, int(valid_token_count) :, :] = 0.0
return latents
@staticmethod
def _expand_sp_token_timestep(
timestep: torch.Tensor,
*,
batch_size: int,
seq_len: int,
valid_token_count: int | None,
) -> torch.Tensor:
"""Expand legacy [B] timesteps to [B, S] and zero SP padding tokens."""
if timestep.ndim >= 2 and int(timestep.shape[1]) == int(seq_len):
ts = timestep
elif timestep.ndim == 1:
ts = timestep.view(batch_size, 1).expand(batch_size, int(seq_len))
else:
ts = timestep
if valid_token_count is not None and int(valid_token_count) < int(seq_len):
ts = ts.clone()
ts[:, int(valid_token_count) :] = 0.0
return ts
def _prepare_ltx2_model_inputs(
self,
ctx: LTX2DenoisingContext,
step: DenoisingStepState,
batch: Req,
server_args: ServerArgs,
sigma: torch.Tensor,
) -> LTX2ModelInputs:
model_inputs = super()._prepare_ltx2_model_inputs(
ctx, step, batch, server_args, sigma
)
if not batch.did_sp_shard_latents:
return model_inputs
batch_size = int(model_inputs.latent_model_input.shape[0])
seq_v = int(model_inputs.latent_model_input.shape[1])
video_valid = batch.sp_video_valid_token_count
video_self_attention_mask = self._build_ltx2_sp_padding_mask(
batch,
seq_len=seq_v,
batch_size=batch_size,
key="sp_video_valid_token_count",
device=model_inputs.latent_model_input.device,
)
video_coords = server_args.pipeline_config.prepare_video_rope_coords_for_sp(
step.current_model,
batch,
model_inputs.latent_model_input,
num_frames=ctx.latent_num_frames_for_model,
height=ctx.latent_height,
width=ctx.latent_width,
)
timestep_video = self._expand_sp_token_timestep(
model_inputs.timestep_video,
batch_size=batch_size,
seq_len=seq_v,
valid_token_count=int(video_valid) if video_valid is not None else None,
)
audio_self_attention_mask = model_inputs.audio_self_attention_mask
audio_coords = model_inputs.audio_coords
timestep_audio = model_inputs.timestep_audio
a2v_cross_attention_mask = model_inputs.a2v_cross_attention_mask
v2a_cross_attention_mask = video_self_attention_mask
if batch.did_sp_shard_audio_latents:
seq_a = int(model_inputs.audio_num_frames_latent)
audio_valid = batch.sp_audio_valid_token_count
audio_self_attention_mask = self._build_ltx2_sp_padding_mask(
batch,
seq_len=seq_a,
batch_size=batch_size,
key="sp_audio_valid_token_count",
device=model_inputs.audio_latent_model_input.device,
)
audio_coords = server_args.pipeline_config.prepare_audio_rope_coords_for_sp(
step.current_model,
batch,
model_inputs.audio_latent_model_input,
num_frames=model_inputs.audio_num_frames_latent,
)
timestep_audio = self._expand_sp_token_timestep(
model_inputs.timestep_audio,
batch_size=batch_size,
seq_len=seq_a,
valid_token_count=int(audio_valid) if audio_valid is not None else None,
)
a2v_cross_attention_mask = audio_self_attention_mask
return dataclasses.replace(
model_inputs,
video_coords=video_coords,
audio_coords=audio_coords,
timestep_video=timestep_video,
timestep_audio=timestep_audio,
video_self_attention_mask=video_self_attention_mask,
audio_self_attention_mask=audio_self_attention_mask,
a2v_cross_attention_mask=a2v_cross_attention_mask,
v2a_cross_attention_mask=v2a_cross_attention_mask,
)
def _build_ltx2_base_model_kwargs(
self,
ctx: LTX2DenoisingContext,
batch: Req,
model_inputs: LTX2ModelInputs,
) -> dict[str, object]:
kwargs = super()._build_ltx2_base_model_kwargs(ctx, batch, model_inputs)
if not batch.did_sp_shard_latents:
return kwargs
kwargs.update(
{
"video_self_attention_mask": model_inputs.video_self_attention_mask,
"audio_self_attention_mask": model_inputs.audio_self_attention_mask,
"a2v_cross_attention_mask": model_inputs.a2v_cross_attention_mask,
"v2a_cross_attention_mask": model_inputs.v2a_cross_attention_mask,
"audio_replicated_for_sp": bool(ctx.replicate_audio_for_sp),
"legacy_ltx23_one_stage_semantics": False,
}
)
return kwargs
@staticmethod
def _dmd_add_noise(
original: torch.Tensor,
noise: torch.Tensor,
sigma: torch.Tensor,
) -> torch.Tensor:
sigma_t = sigma.to(device=original.device, dtype=original.dtype)
if sigma_t.ndim == 1:
sigma_t = sigma_t.reshape(-1, *[1] * (original.ndim - 1))
elif sigma_t.ndim == 2:
sigma_t = sigma_t.reshape(*sigma_t.shape, *[1] * (original.ndim - 2))
return (1.0 - sigma_t) * original + sigma_t * noise
def _sample_sp_consistent_noise(
self,
local_reference: torch.Tensor,
batch: Req,
server_args: ServerArgs,
*,
shard_video: bool,
shard_audio: bool,
) -> torch.Tensor:
"""Sample renoise on the global latent layout, then shard for SP."""
if shard_video:
raw_shape = batch.raw_latent_shape
if not (isinstance(raw_shape, tuple) and len(raw_shape) == 3):
raise ValueError(
"SP DMD renoise requires packed video `batch.raw_latent_shape`."
)
full_reference = torch.empty(
tuple(raw_shape),
device=local_reference.device,
dtype=local_reference.dtype,
)
full_noise = self._randn_like_with_batch_generators(full_reference, batch)
sharded_noise, _ = server_args.pipeline_config.shard_latents_for_sp(
batch, full_noise
)
return sharded_noise
if shard_audio:
orig_audio_len = batch.sp_audio_orig_num_frames
if orig_audio_len <= 0:
raise ValueError(
"SP DMD renoise requires `batch.sp_audio_orig_num_frames`."
)
full_reference = torch.empty(
(
int(local_reference.shape[0]),
int(orig_audio_len),
int(local_reference.shape[2]),
),
device=local_reference.device,
dtype=local_reference.dtype,
)
full_noise = self._randn_like_with_batch_generators(full_reference, batch)
sharded_noise, _ = server_args.pipeline_config.shard_audio_latents_for_sp(
batch, full_noise
)
return sharded_noise
return self._randn_like_with_batch_generators(local_reference, batch)
@staticmethod
def _apply_memory_prefix_to_timestep(
timestep: torch.Tensor,
*,
memory_seq_len: int,
target_seq_len: int,
) -> torch.Tensor:
if memory_seq_len <= 0:
return timestep
batch_size = int(timestep.shape[0])
device = timestep.device
dtype = timestep.dtype
if timestep.ndim == 3:
memory_ts = torch.zeros(
batch_size,
memory_seq_len,
timestep.shape[-1],
device=device,
dtype=dtype,
)
target_ts = timestep[:, :target_seq_len, :]
return torch.cat([memory_ts, target_ts], dim=1)
if timestep.ndim == 2:
memory_ts = torch.zeros(
batch_size, memory_seq_len, device=device, dtype=dtype
)
target_ts = timestep[:, :target_seq_len]
return torch.cat([memory_ts, target_ts], dim=1)
if timestep.ndim == 1:
# JoyEcho legacy one-stage uses scalar [B] timesteps; memory tokens
# must see sigma=0 (clean) while target tokens keep the current sigma.
memory_ts = torch.zeros(
batch_size, memory_seq_len, 1, device=device, dtype=dtype
)
target_ts = timestep.view(batch_size, 1, 1).expand(
batch_size, target_seq_len, 1
)
return torch.cat([memory_ts, target_ts], dim=1)
return timestep
def _build_memory_model_inputs(
self,
model_inputs: LTX2ModelInputs,
batch: Req,
ctx: LTX2DenoisingContext,
server_args: ServerArgs,
step_model,
) -> tuple[LTX2ModelInputs, dict[str, int]]:
memory_info = batch.extra.get("joy_echo_memory")
if not memory_info:
return model_inputs, {}
memory_video = memory_info["memory_video_packed"].to(
device=model_inputs.latent_model_input.device,
dtype=model_inputs.latent_model_input.dtype,
)
memory_audio = memory_info["memory_audio"].to(
device=model_inputs.audio_latent_model_input.device,
dtype=model_inputs.audio_latent_model_input.dtype,
)
target_video = model_inputs.latent_model_input
target_audio = model_inputs.audio_latent_model_input
memory_video_len = int(memory_video.shape[1])
memory_audio_len = int(memory_audio.shape[1])
target_video_len = int(target_video.shape[1])
target_audio_len = int(target_audio.shape[1])
tokens_per_latent_frame = int(ctx.latent_height) * int(ctx.latent_width)
if (
tokens_per_latent_frame <= 0
or memory_video_len % tokens_per_latent_frame != 0
):
num_memory_slots = int(memory_info["num_memory_slots"])
else:
num_memory_slots = memory_video_len // tokens_per_latent_frame
latent_model_input = torch.cat([memory_video, target_video], dim=1)
audio_latent_model_input = torch.cat([memory_audio, target_audio], dim=1)
timestep_video = self._apply_memory_prefix_to_timestep(
model_inputs.timestep_video,
memory_seq_len=memory_video_len,
target_seq_len=target_video_len,
)
timestep_audio = self._apply_memory_prefix_to_timestep(
model_inputs.timestep_audio,
memory_seq_len=memory_audio_len,
target_seq_len=target_audio_len,
)
device = latent_model_input.device
batch_size = int(latent_model_input.shape[0])
sp_world_size = get_sp_world_size()
sp_on = sp_world_size > 1 and batch.did_sp_shard_latents
if sp_on:
target_video_full_len = int(target_video_len) * int(sp_world_size)
raw_shape = batch.raw_latent_shape
if isinstance(raw_shape, tuple) and len(raw_shape) == 3:
target_video_valid_len = int(raw_shape[1])
else:
target_video_valid_len = target_video_full_len
sp_target_start_offset = batch.sp_video_start_frame
else:
target_video_full_len = target_video_len
target_video_valid_len = target_video_len
sp_target_start_offset = 0
a2v_mask = build_paired_memory_cross_mask(
batch_size=batch_size,
query_memory_seq_len=memory_video_len,
query_target_seq_len=target_video_len,
kv_memory_seq_len=memory_audio_len,
kv_target_seq_len=target_audio_len,
num_memory_slots=num_memory_slots,
device=device,
kv_segment_lengths=memory_info.get("memory_audio_segment_lengths"),
)
v2a_mask = build_paired_memory_cross_mask(
batch_size=batch_size,
query_memory_seq_len=memory_audio_len,
query_target_seq_len=target_audio_len,
kv_memory_seq_len=memory_video_len,
kv_target_seq_len=target_video_full_len,
num_memory_slots=num_memory_slots,
device=device,
query_segment_lengths=memory_info.get("memory_audio_segment_lengths"),
)
video_self_attention_mask = None
if sp_on and target_video_full_len > target_video_valid_len:
v2a_mask[:, :, memory_video_len + target_video_valid_len :] = False
if sp_on:
vself_len = memory_video_len + target_video_full_len
video_self_attention_mask = torch.ones(
(batch_size, vself_len), device=device, dtype=torch.bool
)
if target_video_full_len > target_video_valid_len:
video_self_attention_mask[
:, memory_video_len + target_video_valid_len :
] = False
audio_self_attention_mask = build_memory_self_attention_block_mask(
batch_size=batch_size,
memory_seq_len=memory_audio_len,
target_seq_len=target_audio_len,
device=device,
)
config = server_args.pipeline_config
late_layer_ratio = 1.0
memory_position_mode = "reference"
if isinstance(config, JoyEchoPipelineConfig):
late_layer_ratio = float(config.late_layer_ratio)
memory_position_mode = config.memory_position_mode
video_coords = build_memory_video_rope_coords(
rope=step_model.rope,
batch_size=batch_size,
memory_video_len=memory_video_len,
target_num_frames=int(ctx.latent_num_frames_for_model),
latent_height=int(ctx.latent_height),
latent_width=int(ctx.latent_width),
device=device,
fps=float(batch.fps),
memory_position_mode=str(
memory_info.get("memory_position_mode", memory_position_mode)
),
memory_downscale_factor=int(memory_info.get("memory_downscale_factor", 1)),
sp_target_start_offset=sp_target_start_offset,
)
audio_coords = build_memory_audio_rope_coords(
audio_rope=step_model.audio_rope,
batch_size=batch_size,
memory_audio_len=memory_audio_len,
target_audio_len=target_audio_len,
device=device,
memory_position_mode=str(
memory_info.get("memory_position_mode", memory_position_mode)
),
)
return (
LTX2ModelInputs(
latent_model_input=latent_model_input,
audio_latent_model_input=audio_latent_model_input,
audio_num_frames_latent=memory_audio_len + target_audio_len,
video_coords=video_coords,
audio_coords=audio_coords,
timestep_video=timestep_video,
timestep_audio=timestep_audio,
prompt_timestep_video=model_inputs.prompt_timestep_video,
prompt_timestep_audio=model_inputs.prompt_timestep_audio,
video_self_attention_mask=video_self_attention_mask,
audio_self_attention_mask=audio_self_attention_mask,
a2v_cross_attention_mask=a2v_mask,
v2a_cross_attention_mask=v2a_mask,
),
{
"memory_video_len": memory_video_len,
"memory_audio_len": memory_audio_len,
"late_layer_ratio": late_layer_ratio,
"audio_replicated_for_sp": sp_on,
"video_memory_prefix_len": memory_video_len if sp_on else 0,
},
)
def _run_denoising_step(
self,
ctx: LTX2DenoisingContext,
step: DenoisingStepState,
batch: Req,
server_args: ServerArgs,
) -> None:
if ctx.audio_latents is None:
raise ValueError("JoyEcho requires audio latents for denoising.")
if ctx.audio_scheduler is None:
raise ValueError("JoyEcho audio scheduler was not prepared.")
sigmas = ctx.scheduler.sigmas
if not isinstance(sigmas, torch.Tensor):
raise ValueError("Expected scheduler.sigmas to be a tensor for JoyEcho.")
sigma = sigmas[step.step_index].to(
device=ctx.latents.device, dtype=torch.float32
)
sigma_next = sigmas[step.step_index + 1].to(
device=ctx.latents.device, dtype=torch.float32
)
sigma_val = float(sigma.item())
sigma_next_val = float(sigma_next.item())
model_inputs = self._prepare_ltx2_model_inputs(
ctx, step, batch, server_args, sigma
)
model_inputs, memory_meta = self._build_memory_model_inputs(
model_inputs, batch, ctx, server_args, step.current_model
)
prompt_attention_mask = self._get_ltx_prompt_attention_mask(
batch,
is_ltx23_variant=ctx.is_ltx23_variant,
)
base_model_kwargs = self._build_ltx2_base_model_kwargs(ctx, batch, model_inputs)
model_kwargs = self._build_ltx2_model_kwargs(
ctx,
base_model_kwargs,
encoder_hidden_states=batch.prompt_embeds[0],
audio_encoder_hidden_states=batch.audio_prompt_embeds[0],
encoder_attention_mask=prompt_attention_mask,
)
if memory_meta:
# Legacy one-stage LTX2 skips mask kwargs in the base builder; memory
# mode must always pass paired cross/self masks to the DiT.
model_kwargs["late_layer_ratio"] = memory_meta["late_layer_ratio"]
model_kwargs["late_audio_self_attention_mask"] = None
model_kwargs["video_self_attention_mask"] = (
model_inputs.video_self_attention_mask
)
model_kwargs["audio_self_attention_mask"] = (
model_inputs.audio_self_attention_mask
)
model_kwargs["a2v_cross_attention_mask"] = (
model_inputs.a2v_cross_attention_mask
)
model_kwargs["v2a_cross_attention_mask"] = (
model_inputs.v2a_cross_attention_mask
)
model_kwargs["audio_replicated_for_sp"] = memory_meta[
"audio_replicated_for_sp"
]
model_kwargs["video_memory_prefix_len"] = memory_meta[
"video_memory_prefix_len"
]
with self._ltx2_model_forward_context(ctx, step):
model_video, model_audio = step.current_model(**model_kwargs)
if memory_meta:
memory_video_len = memory_meta["memory_video_len"]
memory_audio_len = memory_meta["memory_audio_len"]
model_video = model_video[:, memory_video_len:, :]
if model_audio is not None:
model_audio = model_audio[:, memory_audio_len:, :]
denoised_video = self._ltx2_velocity_to_x0(
ctx.latents, model_video.float(), sigma_val
)
denoised_audio = self._ltx2_velocity_to_x0(
ctx.audio_latents, model_audio.float(), sigma_val
)
denoised_video = self._ltx2_apply_clean_latent_mask(denoised_video, ctx)
if sigma_next_val > 0.0:
video_noise = self._sample_sp_consistent_noise(
ctx.latents,
batch,
server_args,
shard_video=batch.did_sp_shard_latents,
shard_audio=False,
).float()
audio_noise = self._sample_sp_consistent_noise(
ctx.audio_latents,
batch,
server_args,
shard_video=False,
shard_audio=batch.did_sp_shard_audio_latents,
).float()
next_video_latents = self._dmd_add_noise(
denoised_video, video_noise, sigma_next
).to(dtype=ctx.latents.dtype)
next_audio_latents = self._dmd_add_noise(
denoised_audio, audio_noise, sigma_next
).to(dtype=ctx.audio_latents.dtype)
else:
next_video_latents = denoised_video.to(dtype=ctx.latents.dtype)
next_audio_latents = denoised_audio.to(dtype=ctx.audio_latents.dtype)
if batch.did_sp_shard_latents:
next_video_latents = self._zero_sp_shard_padding(
next_video_latents,
valid_token_count=batch.sp_video_valid_token_count,
)
if batch.did_sp_shard_audio_latents:
next_audio_latents = self._zero_sp_shard_padding(
next_audio_latents,
valid_token_count=batch.sp_audio_valid_token_count,
)
ctx.latents = next_video_latents
ctx.audio_latents = next_audio_latents
ctx.latents = self.post_forward_for_ti2v_task(
batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z
)
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/memory.py
+1025
View File
File diff suppressed because it is too large Load Diff
python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/setup.py
+92
View File
@@ -0,0 +1,92 @@
# SPDX-License-Identifier: Apache-2.0
"""JoyEcho pre-denoising setup stages (multi-shot session + sigma schedule)."""
from __future__ import annotations
from typing import TYPE_CHECKING
from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req
from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage
from sglang.multimodal_gen.runtime.server_args import ServerArgs
from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger
if TYPE_CHECKING:
from sglang.multimodal_gen.runtime.pipelines.joy_echo_pipeline import (
JoyEchoPipeline,
)
logger = init_logger(__name__)
class JoyEchoMultishotSetupStage(PipelineStage):
"""Apply official per-shot seeding before input validation and latent noise."""
def __init__(self, pipeline: JoyEchoPipeline) -> None:
super().__init__()
self.pipeline = pipeline
def _maybe_reset_multishot_session(self, batch: Req) -> None:
"""Reset shot index and memory bank at the start of a new ``generate()`` session."""
if not batch.reset_memory_bank:
return
session_id = batch.request_id
if session_id is not None:
if session_id == self.pipeline._multishot_session_id:
return
self.pipeline._multishot_session_id = session_id
self.pipeline.multishot_index = 0
self.pipeline.reset_memory_bank()
logger.info(
"JoyEcho memory bank reset for new multi-shot session (request_id=%s)",
session_id,
)
return
if self.pipeline.multishot_index == 0:
self.pipeline.reset_memory_bank()
logger.info("JoyEcho memory bank reset for new multi-shot session")
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
if not batch.enable_memory_bank:
return batch
self._maybe_reset_multishot_session(batch)
shot_idx = self.pipeline.multishot_index
self.pipeline.multishot_index += 1
base_seed = batch.seed
if isinstance(base_seed, list):
if not base_seed:
raise ValueError("seed list must not be empty for JoyEcho multi-shot")
base_seed = base_seed[0]
# Official inference.py: prompt_seed = int(cfg.seed) + shot_idx
batch.seed = int(base_seed) + shot_idx
logger.info(
"JoyEcho multi-shot setup: shot_idx=%d seed=%d",
shot_idx,
batch.seed,
)
return batch
class JoyEchoSigmaPreparationStage(PipelineStage):
"""Prepare JoyEcho DMD sigma schedule without LTX-2 shift remapping."""
def forward(self, batch: Req, server_args: ServerArgs) -> Req:
batch.extra["ltx2_phase"] = "stage1"
sigmas = batch.sigmas
if sigmas is None:
sampling_sigmas = batch.sampling_params.sigmas
if sampling_sigmas is not None:
sigmas = list(sampling_sigmas)
else:
sigmas = list(server_args.pipeline_config.default_sigmas)
batch.sigmas = list(sigmas)
batch.num_inference_steps = max(len(batch.sigmas) - 1, 1)
return batch
python/sglang/multimodal_gen/runtime/utils/model_overlay.py
+4
View File
@@ -29,6 +29,10 @@ BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {
"overlay_repo_id": "MickJ/LTX-2.3-overlay",
"overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e",
},
"jdopensource/JoyAI-Echo": {
"overlay_repo_id": "Niehen6174/JoyAI-Echo-overlay",
"overlay_revision": "0a19f315c96532b7a5f61bcd765d1fefdd83dc7d",
},
"Efficient-Large-Model/SANA-WM_bidirectional": {
"overlay_repo_id": "sjmshsh/SANA-WM_bidirectional-overlay",
"overlay_revision": "e611beacbcc0cf33c676306ae0eb89f149e044ad",
python/sglang/multimodal_gen/test/server/consistency_threshold.json
+6
View File
@@ -276,6 +276,12 @@
"ssim_threshold": 0.29,
"psnr_threshold": 11.7,
"mean_abs_diff_threshold": 47.0
},
"joy_echo_t2v_2gpu": {
"clip_threshold": 0.78,
"ssim_threshold": 0.48,
"psnr_threshold": 13.0,
"mean_abs_diff_threshold": 45.0
}
},
"default_clip_threshold_image": 0.92,
python/sglang/multimodal_gen/test/server/gpu_cases.py
+15
View File
@@ -20,6 +20,7 @@ from sglang.multimodal_gen.test.server.testcase_configs import (
DiffusionServerArgs,
DiffusionTestCase,
IDEOGRAM4_CI_sampling_params,
JOY_ECHO_T2V_CI_sampling_params,
LINGBOT_WORLD_REALTIME_sampling_params,
MODELOPT_QWEN_IMAGE_2512_NVFP4_CI_sampling_params,
MODELOPT_T2I_CI_sampling_params,
@@ -658,6 +659,20 @@ TWO_GPU_CASES = [
output_size="832x480",
),
),
DiffusionTestCase(
"joy_echo_t2v_2gpu",
DiffusionServerArgs(
model_path="jdopensource/JoyAI-Echo",
extras=["--ulysses-degree=2"],
env_vars={
"PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
},
),
JOY_ECHO_T2V_CI_sampling_params,
run_perf_check=False,
run_consistency_check=True,
run_component_accuracy_check=False,
),
DiffusionTestCase(
"wan2_1_t2v_1.3b_cfg_parallel",
DiffusionServerArgs(
python/sglang/multimodal_gen/test/server/testcase_configs.py
+11
View File
@@ -536,6 +536,17 @@ T2V_sampling_params = DiffusionSamplingParams(
prompt=T2V_PROMPT,
)
JOY_ECHO_T2V_CI_sampling_params = DiffusionSamplingParams(
prompt=T2V_PROMPT,
output_size="640x384",
num_frames=33,
extras={
"num_inference_steps": 8,
"seed": 42,
"enable_memory_bank": False,
},
)
MODELOPT_T2V_CI_sampling_params = DiffusionSamplingParams(
prompt=T2V_PROMPT,
output_size="640x384",
python/sglang/multimodal_gen/test/test_utils.py
+1 -1
View File
@@ -34,7 +34,7 @@ if TYPE_CHECKING:
logger = init_logger(__name__)
SGL_TEST_FILES_CI_DATA_REVISION = "3c6e06ae99001d93f7901bc9b7fdf19ec6c2ce4e"
SGL_TEST_FILES_CI_DATA_REVISION = "4a271ef34602043f19d253f0d30a5f653fe11325"
if current_platform.is_npu():
SGL_TEST_FILES_CI_DATA_REVISION = "670d66a8a290b62c0c3c077b3e9b0f4a4d9a44e7"