diff --git a/docs_new/cookbook/diffusion/JoyEcho/JoyEcho.mdx b/docs_new/cookbook/diffusion/JoyEcho/JoyEcho.mdx new file mode 100644 index 000000000..a28613ef5 --- /dev/null +++ b/docs_new/cookbook/diffusion/JoyEcho/JoyEcho.mdx @@ -0,0 +1,173 @@ +--- +title: JoyEcho +description: Run JoyAI-Echo multi-shot audio–video generation with SGLang Diffusion. +metatags: + description: "Deploy and use JoyAI-Echo long-form audio–video generation with SGLang Diffusion, including single-shot and multi-shot memory-bank workflows." +--- + +## 1. Model Introduction + +[JoyAI-Echo](https://huggingface.co/jdopensource/JoyAI-Echo) (JoyEcho) is a long-form audio–video generation model built on the LTX-2 backbone. Its core idea is a **paired audio–video memory bank**: each shot commits decoded frames and audio latents into a rolling bank, and subsequent shots condition on that memory prefix. This enables **multi-shot, minute-scale generation** with visual and audio continuity across prompts. + +Use `jdopensource/JoyAI-Echo` as `--model-path`. SGLang loads the monolithic release through the built-in [JoyAI-Echo-overlay](https://huggingface.co/Niehen6174/JoyAI-Echo-overlay) materialization path, similar to LTX-2.3-overlay. + +| Aspect | Standard LTX-2.3 | JoyEcho | +| --- | --- | --- | +| Pipeline | `LTX2Pipeline` / `LTX2TwoStageHQPipeline` | `JoyEchoPipeline` (default for this model) | +| Denoising | Multi-step flow matching + CFG | LTX-2 DMD distilled path (8 steps, `guidance_scale=1.0`) | +| Multi-shot | Not supported | Paired audio–video memory bank across shots | +| Sequence parallelism | LTX-2 SP (video/audio sharded) | Ulysses SP (`ulysses_degree=2`): single-shot and multi-shot + memory bank | +| Post-processing | Optional two-stage HQ upscaling | Per-shot mp4 output | + + +Review the model license on the [JoyAI-Echo Hugging Face page](https://huggingface.co/jdopensource/JoyAI-Echo) before production or commercial use. SGLang support does not grant additional model usage rights. + + +## 2. SGLang-diffusion Installation + +Install SGLang with diffusion dependencies: + +```bash +uv pip install "sglang[diffusion]" --prerelease=allow +``` + +For platform-specific setup, see the [SGLang Diffusion installation guide](/docs/sglang-diffusion/installation). + +## 3. Model Deployment + +JoyEcho uses the default `JoyEchoPipeline` registered for `jdopensource/JoyAI-Echo`. A single high-VRAM GPU (for example H100 or H200) is enough for the common 832x480 / 121-frame / 8-step setting. + +```bash +sglang serve \ + --model-path jdopensource/JoyAI-Echo +``` + +Optional environment variable for long runs: + +```bash +export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True +``` + +For multi-GPU serving, tensor parallelism (TP) and **Ulysses sequence parallelism (SP)** are supported. JoyEcho SP uses an **asymmetric layout**: video target latents are time-sharded across ranks, while audio (including memory tokens) is **replicated** on every rank so cross-attention stays temporally aligned. Multi-shot runs with `enable_memory_bank=true` are supported on SP. + +```bash +sglang serve \ + --model-path jdopensource/JoyAI-Echo \ + --num-gpus 2 \ + --ulysses-degree 2 +``` + + +JoyEcho SP currently targets **Ulysses-only** parallelism (`ulysses_degree=2`, `ring_degree=1`). Ring SP is not validated for this pipeline. For `sglang generate`, add `--num-gpus 2 --ulysses-degree 2` to the commands in section 4. + + +## 4. Model Invocation + +### 4.1 Default sampling + +| Setting | Default | +| --- | --- | +| Resolution | 832x480 | +| Frames | 121 | +| FPS | 25 | +| Steps | 8 | +| Guidance scale | 1.0 | +| Seed | 12345 | + +### 4.2 Single-shot text-to-video + +```bash +sglang generate \ + --model-path jdopensource/JoyAI-Echo \ + --prompt "A curious raccoon walks through a sunlit forest path" \ + --height 480 --width 832 --num-frames 121 --fps 25 \ + --num-inference-steps 8 --seed 42 \ + --save-output +``` + +Disable the memory bank for standalone clips with a config file: + +```bash +cat > /tmp/joy_echo_single.json <<'EOF' +{ + "model_path": "jdopensource/JoyAI-Echo", + "prompt": "A curious raccoon walks through a sunlit forest path", + "enable_memory_bank": false, + "seed": 42, + "height": 480, + "width": 832, + "num_frames": 121, + "fps": 25, + "num_inference_steps": 8 +} +EOF + +sglang generate --config /tmp/joy_echo_single.json --save-output +``` + +### 4.3 Multi-shot generation + +JoyEcho does **not** generate all shots in one forward pass. Each shot is one generation request. Continuity is carried by an in-process **memory bank** on the pipeline instance. + +Typical workflow: + +1. **Shot 0** — memory bank is empty; the model generates a standalone A/V clip. +2. **After decode** — decoded video frames and packed audio latents are committed to the memory bank (up to 7 slots by default). +3. **Shot 1+** — prior-shot frames are re-encoded and prepended as a memory prefix before denoising. +4. **Per-shot seeding** — official semantics use `prompt_seed = base_seed + shot_index`. + +Pass multiple prompts as a list in a config file: + +```bash +cat > /tmp/joy_echo_4shot.json <<'EOF' +{ + "model_path": "jdopensource/JoyAI-Echo", + "prompt": [ + "Shot 0: A raccoon wakes up in a cozy attic.", + "Shot 1: The raccoon climbs down and opens the back door.", + "Shot 2: It walks through a rainy alley under neon signs.", + "Shot 3: The raccoon finds a warm bakery window and stops." + ], + "enable_memory_bank": true, + "reset_memory_bank": true, + "seed": 42, + "height": 480, + "width": 832, + "num_frames": 121, + "fps": 25, + "num_inference_steps": 8 +} +EOF + +sglang generate --config /tmp/joy_echo_4shot.json --save-output +``` + +You can also pass prompts from a text file (one prompt per line) with `--prompt-path`: + +```bash +sglang generate \ + --model-path jdopensource/JoyAI-Echo \ + --prompt-path /tmp/joy_echo_shots.txt \ + --seed 42 \ + --height 480 --width 832 --num-frames 121 --fps 25 \ + --num-inference-steps 8 \ + --save-output +``` + +### 4.4 Memory bank controls + +| Parameter | Default | Meaning | +| --- | --- | --- | +| `enable_memory_bank` | `true` | Read/write the paired A/V memory bank between shots. | +| `reset_memory_bank` | `true` | Clear the bank and shot counter at the start of a new session (`request_id` change or first shot). | + +Set `enable_memory_bank=false` when you want independent shots without cross-shot continuity. + +## 5. Practical Tips + +- Use `--num-inference-steps 8` and `--guidance-scale 1.0` to match the official JoyEcho DMD distilled path. +- Multi-shot prompts can be passed as a `prompt` list, via `prompt_path`, or as sequential API calls on the same server instance. +- The memory bank caps at **7 slots**; from shot 8 onward the oldest slots roll off. +- For **2-GPU latency**, try **Ulysses SP** (`--num-gpus 2 --ulysses-degree 2`) on both single-shot and multi-shot runs. Use **TP** when you need a different sharding strategy or more than two GPUs. +- Set `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True` for long multi-shot SP sessions. +- JoyEcho outputs per-shot mp4 files with synchronized audio. There is no built-in two-stage HQ upscaling path like LTX-2.3 HQ. diff --git a/docs_new/docs.json b/docs_new/docs.json index 063f9ab24..856dc534b 100644 --- a/docs_new/docs.json +++ b/docs_new/docs.json @@ -1181,6 +1181,13 @@ "cookbook/diffusion/LTX/LTX2 & LTX2.3" ] }, + { + "group": "JoyEcho", + "tag": "NEW", + "pages": [ + "cookbook/diffusion/JoyEcho/JoyEcho" + ] + }, { "group": "Qwen-Image", "pages": [ diff --git a/python/sglang/multimodal_gen/configs/models/dits/joy_echo.py b/python/sglang/multimodal_gen/configs/models/dits/joy_echo.py new file mode 100644 index 000000000..0847c411c --- /dev/null +++ b/python/sglang/multimodal_gen/configs/models/dits/joy_echo.py @@ -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" diff --git a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py index 240214639..70a59226b 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py @@ -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) diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/joy_echo.py b/python/sglang/multimodal_gen/configs/pipeline_configs/joy_echo.py new file mode 100644 index 000000000..913172df7 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/joy_echo.py @@ -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 diff --git a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py index 20e7d6399..fb03fcd29 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -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", ): diff --git a/python/sglang/multimodal_gen/configs/sample/joy_echo.py b/python/sglang/multimodal_gen/configs/sample/joy_echo.py new file mode 100644 index 000000000..a7a870f55 --- /dev/null +++ b/python/sglang/multimodal_gen/configs/sample/joy_echo.py @@ -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 diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index f22ddb121..e6cb4f86e 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -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( diff --git a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py index 57a737078..4775f8477 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py index a985ad871..e1da51bbf 100644 --- a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py +++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_vae.py @@ -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) diff --git a/python/sglang/multimodal_gen/runtime/pipelines/joy_echo_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/joy_echo_pipeline.py new file mode 100644 index 000000000..75e9a8c85 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines/joy_echo_pipeline.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/__init__.py new file mode 100644 index 000000000..5d529dcc6 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/__init__.py @@ -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", +] diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/denoising.py new file mode 100644 index 000000000..adc42fca3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/denoising.py @@ -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 + ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/memory.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/memory.py new file mode 100644 index 000000000..231931fe3 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/memory.py @@ -0,0 +1,1025 @@ +# SPDX-License-Identifier: Apache-2.0 +"""JoyEcho memory bank utilities and memory-related pipeline stages.""" + +from __future__ import annotations + +import math +import random +from dataclasses import dataclass, field +from typing import Any, Literal, Optional + +import numpy as np +import torch +import torchaudio +from PIL import Image +from torch import Tensor +from torchvision.transforms import functional as TVF + +from sglang.multimodal_gen.configs.pipeline_configs.joy_echo import ( + JoyEchoPipelineConfig, +) +from sglang.multimodal_gen.runtime.distributed import get_local_torch_device +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import OutputBatch, Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.base import PipelineStage +from sglang.multimodal_gen.runtime.pipelines_core.stages.model_specific_stages.ltx_2.decoding_av import ( + LTX2AVDecodingStage, +) +from sglang.multimodal_gen.runtime.server_args import ServerArgs +from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger + +logger = init_logger(__name__) + + +# --- Audio peak-window helpers --- + + +def latent_window_size_to_pixel_window_size( + latent_window_size: int, + *, + downsample_factor: int, + is_causal: bool = True, +) -> int: + if latent_window_size <= 0: + raise ValueError( + f"latent_window_size must be positive, got {latent_window_size}" + ) + if downsample_factor <= 0: + raise ValueError(f"downsample_factor must be positive, got {downsample_factor}") + + pixel_window_size = int(latent_window_size) * int(downsample_factor) + if is_causal: + pixel_window_size = max(pixel_window_size - (int(downsample_factor) - 1), 1) + return pixel_window_size + + +def select_max_response_audio_window_with_bounds( + segment: Tensor, + window_size: int, +) -> tuple[Tensor, Tensor, Tensor]: + if segment.dim() != 4: + raise ValueError( + f"Expected segment shape [B, C, T, F], got {tuple(segment.shape)}" + ) + if window_size <= 0: + raise ValueError(f"window_size must be positive, got {window_size}") + + num_time_steps = segment.shape[2] + if num_time_steps <= 0: + raise ValueError("Cannot select from an empty audio segment") + + scan_stride = max(1, window_size // 4) + offsets = torch.arange(window_size, device=segment.device) + max_start_idx = ( + num_time_steps - window_size + if num_time_steps >= window_size + else num_time_steps - 1 + ) + candidate_start_indices = list(range(0, max_start_idx + 1, scan_stride)) + if candidate_start_indices[-1] != max_start_idx: + candidate_start_indices.append(max_start_idx) + + candidate_windows = [] + candidate_scores = [] + candidate_start_indices_tensor = torch.tensor( + candidate_start_indices, device=segment.device, dtype=torch.long + ) + for start_idx in candidate_start_indices: + gather_indices = (start_idx + offsets).clamp(0, num_time_steps - 1).long() + window = segment.index_select(dim=2, index=gather_indices) + candidate_windows.append(window) + candidate_scores.append(window.float().exp().sum(dim=(1, 2, 3))) + + scores = torch.stack(candidate_scores, dim=1) + best_window_indices = scores.argmax(dim=1) + best_start_indices = candidate_start_indices_tensor[best_window_indices] + best_end_indices = torch.clamp( + best_start_indices + window_size - 1, max=num_time_steps - 1 + ) + selected_windows = torch.cat( + [ + candidate_windows[int(best_window_indices[batch_index])][ + batch_index : batch_index + 1 + ] + for batch_index in range(segment.shape[0]) + ], + dim=0, + ) + return selected_windows, best_start_indices, best_end_indices + + +def select_audio_window_with_bounds( + segment: Tensor, + window_size: int, + *, + mode: Literal["max_response", "random"] = "random", + rng: random.Random | None = None, +) -> tuple[Tensor, Tensor, Tensor]: + if mode == "max_response": + return select_max_response_audio_window_with_bounds(segment, window_size) + if mode == "random": + offsets = torch.arange(window_size, device=segment.device, dtype=torch.long) + num_time_steps = segment.shape[2] + max_start_idx = max(0, num_time_steps - window_size) + start_indices = torch.randint( + 0, max_start_idx + 1, (segment.shape[0],), device=segment.device + ) + selected_windows = torch.cat( + [ + segment[ + batch_index : batch_index + 1, + :, + (start_indices[batch_index] + offsets).clamp(0, num_time_steps - 1), + :, + ] + for batch_index in range(segment.shape[0]) + ], + dim=0, + ) + end_indices = torch.clamp( + start_indices + window_size - 1, max=num_time_steps - 1 + ) + return selected_windows, start_indices, end_indices + raise ValueError(f"Unsupported audio window selection mode: {mode}") + + +def mel_window_bounds_to_seconds( + start_index: int, + end_index: int, + *, + hop_length: int, + sample_rate: int, +) -> tuple[float, float]: + start_time_sec = float(start_index * hop_length) / float(sample_rate) + end_time_sec = float((end_index + 1) * hop_length) / float(sample_rate) + return start_time_sec, end_time_sec + + +def select_video_frame_indices_from_time_range( + *, + num_frames: int, + fps: float, + start_time_sec: float, + end_time_sec: float, + count: int = 1, + mode: Literal["first", "random", "center"] = "center", + rng: random.Random | None = None, +) -> list[int]: + if num_frames <= 0: + raise ValueError(f"num_frames must be positive, got {num_frames}") + + start_frame = int(math.ceil(start_time_sec * fps)) + end_frame = int(math.ceil(end_time_sec * fps)) - 1 + start_frame = max(0, min(start_frame, num_frames - 1)) + end_frame = max(0, min(end_frame, num_frames - 1)) + + if end_frame < start_frame: + center_time_sec = max(0.0, 0.5 * (start_time_sec + end_time_sec)) + center_frame = int(round(center_time_sec * fps)) + candidate_frames = [max(0, min(center_frame, num_frames - 1))] + else: + candidate_frames = list(range(start_frame, end_frame + 1)) + + if mode == "center": + if len(candidate_frames) <= count: + selected = candidate_frames[:] + else: + center_offset = max(0, (len(candidate_frames) - count) // 2) + selected = candidate_frames[center_offset : center_offset + count] + elif mode == "first": + selected = candidate_frames[:count] + else: + rng = rng or random + selected = ( + candidate_frames[:] + if len(candidate_frames) <= count + else sorted(rng.sample(candidate_frames, count)) + ) + + if len(selected) < count: + selected.extend([selected[-1]] * (count - len(selected))) + return selected + + +# --- Slot-aware attention masks --- + + +def memory_slot_ranges(total_seq_len: int, num_slots: int) -> list[tuple[int, int]]: + if total_seq_len <= 0 or num_slots <= 0: + return [] + + ranges: list[tuple[int, int]] = [] + start = 0 + for slot_idx in range(num_slots): + end = round((slot_idx + 1) * total_seq_len / num_slots) + if end > start: + ranges.append((start, end)) + start = end + return ranges + + +def memory_slot_ranges_from_lengths( + lengths: tuple[int, ...] | None, + *, + total_seq_len: int, + num_slots: int, +) -> list[tuple[int, int]]: + if not lengths or len(lengths) != num_slots: + return memory_slot_ranges(total_seq_len, num_slots) + + ranges: list[tuple[int, int]] = [] + start = 0 + for raw_length in lengths: + length = max(0, int(raw_length)) + end = min(start + length, total_seq_len) + if end > start: + ranges.append((start, end)) + start = end + if start != total_seq_len: + return memory_slot_ranges(total_seq_len, num_slots) + return ranges + + +def build_paired_memory_cross_mask( + *, + batch_size: int, + query_memory_seq_len: int, + query_target_seq_len: int, + kv_memory_seq_len: int, + kv_target_seq_len: int, + num_memory_slots: int, + device: torch.device, + query_segment_lengths: tuple[tuple[int, ...], ...] | None = None, + kv_segment_lengths: tuple[tuple[int, ...], ...] | None = None, +) -> torch.Tensor: + query_total_seq_len = query_memory_seq_len + query_target_seq_len + kv_total_seq_len = kv_memory_seq_len + kv_target_seq_len + mask = torch.zeros( + batch_size, + query_total_seq_len, + kv_total_seq_len, + dtype=torch.bool, + device=device, + ) + + for batch_idx in range(batch_size): + query_lengths = ( + query_segment_lengths[batch_idx] + if query_segment_lengths is not None + and batch_idx < len(query_segment_lengths) + else None + ) + kv_lengths = ( + kv_segment_lengths[batch_idx] + if kv_segment_lengths is not None and batch_idx < len(kv_segment_lengths) + else None + ) + query_ranges = memory_slot_ranges_from_lengths( + query_lengths, + total_seq_len=query_memory_seq_len, + num_slots=num_memory_slots, + ) + kv_ranges = memory_slot_ranges_from_lengths( + kv_lengths, + total_seq_len=kv_memory_seq_len, + num_slots=num_memory_slots, + ) + for (q_start, q_end), (k_start, k_end) in zip( + query_ranges, kv_ranges, strict=False + ): + mask[batch_idx, q_start:q_end, k_start:k_end] = True + + if query_target_seq_len > 0 and kv_target_seq_len > 0: + mask[:, query_memory_seq_len:, kv_memory_seq_len:] = True + return mask + + +def build_memory_self_attention_block_mask( + *, + batch_size: int, + memory_seq_len: int, + target_seq_len: int, + device: torch.device, +) -> torch.Tensor | None: + if memory_seq_len <= 0: + return None + + total_seq_len = memory_seq_len + target_seq_len + attention_mask = torch.ones( + batch_size, + total_seq_len, + total_seq_len, + dtype=torch.bool, + device=device, + ) + attention_mask[:, :, :memory_seq_len] = False + attention_mask[:, :memory_seq_len, :] = False + attention_mask[:, :memory_seq_len, :memory_seq_len] = True + return attention_mask + + +# --- Memory VAE encode --- + + +import torch +from PIL import Image +from torchvision.transforms import functional as TVF + + +def frames_to_video_tensor( + frames: list[Image.Image], target_h: int, target_w: int +) -> torch.Tensor: + tensors = [] + for idx, image in enumerate(frames): + if image.size != (target_w, target_h): + raise ValueError( + f"Frame size mismatch at index {idx}: got={image.size}, " + f"expected={(target_w, target_h)}" + ) + tensor = TVF.to_tensor(image) + tensors.append(tensor * 2.0 - 1.0) + return torch.stack(tensors, dim=1).contiguous() + + +@torch.no_grad() +def encode_memory_frames_batch( + *, + video_vae: Any, + batch_memory_frames: list[list[Image.Image | list[Image.Image]]], + target_h: int, + target_w: int, + device: torch.device, + dtype: torch.dtype, + pipeline_config: Any, +) -> torch.Tensor: + """Encode memory PIL frames to packed video latent tokens [B, S_mem, D].""" + if video_vae.encoder is None: + raise RuntimeError("video VAE encoder is not initialized for memory encoding") + + packed_latents = [] + for memory_frames in batch_memory_frames: + if not memory_frames: + raise ValueError("memory_frames cannot be empty when encoding memory video") + per_slot_latents = [] + for memory_item in memory_frames: + is_clip_memory = isinstance(memory_item, list) + frame_video = ( + frames_to_video_tensor( + memory_item if is_clip_memory else [memory_item], + target_h, + target_w, + ) + .unsqueeze(0) + .to(device=device, dtype=dtype) + ) + encode_output = video_vae.encode(frame_video, return_dict=True) + latent = encode_output.sample + if latent.ndim != 5: + raise ValueError( + f"Expected encoded memory latent [B,C,F,H,W], got {tuple(latent.shape)}" + ) + # Align with official encode_memory_frames_batch layout [B, F, C, H, W]. + latent = latent.permute(0, 2, 1, 3, 4).contiguous() + if is_clip_memory: + latent = latent[:, -1:, :, :, :].contiguous() + latent = latent.permute(0, 2, 1, 3, 4).contiguous() + packed = pipeline_config.maybe_pack_latents(latent, 1, None) + per_slot_latents.append(packed) + packed_latents.append(torch.cat(per_slot_latents, dim=1)) + + return torch.cat(packed_latents, dim=0) + + +# --- Memory RoPE coordinates --- + + +# Official ltx_wrapper hardcodes VIDEO_FPS=24.0 for RoPE position conversion. +JOYAI_VIDEO_ROPE_FPS = 24.0 + + +def normalize_memory_position_mode(mode: str) -> str: + normalized = str(mode).lower() + if normalized == "reference": + return "legacy" + if normalized not in {"legacy", "prefix_continuous"}: + raise ValueError( + "memory_position_mode must be one of " + "{'reference', 'legacy', 'prefix_continuous'}, " + f"got {mode}" + ) + return normalized + + +def apply_memory_video_downscale( + video_coords: torch.Tensor, + downscale_factor: int, +) -> torch.Tensor: + if int(downscale_factor) == 1: + return video_coords + scaled = video_coords.clone() + scaled[:, 1, ...] *= int(downscale_factor) + scaled[:, 2, ...] *= int(downscale_factor) + return scaled + + +def build_memory_video_rope_coords( + *, + rope, + batch_size: int, + memory_video_len: int, + target_num_frames: int, + latent_height: int, + latent_width: int, + device: torch.device, + fps: float, + memory_position_mode: str, + memory_downscale_factor: int = 1, + sp_target_start_offset: int = 0, +) -> torch.Tensor: + """Build [memory | target] video RoPE coordinates. + + Under sequence parallelism the target video latents are time-sharded, so + ``target_num_frames`` is the *local* shard frame count and + ``sp_target_start_offset`` is the global frame index of this rank's first + target frame. The memory prefix is replicated (full) on every rank. + """ + tokens_per_latent_frame = int(latent_height) * int(latent_width) + if tokens_per_latent_frame <= 0: + raise ValueError( + f"Invalid latent grid for memory RoPE: {latent_height=} {latent_width=}" + ) + if memory_video_len % tokens_per_latent_frame != 0: + raise ValueError( + "memory_video_len must be a multiple of latent_height * latent_width, " + f"got {memory_video_len=} {latent_height=} {latent_width=}" + ) + + memory_latent_frames = memory_video_len // tokens_per_latent_frame + position_mode = normalize_memory_position_mode(memory_position_mode) + + memory_coords = rope.prepare_video_coords( + batch_size=batch_size, + num_frames=memory_latent_frames, + height=latent_height, + width=latent_width, + device=device, + fps=JOYAI_VIDEO_ROPE_FPS, + start_frame=0, + ) + memory_coords = apply_memory_video_downscale(memory_coords, memory_downscale_factor) + + target_start_frame = ( + memory_latent_frames if position_mode == "prefix_continuous" else 0 + ) + int(sp_target_start_offset) + target_coords = rope.prepare_video_coords( + batch_size=batch_size, + num_frames=target_num_frames, + height=latent_height, + width=latent_width, + device=device, + fps=JOYAI_VIDEO_ROPE_FPS, + start_frame=target_start_frame, + ) + return torch.cat([memory_coords, target_coords], dim=2) + + +def build_memory_audio_rope_coords( + *, + audio_rope, + batch_size: int, + memory_audio_len: int, + target_audio_len: int, + device: torch.device, + memory_position_mode: str, +) -> torch.Tensor: + position_mode = normalize_memory_position_mode(memory_position_mode) + + memory_coords = audio_rope.prepare_audio_coords( + batch_size=batch_size, + num_frames=memory_audio_len, + device=device, + start_frame=0, + ) + target_start_frame = memory_audio_len if position_mode == "prefix_continuous" else 0 + target_coords = audio_rope.prepare_audio_coords( + batch_size=batch_size, + num_frames=target_audio_len, + device=device, + start_frame=target_start_frame, + ) + return torch.cat([memory_coords, target_coords], dim=2) + + +# --- Paired audio-video memory bank --- + + +import torch +import torchaudio +from PIL import Image + + +@dataclass +class MemoryEntry: + frame: Image.Image | list[Image.Image] + audio_latent: Optional[torch.Tensor] = None + metadata: dict[str, Any] = field(default_factory=dict) + + +def video_uint8_to_pil_frames(video_uint8: torch.Tensor) -> list[Image.Image]: + if video_uint8.ndim != 4: + raise ValueError( + f"Expected [F, H, W, C] uint8 video, got shape={tuple(video_uint8.shape)}" + ) + if video_uint8.shape[-1] != 3: + raise ValueError( + f"Expected RGB video with trailing channel dim 3, got shape={tuple(video_uint8.shape)}" + ) + video_uint8 = video_uint8.detach().cpu().contiguous() + return [Image.fromarray(frame.numpy()) for frame in video_uint8] + + +def normalize_audio_waveform_for_media( + audio_waveform: Optional[torch.Tensor], +) -> Optional[torch.Tensor]: + if audio_waveform is None: + return None + + waveform = torch.as_tensor(audio_waveform).detach().cpu().float() + + if waveform.ndim == 3: + if waveform.shape[0] != 1: + raise ValueError( + f"Expected batch size 1 for decoded audio, got shape={tuple(waveform.shape)}" + ) + waveform = waveform[0] + if waveform.ndim == 1: + waveform = waveform.unsqueeze(0) + elif ( + waveform.ndim == 2 + and waveform.shape[0] not in {1, 2} + and waveform.shape[1] + in { + 1, + 2, + } + ): + waveform = waveform.transpose(0, 1) + elif waveform.ndim != 2: + raise ValueError( + f"Expected decoded audio with 1, 2, or 3 dims, got shape={tuple(waveform.shape)}" + ) + + if waveform.shape[0] == 1: + waveform = waveform.repeat(2, 1) + elif waveform.shape[0] > 2: + waveform = waveform[:2] + return waveform.contiguous() + + +class PairedAudioVideoMemoryBank: + def __init__(self, max_size: int, num_fix_frames: int = 0) -> None: + self.max_size = int(max_size) + self.num_fix_frames = max(0, int(num_fix_frames)) + self.memory: list[MemoryEntry] = [] + + @staticmethod + def _prepare_audio_latent( + audio_latent: Optional[torch.Tensor], + ) -> Optional[torch.Tensor]: + if audio_latent is None: + return None + if audio_latent.dim() != 3: + raise ValueError( + f"Expected audio_latent shape [B, T, C], got shape={tuple(audio_latent.shape)}" + ) + return audio_latent.detach().cpu().contiguous() + + @staticmethod + def _select_audio_window( + audio_latent: torch.Tensor, window_size: int + ) -> tuple[torch.Tensor, dict[str, Any]]: + total_frames = int(audio_latent.shape[1]) + window_size = max(1, int(window_size)) + window_len = min(total_frames, window_size) + window_start = max((total_frames - window_len) // 2, 0) + window_end = window_start + window_len + metadata = { + "audio_window_start": int(window_start), + "audio_window_end": int(window_end), + "audio_window_length": int(window_len), + "audio_total_frames": int(total_frames), + } + return audio_latent[:, window_start:window_end].contiguous(), metadata + + @staticmethod + def _waveform_to_mel( + waveform: torch.Tensor, + *, + sample_rate: int, + mel_bins: int, + mel_hop_length: int, + n_fft: int, + ) -> torch.Tensor: + mono = waveform.mean(dim=0, keepdim=True) + mel_transform = torchaudio.transforms.MelSpectrogram( + sample_rate=sample_rate, + n_fft=n_fft, + hop_length=mel_hop_length, + n_mels=mel_bins, + center=True, + power=1.0, + ) + mel = mel_transform(mono) + return mel.unsqueeze(0) + + @staticmethod + def _select_video_clip_for_audio_window( + frames: list[Image.Image], + *, + audio_window_start: int, + audio_window_end: int, + audio_total_frames: int, + video_clip_num_frames: int, + ) -> tuple[list[Image.Image], dict[str, Any]]: + audio_total_frames = max(1, int(audio_total_frames)) + window_center = (float(audio_window_start) + float(audio_window_end - 1)) * 0.5 + center_ratio = window_center / float(max(audio_total_frames - 1, 1)) + center_frame = int(round(center_ratio * float(max(len(frames) - 1, 0)))) + return PairedAudioVideoMemoryBank._select_video_clip_around_frame( + frames, + center_frame=center_frame, + video_clip_num_frames=video_clip_num_frames, + ) + + @staticmethod + def _select_video_clip_around_frame( + frames: list[Image.Image], + *, + center_frame: int, + video_clip_num_frames: int, + ) -> tuple[list[Image.Image], dict[str, Any]]: + video_clip_num_frames = max(1, int(video_clip_num_frames)) + center_frame = max(0, min(int(center_frame), len(frames) - 1)) + left_context = (video_clip_num_frames - 1) // 2 + clip_start = max( + 0, + min( + center_frame - left_context, max(len(frames) - video_clip_num_frames, 0) + ), + ) + clip_end = min(clip_start + video_clip_num_frames, len(frames)) + clip = list(frames[clip_start:clip_end]) + if clip and len(clip) < video_clip_num_frames: + clip.extend([clip[-1]] * (video_clip_num_frames - len(clip))) + metadata = { + "video_clip_start": int(clip_start), + "video_clip_end": int(clip_end), + "video_clip_length": int(len(clip)), + "video_clip_center_frame": int(center_frame), + "video_total_frames": int(len(frames)), + } + return clip, metadata + + def _trim(self) -> None: + if self.max_size <= 0 or len(self.memory) <= self.max_size: + return + fixed = self.memory[: self.num_fix_frames] + tail = self.memory[self.num_fix_frames :] + keep_tail = max(0, self.max_size - len(fixed)) + self.memory = fixed + tail[-keep_tail:] + + def save_memory_slot( + self, + frames: list[Image.Image], + audio_latent: torch.Tensor, + *, + audio_window_size: int, + video_clip_num_frames: int, + audio_waveform: Optional[torch.Tensor] = None, + audio_sample_rate: int = 16000, + video_fps: float = 25.0, + audio_window_selection_mode: str = "max_response", + video_frame_selection_mode: str = "center", + audio_memory_mel_bins: int = 128, + audio_memory_mel_hop_length: int = 160, + audio_memory_n_fft: int = 1024, + audio_memory_downsample_factor: int = 4, + audio_memory_is_causal: bool = True, + ) -> dict[str, Any]: + audio_latent = self._prepare_audio_latent(audio_latent) + if audio_latent is None: + raise ValueError("paired audio memory slot requires audio_latent") + + selection_mode = str(audio_window_selection_mode).lower() + if audio_waveform is not None and selection_mode != "center": + try: + waveform = normalize_audio_waveform_for_media(audio_waveform) + mel = self._waveform_to_mel( + waveform, + sample_rate=audio_sample_rate, + mel_bins=audio_memory_mel_bins, + mel_hop_length=audio_memory_mel_hop_length, + n_fft=audio_memory_n_fft, + ) + pixel_window_size = latent_window_size_to_pixel_window_size( + int(audio_window_size), + downsample_factor=int(audio_memory_downsample_factor), + is_causal=bool(audio_memory_is_causal), + ) + _, window_start_indices, window_end_indices = ( + select_audio_window_with_bounds( + mel.float(), + pixel_window_size, + mode="max_response", + ) + ) + mel_start = int(window_start_indices[0].item()) + mel_end = int(window_end_indices[0].item()) + start_time_sec, end_time_sec = mel_window_bounds_to_seconds( + mel_start, + mel_end, + hop_length=int(audio_memory_mel_hop_length), + sample_rate=int(audio_sample_rate), + ) + total_frames = int(audio_latent.shape[1]) + window_len = min(total_frames, max(1, int(audio_window_size))) + duration_sec = max( + float(waveform.shape[-1]) / float(audio_sample_rate), 1e-6 + ) + center_time_sec = max( + 0.0, min(0.5 * (start_time_sec + end_time_sec), duration_sec) + ) + center_latent = int( + round( + center_time_sec / duration_sec * float(max(total_frames - 1, 0)) + ) + ) + window_start = max( + 0, + min( + center_latent - window_len // 2, + max(total_frames - window_len, 0), + ), + ) + window_end = window_start + window_len + window_latent = audio_latent[:, window_start:window_end].contiguous() + audio_metadata = { + "audio_window_selection_mode": selection_mode, + "audio_window_start": int(window_start), + "audio_window_end": int(window_end), + "audio_window_length": int(window_len), + "audio_total_frames": int(total_frames), + "mel_window_start": int(mel_start), + "mel_window_end": int(mel_end), + "audio_window_start_time_sec": float(start_time_sec), + "audio_window_end_time_sec": float(end_time_sec), + } + selected_frame = select_video_frame_indices_from_time_range( + num_frames=len(frames), + fps=float(video_fps), + start_time_sec=float(start_time_sec), + end_time_sec=float(end_time_sec), + count=1, + mode=str(video_frame_selection_mode).lower(), + )[0] + video_clip, video_metadata = self._select_video_clip_around_frame( + frames, + center_frame=int(selected_frame), + video_clip_num_frames=video_clip_num_frames, + ) + except Exception as exc: + window_latent, audio_metadata = self._select_audio_window( + audio_latent, audio_window_size + ) + audio_metadata["audio_window_selection_mode"] = "center" + audio_metadata["selection_fallback"] = f"{selection_mode}: {exc}" + video_clip, video_metadata = self._select_video_clip_for_audio_window( + frames, + audio_window_start=int(audio_metadata["audio_window_start"]), + audio_window_end=int(audio_metadata["audio_window_end"]), + audio_total_frames=int(audio_metadata["audio_total_frames"]), + video_clip_num_frames=video_clip_num_frames, + ) + else: + window_latent, audio_metadata = self._select_audio_window( + audio_latent, audio_window_size + ) + audio_metadata["audio_window_selection_mode"] = "center" + video_clip, video_metadata = self._select_video_clip_for_audio_window( + frames, + audio_window_start=int(audio_metadata["audio_window_start"]), + audio_window_end=int(audio_metadata["audio_window_end"]), + audio_total_frames=int(audio_metadata["audio_total_frames"]), + video_clip_num_frames=video_clip_num_frames, + ) + + metadata = { + "selection_mode": "paired_audio_window", + **audio_metadata, + **video_metadata, + } + entry = MemoryEntry( + frame=video_clip, audio_latent=window_latent, metadata=metadata + ) + fixed = self.memory[: self.num_fix_frames] + free = self.memory[self.num_fix_frames :] + free.append(entry) + self.memory = fixed + free + self._trim() + return metadata + + def get_memory_frames(self) -> list[Image.Image | list[Image.Image]]: + return [entry.frame for entry in self.memory] + + def get_memory_audio(self) -> Optional[torch.Tensor]: + audio_latents = [entry.audio_latent for entry in self.memory] + if not audio_latents or any(item is None for item in audio_latents): + return None + first = audio_latents[0] + assert first is not None + for audio_latent in audio_latents: + assert audio_latent is not None + if ( + audio_latent.shape[0] != first.shape[0] + or audio_latent.shape[2] != first.shape[2] + ): + raise ValueError( + "All memory audio latents must share batch and channel dimensions" + ) + return torch.cat(audio_latents, dim=1).contiguous() + + def get_memory_audio_segment_lengths(self) -> tuple[tuple[int, ...], ...]: + audio_latents = [entry.audio_latent for entry in self.memory] + if not audio_latents or any(item is None for item in audio_latents): + return () + return ( + tuple( + int(audio_latent.shape[1]) + for audio_latent in audio_latents + if audio_latent is not None + ), + ) + + def __len__(self) -> int: + return len(self.memory) + + +# --- Pipeline stages --- + + +class JoyEchoMemoryBankFetchStage(PipelineStage): + """Prepare memory video/audio prefixes before DMD denoising.""" + + def __init__(self, memory_bank: PairedAudioVideoMemoryBank, vae) -> None: + super().__init__() + self.memory_bank = memory_bank + self.vae = vae + + @staticmethod + def _resolve_vae_encode_dtype(vae) -> torch.dtype: + encoder = vae.encoder + if encoder is not None: + try: + return next(encoder.parameters()).dtype + except StopIteration: + pass + return torch.bfloat16 + + def forward(self, batch: Req, server_args: ServerArgs) -> Req: + batch.extra.pop("joy_echo_memory", None) + config = server_args.pipeline_config + if not isinstance(config, JoyEchoPipelineConfig): + return batch + + if not batch.enable_memory_bank or len(self.memory_bank) == 0: + return batch + + device = get_local_torch_device() + vae_dtype = self._resolve_vae_encode_dtype(self.vae) + latent_dtype = batch.latents.dtype if batch.latents is not None else vae_dtype + + memory_frames = self.memory_bank.get_memory_frames() + memory_video = encode_memory_frames_batch( + video_vae=self.vae, + batch_memory_frames=[memory_frames], + target_h=int(batch.height), + target_w=int(batch.width), + device=device, + dtype=vae_dtype, + pipeline_config=config, + ) + + memory_audio = self.memory_bank.get_memory_audio() + if memory_audio is None: + raise RuntimeError( + "JoyEcho memory bank has video frames but no audio latents" + ) + + memory_audio = memory_audio.to(device=device, dtype=latent_dtype) + batch.extra["joy_echo_memory"] = { + "memory_video_packed": memory_video, + "memory_audio": memory_audio, + "memory_audio_segment_lengths": self.memory_bank.get_memory_audio_segment_lengths(), + "num_memory_slots": len(self.memory_bank), + "paired_audio_memory": True, + "memory_position_mode": str(config.memory_position_mode), + "memory_downscale_factor": 1, + } + logger.info( + "JoyEcho memory fetch: bank_size=%d video_tokens=%d audio_tokens=%d", + len(self.memory_bank), + int(memory_video.shape[1]), + int(memory_audio.shape[1]), + ) + return batch + + +class JoyEchoAVDecodingStage(LTX2AVDecodingStage): + """Decode AV outputs and commit paired memory slots when enabled.""" + + def __init__( + self, + vae, + audio_vae, + vocoder, + memory_bank: PairedAudioVideoMemoryBank, + pipeline=None, + ): + super().__init__(vae, audio_vae, vocoder, pipeline=pipeline) + self.memory_bank = memory_bank + + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: + audio_latent_for_memory = None + if batch.audio_latents is not None: + audio_latent_for_memory = batch.audio_latents.detach().cpu() + # Denoising unpacks audio to [B, C, L, M] before decode; memory bank + # expects the packed [B, L, C] layout used during generation. + if audio_latent_for_memory.dim() == 4: + audio_latent_for_memory = ( + server_args.pipeline_config.maybe_pack_audio_latents( + audio_latent_for_memory, batch.batch_size, batch + ).cpu() + ) + + output_batch = super().forward(batch, server_args) + + config = server_args.pipeline_config + if not isinstance(config, JoyEchoPipelineConfig): + return output_batch + + if not batch.enable_memory_bank or audio_latent_for_memory is None: + return output_batch + + video_np = output_batch.output + if video_np is None: + return output_batch + + if isinstance(video_np, np.ndarray): + if video_np.ndim == 5: + video_uint8 = torch.from_numpy(video_np[0]) + elif video_np.ndim == 4: + video_uint8 = torch.from_numpy(video_np) + else: + logger.warning( + "Unexpected decoded video shape for memory commit: %s", + video_np.shape, + ) + return output_batch + else: + logger.warning( + "Unsupported decoded video type for memory commit: %s", type(video_np) + ) + return output_batch + + if video_uint8.dtype != torch.uint8: + video_uint8 = ( + (video_uint8.clamp(0, 1) * 255).to(torch.uint8) + if video_uint8.is_floating_point() + else video_uint8.to(torch.uint8) + ) + + pil_frames = video_uint8_to_pil_frames(video_uint8) + metadata = self.memory_bank.save_memory_slot( + pil_frames, + audio_latent_for_memory, + audio_window_size=int(config.audio_window_size), + video_clip_num_frames=int(config.memory_video_clip_num_frames), + audio_waveform=output_batch.audio, + audio_sample_rate=int( + output_batch.audio_sample_rate + or config.audio_vae_config.arch_config.sample_rate + ), + video_fps=float(batch.fps), + audio_window_selection_mode=str(config.audio_window_selection_mode), + video_frame_selection_mode=str(config.video_memory_frame_selection_mode), + audio_memory_mel_bins=int(config.audio_mel_bins), + audio_memory_mel_hop_length=int(config.audio_mel_hop_length), + audio_memory_n_fft=int(config.audio_n_fft), + audio_memory_downsample_factor=int(config.audio_downsample_factor), + audio_memory_is_causal=True, + ) + logger.info( + "JoyEcho memory commit: bank_size=%d metadata=%s", + len(self.memory_bank), + metadata, + ) + return output_batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/setup.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/setup.py new file mode 100644 index 000000000..5d91c6761 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/model_specific_stages/joy_echo/setup.py @@ -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 diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index 566ea3d88..30dee0066 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -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", diff --git a/python/sglang/multimodal_gen/test/server/consistency_threshold.json b/python/sglang/multimodal_gen/test/server/consistency_threshold.json index fd9f547ed..8b754e5b9 100644 --- a/python/sglang/multimodal_gen/test/server/consistency_threshold.json +++ b/python/sglang/multimodal_gen/test/server/consistency_threshold.json @@ -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, diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index ecda2ea94..1802383ea 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -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( diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 43064715f..fad5e4198 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -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", diff --git a/python/sglang/multimodal_gen/test/test_utils.py b/python/sglang/multimodal_gen/test/test_utils.py index 6f84df422..9c40c8d84 100644 --- a/python/sglang/multimodal_gen/test/test_utils.py +++ b/python/sglang/multimodal_gen/test/test_utils.py @@ -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"