From 82c41a2d9e0c7a55d001dfcb125b9ab5ef9f662a Mon Sep 17 00:00:00 2001 From: Mick Date: Mon, 6 Apr 2026 12:26:30 +0800 Subject: [PATCH] [diffusion] model: support LTX2.3 (#22111) --- docs/diffusion/compatibility_matrix.md | 5 +- python/sglang/multimodal_gen/README.md | 5 - .../configs/models/adapter/ltx_2_connector.py | 4 + .../configs/models/dits/ltx_2.py | 13 + .../configs/models/vaes/ltx_video.py | 8 +- .../configs/pipeline_configs/ltx_2.py | 36 +- .../multimodal_gen/configs/sample/ltx_2.py | 43 ++ .../configs/sample/sampling_params.py | 12 + .../ltx_2_3/_overlay/materialize.py | 302 ++++++++++ python/sglang/multimodal_gen/registry.py | 32 +- .../entrypoints/diffusion_generator.py | 6 +- .../runtime/entrypoints/utils.py | 7 +- .../component_loaders/adapter_loader.py | 8 +- .../runtime/models/adapter/ltx_2_connector.py | 106 +++- .../runtime/models/dits/ltx_2.py | 245 ++++++-- .../models/vaes/ltx_2_3_condition_encoder.py | 204 +++++++ .../runtime/models/vaes/ltx_2_vae.py | 323 +++++++++- .../runtime/models/vocoder/ltx_2_vocoder.py | 553 ++++++++++++++++++ .../runtime/pipelines/ltx_2_pipeline.py | 65 +- .../runtime/pipelines_core/__init__.py | 1 + .../runtime/pipelines_core/schedule_batch.py | 2 +- .../pipelines_core/stages/decoding_av.py | 12 +- .../pipelines_core/stages/denoising_av.py | 259 +++++--- .../stages/latent_preparation_av.py | 133 ++++- .../runtime/utils/model_overlay.py | 58 +- .../test/server/perf_baselines.json | 48 ++ .../test/server/testcase_configs.py | 10 +- .../test/unit/test_model_overlay_ltx23.py | 392 +++++++++++++ 28 files changed, 2686 insertions(+), 206 deletions(-) create mode 100644 python/sglang/multimodal_gen/model_overlays/ltx_2_3/_overlay/materialize.py create mode 100644 python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_3_condition_encoder.py create mode 100644 python/sglang/multimodal_gen/test/unit/test_model_overlay_ltx23.py diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index 03ebd5098..038e43892 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/docs/diffusion/compatibility_matrix.md @@ -33,12 +33,15 @@ default parameters when initializing and generating videos. | TurboWan2.1 T2V 14B | `IPostYellow/TurboWan2.1-T2V-14B-Diffusers` | 480p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.1 T2V 14B 720P | `IPostYellow/TurboWan2.1-T2V-14B-720P-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | | TurboWan2.2 I2V A14B | `IPostYellow/TurboWan2.2-I2V-A14B-Diffusers` | 720p | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| LTX-2 | `Lightricks/LTX-2` | 1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| LTX-2 | `Lightricks/LTX-2` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| LTX-2.3 | `Lightricks/LTX-2.3` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | **Note**: 1. Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. 2. SageSLA is based on SpargeAttn. Install it first with `pip install git+https://github.com/thu-ml/SpargeAttn.git --no-build-isolation` 3. LTX-2 two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. The spatial upsampler and distilled LoRA are auto-resolved from the model snapshot by default, and can still be overridden with `--spatial-upsampler-path` and `--distilled-lora-path`. +4. `Lightricks/LTX-2.3` is supported through the bundled native overlay materialization path. One-stage generation uses the default `LTX2Pipeline`; two-stage generation uses `--pipeline-class-name LTX2TwoStagePipeline`. +5. For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`. One-stage generation is validated at `768×512`; two-stage generation is validated at `1536×1024`. ### Image Generation Models diff --git a/python/sglang/multimodal_gen/README.md b/python/sglang/multimodal_gen/README.md index 471b91a1a..385383a70 100644 --- a/python/sglang/multimodal_gen/README.md +++ b/python/sglang/multimodal_gen/README.md @@ -76,11 +76,6 @@ sglang generate --model-path Wan-AI/Wan2.1-T2V-1.3B-Diffusers \ --save-output ``` -For LTX-2 two-stage generation, use `--pipeline-class-name LTX2TwoStagePipeline`. The -spatial upsampler and distilled LoRA are auto-resolved from the same model snapshot by -default, and can still be overridden with `--spatial-upsampler-path` and -`--distilled-lora-path` when needed. - ### LoRA support Apply LoRA adapters via `--lora-path`: diff --git a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py index d8f6c43ef..03e2bcf9b 100644 --- a/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/configs/models/adapter/ltx_2_connector.py @@ -12,13 +12,17 @@ class LTX2ConnectorArchConfig(AdapterArchConfig): audio_connector_num_attention_heads: int = 30 audio_connector_num_layers: int = 2 audio_connector_num_learnable_registers: int = 128 + audio_feature_extractor_out_features: int = 0 caption_channels: int = 3840 causal_temporal_positioning: bool = False connector_rope_base_seq_len: int = 4096 + connector_apply_gated_attention: bool = False + feature_extractor_in_features: int = 0 rope_double_precision: bool = True rope_theta: float = 10000.0 rope_type: str = "split" text_proj_in_factor: int = 49 + video_feature_extractor_out_features: int = 0 video_connector_attention_head_dim: int = 128 video_connector_num_attention_heads: int = 30 video_connector_num_layers: int = 2 diff --git a/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py b/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py index d2bee4ba0..5378883a1 100644 --- a/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/models/dits/ltx_2.py @@ -63,6 +63,7 @@ class LTX2ArchConfig(DiTArchConfig): # We use upstream variable names (patchify_proj, adaln_single) but HF uses different keys. # # HF key -> SGLang key (upstream naming) + r"^model\.diffusion_model\.(.*)$": r"\1", r"^proj_in\.(.*)$": r"patchify_proj.\1", r"^time_embed\.(.*)$": r"adaln_single.\1", r"^audio_proj_in\.(.*)$": r"audio_patchify_proj.\1", @@ -123,6 +124,10 @@ class LTX2ArchConfig(DiTArchConfig): attention_type: LTX2AttentionFunction = LTX2AttentionFunction.DEFAULT rope_type: LTX2RopeType = LTX2RopeType.INTERLEAVED double_precision_rope: bool = False + quantize_video_rope_coords_to_hidden_dtype: bool = False + apply_gated_attention: bool = False + cross_attention_adaln: bool = False + caption_proj_before_connector: bool = False # Video parameters num_attention_heads: int = 32 @@ -147,6 +152,14 @@ class LTX2ArchConfig(DiTArchConfig): audio_positional_embedding_max_pos: list[int] | None = None av_ca_timestep_scale_multiplier: int = 1 + # 2.3 connector-related fields may show up in transformer/config.json. + connector_attention_head_dim: int = 128 + connector_num_attention_heads: int = 30 + connector_num_layers: int = 2 + audio_connector_attention_head_dim: int = 128 + audio_connector_num_attention_heads: int = 30 + audio_connector_num_layers: int = 2 + # SGLang-specific parameters patch_size: tuple[int, int, int] = (1, 2, 2) text_len: int = 512 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 02964d48a..92d29537d 100644 --- a/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py +++ b/python/sglang/multimodal_gen/configs/models/vaes/ltx_video.py @@ -1,6 +1,6 @@ # SPDX-License-Identifier: Apache-2.0 from dataclasses import dataclass, field -from typing import List +from typing import Any, List from sglang.multimodal_gen.configs.models.vaes.base import VAEArchConfig, VAEConfig @@ -52,6 +52,12 @@ class LTXVideoVAEArchConfig(VAEArchConfig): decoder_causal: bool = False decoder_spatial_padding_mode: str = "reflect" + # Native LTX variant metadata. + ltx_variant: str = "ltx_2" + condition_encoder_subdir: str = "" + video_decoder_variant: str = "ltx_2" + video_decoder_config: dict[str, Any] = field(default_factory=dict) + @dataclass class LTXVideoVAEConfig(VAEConfig): 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 827633c3d..5029e68ba 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -93,20 +93,48 @@ def pack_text_embeds( return normalized_hidden_states +def pack_text_embeds_v2( + text_hidden_states: torch.Tensor, + attention_mask: torch.Tensor, + eps: float = 1e-6, +) -> torch.Tensor: + """ + LTX-2.3 feature extractor pre-processing. + + Upstream `FeatureExtractorV2` applies per-token RMS normalization on each + Gemma layer and then flattens `[hidden_dim, num_layers]` into the channel + dimension, zeroing out padded positions afterwards. + """ + + variance = torch.mean(text_hidden_states**2, dim=2, keepdim=True) + normalized_hidden_states = text_hidden_states * torch.rsqrt(variance + eps) + normalized_hidden_states = normalized_hidden_states.flatten(2) + mask = attention_mask.bool().unsqueeze(-1) + return torch.where( + mask, normalized_hidden_states, torch.zeros_like(normalized_hidden_states) + ) + + +def is_ltx23_native_variant(arch_config: object) -> bool: + return str(getattr(arch_config, "ltx_variant", "ltx_2")) == "ltx_2_3" + + def _gemma_postprocess_func( outputs: BaseEncoderOutput, text_inputs: dict, pipeline_config: Optional["LTX2PipelineConfig"] = None, ) -> torch.Tensor: - _ = pipeline_config # LTX-2 requires all hidden states concatenated for the connector if hasattr(outputs, "hidden_states") and outputs.hidden_states is not None: - # outputs.hidden_states is a tuple of tensors - # We need to stack them along the last dimension and pack them hidden_states = torch.stack(outputs.hidden_states, dim=-1) attention_mask = text_inputs["attention_mask"] + if ( + pipeline_config is not None + and pipeline_config.dit_config.arch_config.caption_proj_before_connector + ): + return pack_text_embeds_v2(hidden_states, attention_mask) + sequence_lengths = attention_mask.sum(dim=-1) - # Assuming left padding for Gemma as per Diffusers return pack_text_embeds(hidden_states, sequence_lengths, padding_side="left") else: raise AttributeError( diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index ef0b35981..dec4f918d 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -1,4 +1,6 @@ import dataclasses +from dataclasses import field +from typing import Any from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams @@ -39,3 +41,44 @@ class LTX2SamplingParams(SamplingParams): "pauses, incorrect timing, unnatural transitions, inconsistent framing, tilted camera, flat lighting, " "inconsistent tone, cinematic oversaturation, stylized filters, or AI artifacts." ) + + +@dataclasses.dataclass +class LTX23SamplingParams(LTX2SamplingParams): + """Sampling parameters matching official LTX-2.3 one-stage defaults.""" + + generator_device: str = "cuda" + guidance_scale: float = 3.0 + num_inference_steps: int = 30 + + video_cfg_scale: float = 3.0 + video_stg_scale: float = 1.0 + video_rescale_scale: float = 0.7 + video_modality_scale: float = 3.0 + video_skip_step: int = 0 + video_stg_blocks: list[int] = field(default_factory=lambda: [28]) + + audio_cfg_scale: float = 7.0 + audio_stg_scale: float = 1.0 + audio_rescale_scale: float = 0.7 + audio_modality_scale: float = 3.0 + audio_skip_step: int = 0 + audio_stg_blocks: list[int] = field(default_factory=lambda: [28]) + + def build_request_extra(self) -> dict[str, Any]: + extra = super().build_request_extra() + extra["ltx2_stage1_guider_params"] = { + "video_cfg_scale": self.video_cfg_scale, + "video_stg_scale": self.video_stg_scale, + "video_rescale_scale": self.video_rescale_scale, + "video_modality_scale": self.video_modality_scale, + "video_skip_step": self.video_skip_step, + "video_stg_blocks": self.video_stg_blocks, + "audio_cfg_scale": self.audio_cfg_scale, + "audio_stg_scale": self.audio_stg_scale, + "audio_rescale_scale": self.audio_rescale_scale, + "audio_modality_scale": self.audio_modality_scale, + "audio_skip_step": self.audio_skip_step, + "audio_stg_blocks": self.audio_stg_blocks, + } + return extra diff --git a/python/sglang/multimodal_gen/configs/sample/sampling_params.py b/python/sglang/multimodal_gen/configs/sample/sampling_params.py index 617e2b5e2..053019a41 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -253,6 +253,18 @@ class SamplingParams: if env_steps is not None and self.num_inference_steps is not None: self.num_inference_steps = int(env_steps) + def build_request_extra(self) -> dict[str, Any]: + """Return optional request-scoped extras for downstream pipeline stages.""" + extra = {} + diffusers_kwargs = getattr(self, "diffusers_kwargs", None) + if diffusers_kwargs: + extra["diffusers_kwargs"] = diffusers_kwargs + return extra + + def apply_request_extra(self, req: Any) -> None: + """Merge request extras (model specific, e.g., LTX2.3) into an already-created pipeline request.""" + req.extra.update(self.build_request_extra()) + def _adjust_output_quality(self, output_quality: str, data_type: DataType) -> int: """Convert output_quality string to compression level.""" output_quality_mapper = {"maximum": 100, "high": 90, "medium": 55, "low": 35} diff --git a/python/sglang/multimodal_gen/model_overlays/ltx_2_3/_overlay/materialize.py b/python/sglang/multimodal_gen/model_overlays/ltx_2_3/_overlay/materialize.py new file mode 100644 index 000000000..73fadfc22 --- /dev/null +++ b/python/sglang/multimodal_gen/model_overlays/ltx_2_3/_overlay/materialize.py @@ -0,0 +1,302 @@ +import json +import os + +from huggingface_hub import snapshot_download +from safetensors import safe_open +from safetensors.torch import save_file + +from sglang.multimodal_gen.runtime.utils.model_overlay import ( + _copytree_link_or_copy, + _ensure_dir, + _link_or_copy_file, +) + +AUXILIARY_MODEL_ID = "Lightricks/LTX-2" +CONFIG_DONOR_MODEL_ID = "FastVideo/LTX-2.3-Distilled-Diffusers" + +AUXILIARY_PATTERNS = [ + "audio_vae/**", + "scheduler/**", + "text_encoder/**", + "tokenizer/**", + "vae/config.json", + "vae/diffusion_pytorch_model.safetensors", +] + +CONFIG_DONOR_PATTERNS = [ + "transformer/config.json", + "text_encoder/config.json", + "vae/**", + "vocoder/**", +] + +MONOLITH_PREFIX = "model.diffusion_model." +VIDEO_CONNECTOR_PREFIX = f"{MONOLITH_PREFIX}video_embeddings_connector." +AUDIO_CONNECTOR_PREFIX = f"{MONOLITH_PREFIX}audio_embeddings_connector." +TEXT_PROJ_IN_PREFIX = f"{MONOLITH_PREFIX}text_proj_in." +VIDEO_AGGREGATE_PREFIX = "text_embedding_projection.video_aggregate_embed." +AUDIO_AGGREGATE_PREFIX = "text_embedding_projection.audio_aggregate_embed." + + +def _load_json(path: str) -> dict: + with open(path) as f: + return json.load(f) + + +def _write_json(path: str, payload: dict) -> None: + with open(path, "w") as f: + json.dump(payload, f, indent=2) + f.write("\n") + + +def _rename_connector_key(key: str) -> str | None: + if key.startswith(VIDEO_CONNECTOR_PREFIX): + suffix = key[len(VIDEO_CONNECTOR_PREFIX) :] + suffix = suffix.replace("transformer_1d_blocks", "transformer_blocks") + suffix = suffix.replace(".attn1.q_norm.", ".attn1.norm_q.") + suffix = suffix.replace(".attn1.k_norm.", ".attn1.norm_k.") + return f"video_connector.{suffix}" + if key.startswith(AUDIO_CONNECTOR_PREFIX): + suffix = key[len(AUDIO_CONNECTOR_PREFIX) :] + suffix = suffix.replace("transformer_1d_blocks", "transformer_blocks") + suffix = suffix.replace(".attn1.q_norm.", ".attn1.norm_q.") + suffix = suffix.replace(".attn1.k_norm.", ".attn1.norm_k.") + return f"audio_connector.{suffix}" + if key.startswith(TEXT_PROJ_IN_PREFIX): + return key[len(MONOLITH_PREFIX) :] + if key.startswith(VIDEO_AGGREGATE_PREFIX): + return f"video_aggregate_embed.{key[len(VIDEO_AGGREGATE_PREFIX):]}" + if key.startswith(AUDIO_AGGREGATE_PREFIX): + return f"audio_aggregate_embed.{key[len(AUDIO_AGGREGATE_PREFIX):]}" + return None + + +def _repack_transformer_weights(source_path: str, output_path: str) -> None: + tensors = {} + with safe_open(source_path, framework="pt") as f: + for key in f.keys(): + if not key.startswith(MONOLITH_PREFIX): + continue + if key.startswith(VIDEO_CONNECTOR_PREFIX): + continue + if key.startswith(AUDIO_CONNECTOR_PREFIX): + continue + if key.startswith(TEXT_PROJ_IN_PREFIX): + continue + tensors[key[len(MONOLITH_PREFIX) :]] = f.get_tensor(key) + if not tensors: + raise ValueError("No transformer tensors found in LTX-2.3 source checkpoint.") + save_file(tensors, output_path) + + +def _repack_connectors_weights(source_path: str, output_path: str) -> None: + tensors = {} + with safe_open(source_path, framework="pt") as f: + for key in f.keys(): + renamed = _rename_connector_key(key) + if renamed is None: + continue + tensors[renamed] = f.get_tensor(key) + if not tensors: + raise ValueError("No connector tensors found in LTX-2.3 source checkpoint.") + save_file(tensors, output_path) + + +def _build_transformer_config(config_donor_dir: str) -> dict: + config = _load_json(os.path.join(config_donor_dir, "transformer", "config.json")) + config["_class_name"] = "LTX2VideoTransformer3DModel" + config["force_sdpa_v2a_cross_attention"] = True + config["quantize_video_rope_coords_to_hidden_dtype"] = True + return config + + +def _build_connectors_config(config_donor_dir: str) -> dict: + text_encoder_config = _load_json( + os.path.join(config_donor_dir, "text_encoder", "config.json") + ) + return { + "_class_name": "LTX2TextConnectors", + "_diffusers_version": "0.37.0.dev0", + "audio_connector_attention_head_dim": text_encoder_config[ + "audio_connector_attention_head_dim" + ], + "audio_connector_num_attention_heads": text_encoder_config[ + "audio_connector_num_attention_heads" + ], + "audio_connector_num_layers": text_encoder_config["audio_connector_num_layers"], + "audio_connector_num_learnable_registers": text_encoder_config[ + "connector_num_learnable_registers" + ], + "audio_feature_extractor_out_features": text_encoder_config[ + "audio_feature_extractor_out_features" + ], + "caption_channels": text_encoder_config["hidden_size"], + "causal_temporal_positioning": False, + "connector_apply_gated_attention": text_encoder_config[ + "connector_apply_gated_attention" + ], + "feature_extractor_in_features": text_encoder_config[ + "feature_extractor_in_features" + ], + "connector_rope_base_seq_len": text_encoder_config[ + "connector_positional_embedding_max_pos" + ][0], + "rope_double_precision": text_encoder_config["connector_double_precision_rope"], + "rope_theta": text_encoder_config["connector_positional_embedding_theta"], + "rope_type": text_encoder_config["connector_rope_type"], + "text_proj_in_factor": text_encoder_config["feature_extractor_in_features"] + // text_encoder_config["hidden_size"], + "video_feature_extractor_out_features": text_encoder_config[ + "video_feature_extractor_out_features" + ], + "video_connector_attention_head_dim": text_encoder_config[ + "connector_attention_head_dim" + ], + "video_connector_num_attention_heads": text_encoder_config[ + "connector_num_attention_heads" + ], + "video_connector_num_layers": text_encoder_config["connector_num_layers"], + "video_connector_num_learnable_registers": text_encoder_config[ + "connector_num_learnable_registers" + ], + } + + +def _build_vae_config(auxiliary_dir: str, config_donor_dir: str) -> dict: + config = _load_json(os.path.join(auxiliary_dir, "vae", "config.json")) + config["ltx_variant"] = "ltx_2_3" + config["condition_encoder_subdir"] = "ltx23_image_encoder" + config["video_decoder_variant"] = "ltx_2_3" + config["video_decoder_config"] = _load_json( + os.path.join(config_donor_dir, "vae", "config.json") + )["vae"] + return config + + +def _repack_ltx23_image_encoder_weights(source_path: str, output_path: str) -> None: + tensors = {} + with safe_open(source_path, framework="pt") as f: + for key in f.keys(): + if key.startswith("encoder."): + tensors[key[len("encoder.") :]] = f.get_tensor(key) + continue + if key.startswith("per_channel_statistics."): + tensors[key] = f.get_tensor(key) + if not tensors: + raise ValueError("No LTX-2.3 image-encoder tensors found in donor checkpoint.") + save_file(tensors, output_path) + + +def _repack_ltx23_video_decoder_weights( + auxiliary_encoder_path: str, + donor_decoder_path: str, + output_path: str, +) -> None: + tensors = {} + with safe_open(auxiliary_encoder_path, framework="pt") as f: + for key in f.keys(): + if key.startswith("encoder."): + tensors[key] = f.get_tensor(key) + with safe_open(donor_decoder_path, framework="pt") as f: + for key in f.keys(): + if key.startswith("decoder."): + tensors[key] = f.get_tensor(key) + continue + if key == "per_channel_statistics.mean-of-means": + tensor = f.get_tensor(key) + tensors["decoder.per_channel_statistics.mean_of_means"] = tensor + tensors["latents_mean"] = tensor.clone() + continue + if key == "per_channel_statistics.std-of-means": + tensor = f.get_tensor(key) + tensors["decoder.per_channel_statistics.std_of_means"] = tensor + tensors["latents_std"] = tensor.clone() + continue + if not tensors: + raise ValueError("No LTX-2.3 decoder tensors found in donor checkpoint.") + save_file(tensors, output_path) + + +def materialize( + *, + overlay_dir: str, + source_dir: str, + output_dir: str, + manifest: dict, +) -> None: + _ = overlay_dir, manifest + + auxiliary_dir = snapshot_download( + repo_id=AUXILIARY_MODEL_ID, + allow_patterns=AUXILIARY_PATTERNS, + max_workers=8, + ) + config_donor_dir = snapshot_download( + repo_id=CONFIG_DONOR_MODEL_ID, + allow_patterns=CONFIG_DONOR_PATTERNS, + max_workers=8, + ) + + for component_name in ("audio_vae", "scheduler", "text_encoder", "tokenizer"): + _copytree_link_or_copy( + os.path.join(auxiliary_dir, component_name), + os.path.join(output_dir, component_name), + ) + _copytree_link_or_copy( + os.path.join(config_donor_dir, "vocoder"), + os.path.join(output_dir, "vocoder"), + ) + + source_checkpoint = os.path.join(source_dir, "ltx-2.3-22b-dev.safetensors") + + transformer_dir = os.path.join(output_dir, "transformer") + _ensure_dir(transformer_dir) + _write_json( + os.path.join(transformer_dir, "config.json"), + _build_transformer_config(config_donor_dir), + ) + _repack_transformer_weights( + source_checkpoint, os.path.join(transformer_dir, "model.safetensors") + ) + + connectors_dir = os.path.join(output_dir, "connectors") + _ensure_dir(connectors_dir) + _write_json( + os.path.join(connectors_dir, "config.json"), + _build_connectors_config(config_donor_dir), + ) + _repack_connectors_weights( + source_checkpoint, os.path.join(connectors_dir, "model.safetensors") + ) + + vae_dir = os.path.join(output_dir, "vae") + _ensure_dir(vae_dir) + _write_json( + os.path.join(vae_dir, "config.json"), + _build_vae_config(auxiliary_dir, config_donor_dir), + ) + _repack_ltx23_video_decoder_weights( + os.path.join(auxiliary_dir, "vae", "diffusion_pytorch_model.safetensors"), + os.path.join(config_donor_dir, "vae", "model.safetensors"), + os.path.join(vae_dir, "model.safetensors"), + ) + + image_encoder_dir = os.path.join(vae_dir, "ltx23_image_encoder") + _ensure_dir(image_encoder_dir) + _link_or_copy_file( + os.path.join(config_donor_dir, "vae", "config.json"), + os.path.join(image_encoder_dir, "config.json"), + ) + _repack_ltx23_image_encoder_weights( + os.path.join(config_donor_dir, "vae", "model.safetensors"), + os.path.join(image_encoder_dir, "model.safetensors"), + ) + + _link_or_copy_file( + os.path.join(source_dir, "ltx-2.3-22b-distilled-lora-384.safetensors"), + os.path.join(output_dir, "ltx-2.3-22b-distilled-lora-384.safetensors"), + ) + _link_or_copy_file( + os.path.join(source_dir, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"), + os.path.join(output_dir, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"), + ) diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index bfe4e97bf..a3928865e 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -90,7 +90,10 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( HunyuanSamplingParams, ) from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams -from sglang.multimodal_gen.configs.sample.ltx_2 import LTX2SamplingParams +from sglang.multimodal_gen.configs.sample.ltx_2 import ( + LTX2SamplingParams, + LTX23SamplingParams, +) from sglang.multimodal_gen.configs.sample.mova import ( MOVA_360P_SamplingParams, MOVA_720P_SamplingParams, @@ -155,7 +158,18 @@ def _discover_and_register_pipelines(): package.__path__, package.__name__ + "." ): if not ispkg: - pipeline_module = importlib.import_module(module_name) + try: + pipeline_module = importlib.import_module(module_name) + except Exception as exc: + logger.warning( + "Skipping pipeline module %s during discovery due to import failure: %s", + module_name, + exc, + ) + logger.debug( + "Pipeline import failure details for %s", module_name, exc_info=True + ) + continue if hasattr(pipeline_module, "EntryClass"): entry_cls = pipeline_module.EntryClass entry_cls_list = ( @@ -594,12 +608,18 @@ def _register_configs(): register_configs( sampling_param_cls=LTX2SamplingParams, pipeline_config_cls=LTX2PipelineConfig, - hf_model_paths=[ - "Lightricks/LTX-2", - ], + hf_model_paths=["Lightricks/LTX-2"], model_detectors=[ lambda path: "ltx" in path.lower() and "video" in path.lower(), - lambda path: "ltx-2" in path.lower(), + lambda path: "ltx-2" in path.lower() and "ltx-2.3" not in path.lower(), + ], + ) + register_configs( + sampling_param_cls=LTX23SamplingParams, + pipeline_config_cls=LTX2PipelineConfig, + hf_model_paths=["Lightricks/LTX-2.3"], + model_detectors=[ + lambda path: "ltx-2.3" in path.lower(), ], ) diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py index 6e21f13b4..94190ec1f 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/diffusion_generator.py @@ -552,13 +552,15 @@ class DiffGenerator: self.shutdown() def __del__(self): - if self.owns_scheduler_client: + owns_scheduler_client = bool(getattr(self, "owns_scheduler_client", False)) + local_scheduler_process = getattr(self, "local_scheduler_process", None) + if owns_scheduler_client: logger.warning( "Generator was garbage collected without being shut down. " "Attempting to shut down the local server and client." ) self.shutdown() - elif self.local_scheduler_process: + elif local_scheduler_process: logger.warning( "Generator was garbage collected without being shut down. " "Attempting to shut down the local server." diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py index 5450cc409..06eafa2df 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/utils.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/utils.py @@ -288,12 +288,7 @@ def prepare_request( sampling_params=sampling_params, VSA_sparsity=server_args.attention_backend_config.VSA_sparsity, ) - try: - diffusers_kwargs = sampling_params.diffusers_kwargs - except AttributeError: - diffusers_kwargs = None - if diffusers_kwargs: - req.extra["diffusers_kwargs"] = diffusers_kwargs + sampling_params.apply_request_extra(req) req.adjust_size(server_args) diff --git a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py index 7d073d5a0..946f847a3 100644 --- a/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py +++ b/python/sglang/multimodal_gen/runtime/loader/component_loaders/adapter_loader.py @@ -1,5 +1,8 @@ from safetensors.torch import load_file as safetensors_load_file +from sglang.multimodal_gen.configs.models.adapter.ltx_2_connector import ( + LTX2ConnectorConfig, +) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( ComponentLoader, @@ -50,10 +53,9 @@ class AdapterLoader(ComponentLoader): target_device = get_local_torch_device() default_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.dit_precision] - from types import SimpleNamespace - with set_default_torch_dtype(default_dtype), skip_init_modules(): - connector_cfg = SimpleNamespace(**config) + connector_cfg = LTX2ConnectorConfig() + connector_cfg.update_model_arch(config) model = model_cls(connector_cfg).to( device=target_device, dtype=default_dtype ) diff --git a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py index 7bc5cd4db..e629e2aad 100644 --- a/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py +++ b/python/sglang/multimodal_gen/runtime/models/adapter/ltx_2_connector.py @@ -1,3 +1,4 @@ +import math from typing import Optional, Tuple, Union import torch @@ -89,6 +90,7 @@ class LTX2Attention(torch.nn.Module): norm_eps: float = 1e-6, norm_elementwise_affine: bool = True, rope_type: str = "interleaved", + apply_gated_attention: bool = False, processor=None, ): super().__init__() @@ -125,6 +127,9 @@ class LTX2Attention(torch.nn.Module): self.to_v = torch.nn.Linear( self.cross_attention_dim, self.inner_kv_dim, bias=bias ) + self.to_gate_logits = None + if apply_gated_attention: + self.to_gate_logits = torch.nn.Linear(query_dim, heads, bias=True) self.to_out = torch.nn.ModuleList([]) self.to_out.append(torch.nn.Linear(self.inner_dim, self.out_dim, bias=out_bias)) self.to_out.append(torch.nn.Dropout(dropout)) @@ -153,6 +158,7 @@ class LTX2Attention(torch.nn.Module): query_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, key_rotary_emb: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, ) -> torch.Tensor: + gate_input = hidden_states if encoder_hidden_states is None: encoder_hidden_states = hidden_states @@ -199,6 +205,15 @@ class LTX2Attention(torch.nn.Module): hidden_states = hidden_states.transpose(1, 2).flatten(2, 3) hidden_states = hidden_states.to(query.dtype) + if self.to_gate_logits is not None: + gate_logits = self.to_gate_logits(gate_input) + b, t, _ = hidden_states.shape + hidden_states = hidden_states.view(b, t, self.heads, self.head_dim) + hidden_states = hidden_states * ( + 2.0 * torch.sigmoid(gate_logits).unsqueeze(-1) + ) + hidden_states = hidden_states.view(b, t, self.heads * self.head_dim) + hidden_states = self.to_out[0](hidden_states) hidden_states = self.to_out[1](hidden_states) return hidden_states @@ -317,6 +332,7 @@ class LTX2TransformerBlock1d(nn.Module): activation_fn: str = "gelu-approximate", eps: float = 1e-6, rope_type: str = "interleaved", + apply_gated_attention: bool = False, ): super().__init__() @@ -327,6 +343,7 @@ class LTX2TransformerBlock1d(nn.Module): kv_heads=num_attention_heads, dim_head=attention_head_dim, rope_type=rope_type, + apply_gated_attention=apply_gated_attention, ) self.norm2 = torch.nn.RMSNorm(dim, eps=eps, elementwise_affine=False) @@ -373,6 +390,7 @@ class LTX2ConnectorTransformer1d(nn.Module): eps: float = 1e-6, causal_temporal_positioning: bool = False, rope_type: str = "interleaved", + apply_gated_attention: bool = False, ): super().__init__() self.num_attention_heads = num_attention_heads @@ -403,6 +421,7 @@ class LTX2ConnectorTransformer1d(nn.Module): num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, rope_type=rope_type, + apply_gated_attention=apply_gated_attention, ) for _ in range(num_layers) ] @@ -516,10 +535,37 @@ class LTX2TextConnectors(nn.Module): rope_double_precision = config.rope_double_precision causal_temporal_positioning = config.causal_temporal_positioning rope_type = config.rope_type - - self.text_proj_in = nn.Linear( - caption_channels * text_proj_in_factor, caption_channels, bias=False + connector_apply_gated_attention = config.connector_apply_gated_attention + feature_extractor_in_features = config.feature_extractor_in_features + video_feature_extractor_out_features = ( + config.video_feature_extractor_out_features ) + audio_feature_extractor_out_features = ( + config.audio_feature_extractor_out_features + ) + + self.text_proj_in: nn.Linear | None = None + self.video_aggregate_embed: nn.Linear | None = None + self.audio_aggregate_embed: nn.Linear | None = None + if ( + feature_extractor_in_features > 0 + and video_feature_extractor_out_features > 0 + and audio_feature_extractor_out_features > 0 + ): + self.video_aggregate_embed = nn.Linear( + feature_extractor_in_features, + video_feature_extractor_out_features, + bias=True, + ) + self.audio_aggregate_embed = nn.Linear( + feature_extractor_in_features, + audio_feature_extractor_out_features, + bias=True, + ) + else: + self.text_proj_in = nn.Linear( + caption_channels * text_proj_in_factor, caption_channels, bias=False + ) self.video_connector = LTX2ConnectorTransformer1d( num_attention_heads=video_connector_num_attention_heads, attention_head_dim=video_connector_attention_head_dim, @@ -530,6 +576,7 @@ class LTX2TextConnectors(nn.Module): rope_double_precision=rope_double_precision, causal_temporal_positioning=causal_temporal_positioning, rope_type=rope_type, + apply_gated_attention=connector_apply_gated_attention, ) self.audio_connector = LTX2ConnectorTransformer1d( num_attention_heads=audio_connector_num_attention_heads, @@ -541,8 +588,15 @@ class LTX2TextConnectors(nn.Module): rope_double_precision=rope_double_precision, causal_temporal_positioning=causal_temporal_positioning, rope_type=rope_type, + apply_gated_attention=connector_apply_gated_attention, ) + @staticmethod + def _rescale_v2_features( + x: torch.Tensor, target_dim: int, source_dim: int + ) -> torch.Tensor: + return x * math.sqrt(target_dim / source_dim) + def forward( self, text_encoder_hidden_states: torch.Tensor, @@ -557,12 +611,6 @@ class LTX2TextConnectors(nn.Module): ) attention_mask = attention_mask.to(text_dtype) * torch.finfo(text_dtype).max - # Ensure input dtype matches the layer's weight dtype - if text_encoder_hidden_states.dtype != self.text_proj_in.weight.dtype: - text_encoder_hidden_states = text_encoder_hidden_states.to( - self.text_proj_in.weight.dtype - ) - # Ensure sequence length is divisible by num_learnable_registers (128) seq_len = text_encoder_hidden_states.shape[1] num_learnable_registers = self.video_connector.num_learnable_registers @@ -579,10 +627,44 @@ class LTX2TextConnectors(nn.Module): # Pad with a large negative value to mask out the new tokens attention_mask = F.pad(attention_mask, (0, pad_len), value=-1000000.0) - text_encoder_hidden_states = self.text_proj_in(text_encoder_hidden_states) + if ( + self.video_aggregate_embed is not None + and self.audio_aggregate_embed is not None + ): + video_hidden_states = text_encoder_hidden_states + audio_hidden_states = text_encoder_hidden_states + if video_hidden_states.dtype != self.video_aggregate_embed.weight.dtype: + video_hidden_states = video_hidden_states.to( + self.video_aggregate_embed.weight.dtype + ) + if audio_hidden_states.dtype != self.audio_aggregate_embed.weight.dtype: + audio_hidden_states = audio_hidden_states.to( + self.audio_aggregate_embed.weight.dtype + ) + source_dim = self.video_aggregate_embed.out_features + video_hidden_states = self._rescale_v2_features( + video_hidden_states, + self.video_aggregate_embed.out_features, + source_dim, + ) + audio_hidden_states = self._rescale_v2_features( + audio_hidden_states, + self.audio_aggregate_embed.out_features, + source_dim, + ) + video_hidden_states = self.video_aggregate_embed(video_hidden_states) + audio_hidden_states = self.audio_aggregate_embed(audio_hidden_states) + else: + assert self.text_proj_in is not None + if text_encoder_hidden_states.dtype != self.text_proj_in.weight.dtype: + text_encoder_hidden_states = text_encoder_hidden_states.to( + self.text_proj_in.weight.dtype + ) + video_hidden_states = self.text_proj_in(text_encoder_hidden_states) + audio_hidden_states = video_hidden_states video_text_embedding, new_attn_mask = self.video_connector( - text_encoder_hidden_states, attention_mask + video_hidden_states, attention_mask ) attn_mask = (new_attn_mask < 1e-6).to(torch.int64) @@ -593,7 +675,7 @@ class LTX2TextConnectors(nn.Module): new_attn_mask = attn_mask.squeeze(-1) audio_text_embedding, _ = self.audio_connector( - text_encoder_hidden_states, attention_mask + audio_hidden_states, attention_mask ) return video_text_embedding, audio_text_embedding, new_attn_mask 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 eae637529..c2de1f643 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -37,6 +37,15 @@ from sglang.multimodal_gen.runtime.utils.logging_utils import init_logger logger = init_logger(__name__) +ADALN_NUM_BASE_PARAMS = 6 +ADALN_NUM_CROSS_ATTN_PARAMS = 3 + + +def adaln_embedding_coefficient(cross_attention_adaln: bool) -> int: + return ADALN_NUM_BASE_PARAMS + ( + ADALN_NUM_CROSS_ATTN_PARAMS if cross_attention_adaln else 0 + ) + def apply_interleaved_rotary_emb( x: torch.Tensor, freqs: Tuple[torch.Tensor, torch.Tensor] @@ -447,6 +456,7 @@ class LTX2Attention(nn.Module): norm_eps: float = 1e-6, qk_norm: bool = True, use_local_attention: bool = False, + apply_gated_attention: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, @@ -461,6 +471,8 @@ class LTX2Attention(nn.Module): self.norm_eps = float(norm_eps) self.qk_norm = bool(qk_norm) self.use_local_attention = bool(use_local_attention) + self.apply_gated_attention = bool(apply_gated_attention) + self.prefix = prefix tp_size = get_tp_world_size() if tp_size <= 0: @@ -499,6 +511,15 @@ class LTX2Attention(nn.Module): gather_output=False, quant_config=quant_config, ) + self.to_gate_logits: ColumnParallelLinear | None = None + if self.apply_gated_attention: + self.to_gate_logits = ColumnParallelLinear( + self.query_dim, + self.heads, + bias=True, + gather_output=False, + quant_config=quant_config, + ) self.q_norm: nn.Module | None = None self.k_norm: nn.Module | None = None @@ -561,6 +582,7 @@ class LTX2Attention(nn.Module): perturbation_mask: torch.Tensor | None = None, all_perturbed: bool = False, ) -> torch.Tensor: + gate_input = x context_ = x if context is None else context v, _ = self.to_v(context_) use_attention = not all_perturbed @@ -609,9 +631,17 @@ class LTX2Attention(nn.Module): if not use_attention: out = v - out = out.flatten(2) - out, _ = self.to_out[0](out) - return out + if self.to_gate_logits is not None: + gate_logits, _ = self.to_gate_logits(gate_input) + b, t = out.shape[:2] + out = out.view(b, t, self.local_heads, self.dim_head) + out = out * (2.0 * torch.sigmoid(gate_logits).unsqueeze(-1)) + out = out.view(b, t, self.local_heads * self.dim_head) + + out_flat = out.flatten(2) + out_proj, _ = self.to_out[0](out_flat) + + return out_proj def _slice_rope_for_tp( self, @@ -688,6 +718,10 @@ class LTX2TransformerBlock(nn.Module): audio_cross_attention_dim: int, qk_norm: bool = True, norm_eps: float = 1e-6, + apply_gated_attention: bool = False, + cross_attention_adaln: bool = False, + use_local_av_cross_attention: bool = False, + force_sdpa_v2a_cross_attention: bool = False, supported_attention_backends: set[AttentionBackendEnum] | None = None, prefix: str = "", quant_config: QuantizationConfig | None = None, @@ -695,6 +729,9 @@ class LTX2TransformerBlock(nn.Module): super().__init__() self.idx = idx self.norm_eps = norm_eps + # LTX2.3 + self.cross_attention_adaln = cross_attention_adaln + self.use_local_av_cross_attention = use_local_av_cross_attention # 1. Self-Attention (video and audio) self.attn1 = LTX2Attention( @@ -703,6 +740,7 @@ class LTX2TransformerBlock(nn.Module): dim_head=attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, + apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn1", quant_config=quant_config, @@ -713,6 +751,7 @@ class LTX2TransformerBlock(nn.Module): dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, + apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn1", quant_config=quant_config, @@ -729,6 +768,7 @@ class LTX2TransformerBlock(nn.Module): norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=True, + apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.attn2", quant_config=quant_config, @@ -741,6 +781,7 @@ class LTX2TransformerBlock(nn.Module): norm_eps=norm_eps, qk_norm=qk_norm, use_local_attention=True, + apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_attn2", quant_config=quant_config, @@ -754,6 +795,8 @@ class LTX2TransformerBlock(nn.Module): dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, + use_local_attention=use_local_av_cross_attention, + apply_gated_attention=apply_gated_attention, supported_attention_backends=supported_attention_backends, prefix=f"{prefix}.audio_to_video_attn", quant_config=quant_config, @@ -765,7 +808,13 @@ class LTX2TransformerBlock(nn.Module): dim_head=audio_attention_head_dim, norm_eps=norm_eps, qk_norm=qk_norm, - supported_attention_backends=supported_attention_backends, + use_local_attention=use_local_av_cross_attention, + apply_gated_attention=apply_gated_attention, + supported_attention_backends=( + {AttentionBackendEnum.TORCH_SDPA} + if force_sdpa_v2a_cross_attention + else supported_attention_backends + ), prefix=f"{prefix}.video_to_audio_attn", quant_config=quant_config, ) @@ -777,14 +826,23 @@ class LTX2TransformerBlock(nn.Module): ) # 5. Modulation Parameters - self.scale_shift_table = nn.Parameter(torch.randn(6, dim) / dim**0.5) + num_ada_params = adaln_embedding_coefficient(cross_attention_adaln) + self.scale_shift_table = nn.Parameter( + torch.randn(num_ada_params, dim) / dim**0.5 + ) self.audio_scale_shift_table = nn.Parameter( - torch.randn(6, audio_dim) / audio_dim**0.5 + torch.randn(num_ada_params, audio_dim) / audio_dim**0.5 ) self.video_a2v_cross_attn_scale_shift_table = nn.Parameter(torch.randn(5, dim)) self.audio_a2v_cross_attn_scale_shift_table = nn.Parameter( torch.randn(5, audio_dim) ) + if self.cross_attention_adaln: + # LTX2.3 + self.prompt_scale_shift_table = nn.Parameter(torch.randn(2, dim)) + self.audio_prompt_scale_shift_table = nn.Parameter( + torch.randn(2, audio_dim) + ) def get_ada_values( self, @@ -813,6 +871,8 @@ class LTX2TransformerBlock(nn.Module): audio_encoder_hidden_states: torch.Tensor, temb: torch.Tensor, temb_audio: torch.Tensor, + temb_prompt: torch.Tensor | None, + temb_audio_prompt: torch.Tensor | None, temb_ca_scale_shift: torch.Tensor, temb_ca_audio_scale_shift: torch.Tensor, temb_ca_gate: torch.Tensor, @@ -860,21 +920,70 @@ class LTX2TransformerBlock(nn.Module): ) audio_hidden_states = audio_hidden_states + attn_audio_hidden_states * agate_msa # 2. Prompt Cross-Attention - norm_hidden_states = rms_norm(hidden_states, self.norm_eps) - attn_hidden_states = self.attn2( - norm_hidden_states, - context=encoder_hidden_states, - mask=encoder_attention_mask, - ) - hidden_states = hidden_states + attn_hidden_states + if self.cross_attention_adaln: + # LTX2.3 + if temb_prompt is None or temb_audio_prompt is None: + raise ValueError( + "cross_attention_adaln requires prompt modulation tensors." + ) + vshift_q, vscale_q, vgate_q = self.get_ada_values( + self.scale_shift_table, batch_size, temb, slice(6, 9) + ) + v_prompt_shift, v_prompt_scale = self.get_ada_values( + self.prompt_scale_shift_table, batch_size, temb_prompt, slice(None) + ) + norm_hidden_states = ( + rms_norm(hidden_states, self.norm_eps) * (1 + vscale_q) + vshift_q + ) + mod_encoder_hidden_states = ( + encoder_hidden_states * (1 + v_prompt_scale) + v_prompt_shift + ) + attn_hidden_states = self.attn2( + norm_hidden_states, + context=mod_encoder_hidden_states, + mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states * vgate_q - norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) - attn_audio_hidden_states = self.audio_attn2( - norm_audio_hidden_states, - context=audio_encoder_hidden_states, - mask=audio_encoder_attention_mask, - ) - audio_hidden_states = audio_hidden_states + attn_audio_hidden_states + ashift_q, ascale_q, agate_q = self.get_ada_values( + self.audio_scale_shift_table, batch_size, temb_audio, slice(6, 9) + ) + a_prompt_shift, a_prompt_scale = self.get_ada_values( + self.audio_prompt_scale_shift_table, + batch_size, + temb_audio_prompt, + slice(None), + ) + norm_audio_hidden_states = ( + rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_q) + ashift_q + ) + mod_audio_encoder_hidden_states = ( + audio_encoder_hidden_states * (1 + a_prompt_scale) + a_prompt_shift + ) + attn_audio_hidden_states = self.audio_attn2( + norm_audio_hidden_states, + context=mod_audio_encoder_hidden_states, + mask=audio_encoder_attention_mask, + ) + audio_hidden_states = ( + audio_hidden_states + attn_audio_hidden_states * agate_q + ) + else: + norm_hidden_states = rms_norm(hidden_states, self.norm_eps) + attn_hidden_states = self.attn2( + norm_hidden_states, + context=encoder_hidden_states, + mask=encoder_attention_mask, + ) + hidden_states = hidden_states + attn_hidden_states + + norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) + attn_audio_hidden_states = self.audio_attn2( + norm_audio_hidden_states, + context=audio_encoder_hidden_states, + mask=audio_encoder_attention_mask, + ) + audio_hidden_states = audio_hidden_states + attn_audio_hidden_states # 3. Audio-to-Video and Video-to-Audio Cross-Attention norm_hidden_states = rms_norm(hidden_states, self.norm_eps) norm_audio_hidden_states = rms_norm(audio_hidden_states, self.norm_eps) @@ -976,7 +1085,7 @@ class LTX2TransformerBlock(nn.Module): ) # 4. Feedforward vshift_mlp, vscale_mlp, vgate_mlp = self.get_ada_values( - self.scale_shift_table, batch_size, temb, slice(3, None) + self.scale_shift_table, batch_size, temb, slice(3, 6) ) norm_hidden_states = ( rms_norm(hidden_states, self.norm_eps) * (1 + vscale_mlp) + vshift_mlp @@ -985,7 +1094,7 @@ class LTX2TransformerBlock(nn.Module): hidden_states = hidden_states + ff_output * vgate_mlp ashift_mlp, ascale_mlp, agate_mlp = self.get_ada_values( - self.audio_scale_shift_table, batch_size, temb_audio, slice(3, None) + self.audio_scale_shift_table, batch_size, temb_audio, slice(3, 6) ) norm_audio_hidden_states = ( rms_norm(audio_hidden_states, self.norm_eps) * (1 + ascale_mlp) + ashift_mlp @@ -1003,6 +1112,12 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): reverse_param_names_mapping = LTX2ArchConfig().reverse_param_names_mapping lora_param_names_mapping = LTX2ArchConfig().lora_param_names_mapping + @staticmethod + def _collapse_prompt_timestep(timestep: torch.Tensor) -> torch.Tensor: + if timestep.ndim <= 1: + return timestep + return timestep.amax(dim=tuple(range(1, timestep.ndim))) + def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None: """Validate TP-related dimension constraints (fail-fast).""" if tp_size < 1: @@ -1089,20 +1204,38 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): ) # 2. Prompt embeddings - self.caption_projection = LTX2TextProjection( - in_features=arch.caption_channels, hidden_size=self.hidden_size - ) - self.audio_caption_projection = LTX2TextProjection( - in_features=arch.caption_channels, hidden_size=self.audio_hidden_size - ) + self.caption_projection: LTX2TextProjection | None = None + self.audio_caption_projection: LTX2TextProjection | None = None + if not arch.caption_proj_before_connector: + self.caption_projection = LTX2TextProjection( + in_features=arch.caption_channels, hidden_size=self.hidden_size + ) + self.audio_caption_projection = LTX2TextProjection( + in_features=arch.caption_channels, hidden_size=self.audio_hidden_size + ) # 3. Timestep Modulation Params and Embedding self.adaln_single = LTX2AdaLayerNormSingle( - self.hidden_size, embedding_coefficient=6 + self.hidden_size, + embedding_coefficient=adaln_embedding_coefficient( + arch.cross_attention_adaln + ), ) self.audio_adaln_single = LTX2AdaLayerNormSingle( - self.audio_hidden_size, embedding_coefficient=6 + self.audio_hidden_size, + embedding_coefficient=adaln_embedding_coefficient( + arch.cross_attention_adaln + ), ) + self.prompt_adaln_single: LTX2AdaLayerNormSingle | None = None + self.audio_prompt_adaln_single: LTX2AdaLayerNormSingle | None = None + if arch.cross_attention_adaln: + self.prompt_adaln_single = LTX2AdaLayerNormSingle( + self.hidden_size, embedding_coefficient=2 + ) + self.audio_prompt_adaln_single = LTX2AdaLayerNormSingle( + self.audio_hidden_size, embedding_coefficient=2 + ) # Global Cross Attention Modulation Parameters self.av_ca_video_scale_shift_adaln_single = LTX2AdaLayerNormSingle( @@ -1141,6 +1274,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): rope_double_precision = bool( hf_config.get("rope_double_precision", arch.double_precision_rope) ) + self.quantize_video_rope_coords_to_hidden_dtype = bool( + hf_config.get("quantize_video_rope_coords_to_hidden_dtype", False) + ) causal_offset = int(hf_config.get("causal_offset", 1)) pos_embed_max_pos = int(arch.positional_embedding_max_pos[0]) @@ -1231,6 +1367,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_cross_attention_dim=arch.audio_cross_attention_dim, norm_eps=self.norm_eps, qk_norm=True, # Always True in LTX2 + apply_gated_attention=arch.apply_gated_attention, + cross_attention_adaln=arch.cross_attention_adaln, + use_local_av_cross_attention=bool( + getattr(arch, "use_local_av_cross_attention", False) + ), + force_sdpa_v2a_cross_attention=bool( + getattr(arch, "force_sdpa_v2a_cross_attention", False) + ), supported_attention_backends=self._supported_attention_backends, prefix=config.prefix, quant_config=quant_config, @@ -1336,6 +1480,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): device=audio_hidden_states.device, ) + if self.quantize_video_rope_coords_to_hidden_dtype: + video_coords = video_coords.to( + device=hidden_states.device, dtype=hidden_states.dtype + ) + else: + video_coords = video_coords.to(device=hidden_states.device) + audio_coords = audio_coords.to(device=audio_hidden_states.device) + video_rotary_emb = self.rope(video_coords, device=hidden_states.device) audio_rotary_emb = self.audio_rope( audio_coords, device=audio_hidden_states.device @@ -1367,12 +1519,25 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_embedded_timestep = audio_embedded_timestep.view( batch_size, -1, audio_embedded_timestep.size(-1) ) + temb_prompt = None + temb_audio_prompt = None + if self.prompt_adaln_single is not None: + prompt_timestep = self._collapse_prompt_timestep(timestep) + temb_prompt, _ = self.prompt_adaln_single( + prompt_timestep.flatten(), hidden_dtype=hidden_states.dtype + ) + temb_prompt = temb_prompt.view(batch_size, -1, temb_prompt.size(-1)) + if self.audio_prompt_adaln_single is not None: + audio_prompt_timestep = self._collapse_prompt_timestep(audio_timestep) + temb_audio_prompt, _ = self.audio_prompt_adaln_single( + audio_prompt_timestep.flatten(), + hidden_dtype=audio_hidden_states.dtype, + ) + temb_audio_prompt = temb_audio_prompt.view( + batch_size, -1, temb_audio_prompt.size(-1) + ) # 3.2. Prepare global modality cross attention modulation parameters - ts_ca_mult = ( - self.av_ca_timestep_scale_multiplier / self.timestep_scale_multiplier - ) - hidden_dtype = hidden_states.dtype temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single( timestep.flatten(), hidden_dtype=hidden_dtype @@ -1403,10 +1568,12 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): ) # 4. Prepare prompt embeddings - encoder_hidden_states = self.caption_projection(encoder_hidden_states) - audio_encoder_hidden_states = self.audio_caption_projection( - audio_encoder_hidden_states - ) + if self.caption_projection is not None: + encoder_hidden_states = self.caption_projection(encoder_hidden_states) + if self.audio_caption_projection is not None: + audio_encoder_hidden_states = self.audio_caption_projection( + audio_encoder_hidden_states + ) # 5. Run blocks skip_video_self_attn_blocks = set(skip_video_self_attn_blocks or ()) skip_audio_self_attn_blocks = set(skip_audio_self_attn_blocks or ()) @@ -1421,6 +1588,8 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): # under ForwardPattern.Pattern_0. temb=temb, temb_audio=temb_audio, + temb_prompt=temb_prompt, + temb_audio_prompt=temb_audio_prompt, temb_ca_scale_shift=temb_ca_scale_shift, temb_ca_audio_scale_shift=temb_ca_audio_scale_shift, temb_ca_gate=temb_ca_gate, diff --git a/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_3_condition_encoder.py b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_3_condition_encoder.py new file mode 100644 index 000000000..587d0e573 --- /dev/null +++ b/python/sglang/multimodal_gen/runtime/models/vaes/ltx_2_3_condition_encoder.py @@ -0,0 +1,204 @@ +from typing import Any + +import torch +import torch.nn as nn + +from sglang.multimodal_gen.runtime.models.vaes.ltx_2_vae import ( + LTX2VideoCausalConv3d, + LTX2VideoResnetBlock3d, + LTXVideoDownsampler3d, +) + + +def _patchify_video(sample: torch.Tensor, patch_size: int) -> torch.Tensor: + if patch_size == 1: + return sample + batch_size, channels, num_frames, height, width = sample.shape + sample = sample.reshape( + batch_size, + channels, + num_frames, + 1, + height // patch_size, + patch_size, + width // patch_size, + patch_size, + ) + return sample.permute(0, 1, 3, 7, 5, 2, 4, 6).flatten(1, 4) + + +class LTX23VideoPixelNorm(nn.Module): + def __init__(self, dim: int = 1, eps: float = 1e-8) -> None: + super().__init__() + self.dim = dim + self.eps = eps + + def forward(self, x: torch.Tensor) -> torch.Tensor: + mean_sq = torch.mean(x**2, dim=self.dim, keepdim=True) + rms = torch.sqrt(mean_sq + self.eps) + return x / rms + + +class LTX23PerChannelStatistics(nn.Module): + def __init__(self, latent_channels: int) -> None: + super().__init__() + self.register_buffer("std-of-means", torch.empty(latent_channels)) + self.register_buffer("mean-of-means", torch.empty(latent_channels)) + + def normalize(self, x: torch.Tensor) -> torch.Tensor: + mean = self.get_buffer("mean-of-means").view(1, -1, 1, 1, 1).to(x) + std = self.get_buffer("std-of-means").view(1, -1, 1, 1, 1).to(x) + return (x - mean) / std + + +class LTX23VideoResBlockStack(nn.Module): + def __init__( + self, channels: int, num_layers: int, spatial_padding_mode: str + ) -> None: + super().__init__() + self.res_blocks = nn.ModuleList( + [ + LTX2VideoResnetBlock3d( + in_channels=channels, + out_channels=channels, + spatial_padding_mode=spatial_padding_mode, + ) + for _ in range(num_layers) + ] + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + for res_block in self.res_blocks: + hidden_states = res_block(hidden_states, causal=True) + return hidden_states + + +def _make_ltx23_encoder_block( + block_name: str, + block_config: dict[str, Any], + in_channels: int, + spatial_padding_mode: str, +) -> tuple[nn.Module, int]: + if block_name == "res_x": + return ( + LTX23VideoResBlockStack( + channels=in_channels, + num_layers=int(block_config["num_layers"]), + spatial_padding_mode=spatial_padding_mode, + ), + in_channels, + ) + + multiplier = int(block_config.get("multiplier", 2)) + stride_map = { + "compress_space_res": (1, 2, 2), + "compress_time_res": (2, 1, 1), + "compress_all_res": (2, 2, 2), + } + stride = stride_map.get(block_name) + if stride is None: + raise ValueError(f"Unsupported LTX-2.3 encoder block: {block_name}") + out_channels = in_channels * multiplier + return ( + LTXVideoDownsampler3d( + in_channels=in_channels, + out_channels=out_channels, + stride=stride, + spatial_padding_mode=spatial_padding_mode, + ), + out_channels, + ) + + +class LTX23VideoConditionEncoder(nn.Module): + def __init__(self, config: dict[str, Any]) -> None: + super().__init__() + + vae_config = config.get("vae", config) + latent_channels = int(vae_config["latent_channels"]) + patch_size = int(vae_config.get("patch_size", 4)) + spatial_padding_mode = str(vae_config.get("spatial_padding_mode", "zeros")) + encoder_blocks = list(vae_config["encoder_blocks"]) + latent_log_var = str(vae_config.get("latent_log_var", "uniform")) + + self.patch_size = patch_size + self.latency_channels = latent_channels + self.latent_log_var = latent_log_var + self.per_channel_statistics = LTX23PerChannelStatistics(latent_channels) + + feature_channels = latent_channels + self.conv_in = LTX2VideoCausalConv3d( + in_channels=int(vae_config.get("in_channels", 3)) * patch_size**2, + out_channels=feature_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + self.down_blocks = nn.ModuleList() + for block_name, block_params in encoder_blocks: + block_config = ( + {"num_layers": block_params} + if isinstance(block_params, int) + else dict(block_params) + ) + block, feature_channels = _make_ltx23_encoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + spatial_padding_mode=spatial_padding_mode, + ) + self.down_blocks.append(block) + + self.conv_norm_out = LTX23VideoPixelNorm(dim=1, eps=1e-8) + self.conv_act = nn.SiLU() + + conv_out_channels = latent_channels + if latent_log_var == "per_channel": + conv_out_channels *= 2 + elif latent_log_var in {"uniform", "constant"}: + conv_out_channels += 1 + elif latent_log_var != "none": + raise ValueError(f"Unsupported latent_log_var: {latent_log_var}") + + self.conv_out = LTX2VideoCausalConv3d( + in_channels=feature_channels, + out_channels=conv_out_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + def forward(self, sample: torch.Tensor) -> torch.Tensor: + frames_count = int(sample.shape[2]) + if (frames_count - 1) % 8 != 0: + frames_to_crop = (frames_count - 1) % 8 + sample = sample[:, :, :-frames_to_crop, ...] + + hidden_states = _patchify_video(sample, self.patch_size) + hidden_states = self.conv_in(hidden_states, causal=True) + + for block in self.down_blocks: + hidden_states = block(hidden_states) + + hidden_states = self.conv_norm_out(hidden_states) + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states, causal=True) + + if self.latent_log_var == "uniform": + means = hidden_states[:, :-1, ...] + logvar = hidden_states[:, -1:, ...] + hidden_states = torch.cat( + [ + means, + logvar.repeat(1, means.shape[1], *([1] * (means.ndim - 2))), + ], + dim=1, + ) + elif self.latent_log_var == "constant": + means = hidden_states[:, :-1, ...] + logvar = torch.full_like(means, -30.0) + hidden_states = torch.cat([means, logvar], dim=1) + + means, _ = torch.chunk(hidden_states, 2, dim=1) + return self.per_channel_statistics.normalize(means) 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 53176ed98..9ddaea282 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 @@ -394,6 +394,18 @@ class LTXVideoUpsampler3d(nn.Module): return hidden_states +class LTX23PerChannelStatistics(nn.Module): + def __init__(self, latent_channels: int) -> None: + super().__init__() + self.register_buffer("mean_of_means", torch.empty(latent_channels)) + self.register_buffer("std_of_means", torch.empty(latent_channels)) + + def un_normalize(self, x: torch.Tensor) -> torch.Tensor: + mean = self.mean_of_means.view(1, -1, 1, 1, 1).to(x) + std = self.std_of_means.view(1, -1, 1, 1, 1).to(x) + return x * std + mean + + # Like LTX 1.0 LTXVideo095DownBlock3D, but with the updated LTX2VideoResnetBlock3d class LTX2VideoDownBlock3D(nn.Module): r""" @@ -609,6 +621,64 @@ class LTX2VideoMidBlock3d(nn.Module): return hidden_states +class LTX23VideoMidBlock3d(nn.Module): + def __init__( + self, + in_channels: int, + num_layers: int = 1, + dropout: float = 0.0, + resnet_eps: float = 1e-6, + resnet_act_fn: str = "swish", + inject_noise: bool = False, + timestep_conditioning: bool = False, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + self.time_embedder = None + if timestep_conditioning: + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + in_channels * 4, 0 + ) + + self.res_blocks = nn.ModuleList( + [ + LTX2VideoResnetBlock3d( + in_channels=in_channels, + out_channels=in_channels, + dropout=dropout, + eps=resnet_eps, + non_linearity=resnet_act_fn, + inject_noise=inject_noise, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + for _ in range(num_layers) + ] + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + causal: bool = True, + ) -> torch.Tensor: + if self.time_embedder is not None: + temb = self.time_embedder( + timestep=temb.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=hidden_states.size(0), + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(hidden_states.size(0), -1, 1, 1, 1) + + for res_block in self.res_blocks: + hidden_states = res_block(hidden_states, temb, causal=causal) + + return hidden_states + + # Like LTXVideoUpBlock3d but with no conv_in and the updated LTX2VideoResnetBlock3d class LTX2VideoUpBlock3d(nn.Module): r""" @@ -1104,6 +1174,192 @@ class LTX2VideoDecoder3d(nn.Module): return hidden_states +def _make_ltx23_decoder_block( + block_name: str, + block_config: dict, + in_channels: int, + resnet_norm_eps: float, + timestep_conditioning: bool, + spatial_padding_mode: str, +) -> tuple[nn.Module, int]: + out_channels = in_channels + if block_name == "res_x": + block = LTX23VideoMidBlock3d( + in_channels=in_channels, + num_layers=int(block_config["num_layers"]), + resnet_eps=resnet_norm_eps, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "res_x_y": + out_channels = in_channels // int(block_config.get("multiplier", 2)) + block = LTX2VideoResnetBlock3d( + in_channels=in_channels, + out_channels=out_channels, + eps=resnet_norm_eps, + timestep_conditioning=False, + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_time": + out_channels = in_channels // int(block_config.get("multiplier", 1)) + block = LTXVideoUpsampler3d( + in_channels=in_channels, + stride=(2, 1, 1), + residual=False, + upscale_factor=int(block_config.get("multiplier", 1)), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_space": + out_channels = in_channels // int(block_config.get("multiplier", 1)) + block = LTXVideoUpsampler3d( + in_channels=in_channels, + stride=(1, 2, 2), + residual=False, + upscale_factor=int(block_config.get("multiplier", 1)), + spatial_padding_mode=spatial_padding_mode, + ) + elif block_name == "compress_all": + out_channels = in_channels // int(block_config.get("multiplier", 1)) + block = LTXVideoUpsampler3d( + in_channels=in_channels, + stride=(2, 2, 2), + residual=bool(block_config.get("residual", False)), + upscale_factor=int(block_config.get("multiplier", 1)), + spatial_padding_mode=spatial_padding_mode, + ) + else: + raise ValueError(f"Unsupported LTX-2.3 decoder block: {block_name}") + + return block, out_channels + + +class LTX23VideoDecoder3d(nn.Module): + def __init__( + self, + in_channels: int = 128, + out_channels: int = 3, + decoder_blocks: tuple[tuple[str, dict], ...] = (), + patch_size: int = 4, + patch_size_t: int = 1, + resnet_norm_eps: float = 1e-6, + is_causal: bool = False, + timestep_conditioning: bool = False, + base_channels: int = 128, + spatial_padding_mode: str = "zeros", + ) -> None: + super().__init__() + + self.patch_size = patch_size + self.patch_size_t = patch_size_t + self.out_channels = out_channels * patch_size**2 + self.is_causal = is_causal + self.per_channel_statistics = LTX23PerChannelStatistics(in_channels) + + feature_channels = base_channels * 8 + self.conv_in = LTX2VideoCausalConv3d( + in_channels=in_channels, + out_channels=feature_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + self.up_blocks = nn.ModuleList([]) + for block_name, block_params in reversed(tuple(decoder_blocks)): + block_config = ( + {"num_layers": block_params} + if isinstance(block_params, int) + else dict(block_params) + ) + block, feature_channels = _make_ltx23_decoder_block( + block_name=block_name, + block_config=block_config, + in_channels=feature_channels, + resnet_norm_eps=resnet_norm_eps, + timestep_conditioning=timestep_conditioning, + spatial_padding_mode=spatial_padding_mode, + ) + self.up_blocks.append(block) + + self.norm_out = PerChannelRMSNorm() + self.conv_act = nn.SiLU() + self.conv_out = LTX2VideoCausalConv3d( + in_channels=feature_channels, + out_channels=self.out_channels, + kernel_size=3, + stride=1, + spatial_padding_mode=spatial_padding_mode, + ) + + self.time_embedder = None + self.scale_shift_table = None + self.timestep_scale_multiplier = None + if timestep_conditioning: + self.timestep_scale_multiplier = nn.Parameter( + torch.tensor(1000.0, dtype=torch.float32) + ) + self.time_embedder = PixArtAlphaCombinedTimestepSizeEmbeddings( + feature_channels * 2, 0 + ) + self.scale_shift_table = nn.Parameter( + torch.randn(2, feature_channels) / feature_channels**0.5 + ) + + def forward( + self, + hidden_states: torch.Tensor, + temb: Optional[torch.Tensor] = None, + causal: Optional[bool] = None, + ) -> torch.Tensor: + causal = self.is_causal if causal is None else causal + + hidden_states = self.per_channel_statistics.un_normalize(hidden_states) + hidden_states = self.conv_in(hidden_states, causal=causal) + + if self.timestep_scale_multiplier is not None and temb is not None: + temb = temb * self.timestep_scale_multiplier + + for up_block in self.up_blocks: + if isinstance(up_block, LTX23VideoMidBlock3d): + hidden_states = up_block(hidden_states, temb, causal=causal) + elif isinstance(up_block, LTX2VideoResnetBlock3d): + hidden_states = up_block(hidden_states, None, causal=causal) + else: + hidden_states = up_block(hidden_states, causal=causal) + + hidden_states = self.norm_out(hidden_states) + + if self.time_embedder is not None and temb is not None: + temb = self.time_embedder( + timestep=temb.flatten(), + resolution=None, + aspect_ratio=None, + batch_size=hidden_states.size(0), + hidden_dtype=hidden_states.dtype, + ) + temb = temb.view(hidden_states.size(0), -1, 1, 1, 1).unflatten(1, (2, -1)) + temb = temb + self.scale_shift_table[None, ..., None, None, None] + shift, scale = temb.unbind(dim=1) + hidden_states = hidden_states * (1 + scale) + shift + + hidden_states = self.conv_act(hidden_states) + hidden_states = self.conv_out(hidden_states, causal=causal) + + p = self.patch_size + p_t = self.patch_size_t + batch_size, _, num_frames, height, width = hidden_states.shape + hidden_states = hidden_states.reshape( + batch_size, -1, p_t, p, p, num_frames, height, width + ) + hidden_states = ( + hidden_states.permute(0, 1, 5, 2, 6, 4, 7, 3) + .flatten(6, 7) + .flatten(4, 5) + .flatten(2, 3) + ) + return hidden_states + + class AutoencoderKLLTX2Video(ParallelTiledVAE): r""" A VAE model with KL loss for encoding videos into latents and decoding latent representations into videos. @@ -1157,6 +1413,10 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE): timestep_conditioning = getattr( config.arch_config, "timestep_conditioning", False ) + use_ltx23_video_decoder = ( + str(getattr(config.arch_config, "video_decoder_variant", "ltx_2")) + == "ltx_2_3" + ) decoder_causal = config.arch_config.decoder_causal decoder_spatial_padding_mode = config.arch_config.decoder_spatial_padding_mode @@ -1175,22 +1435,53 @@ class AutoencoderKLLTX2Video(ParallelTiledVAE): encoder_spatial_padding_mode, ) - self.decoder = LTX2VideoDecoder3d( - in_channels=latent_channels, - out_channels=out_channels, - block_out_channels=decoder_block_out_channels, - spatio_temporal_scaling=decoder_spatio_temporal_scaling, - layers_per_block=decoder_layers_per_block, - patch_size=patch_size, - patch_size_t=patch_size_t, - resnet_norm_eps=resnet_norm_eps, - is_causal=decoder_causal, - inject_noise=decoder_inject_noise, - timestep_conditioning=timestep_conditioning, - upsample_residual=upsample_residual, - upsample_factor=upsample_factor, - spatial_padding_mode=decoder_spatial_padding_mode, - ) + if use_ltx23_video_decoder: + video_decoder_config = dict(config.arch_config.video_decoder_config) + if not video_decoder_config: + raise ValueError( + "LTX-2.3 native video decoder requires video_decoder_config." + ) + self.decoder = LTX23VideoDecoder3d( + in_channels=latent_channels, + out_channels=out_channels, + decoder_blocks=tuple(video_decoder_config["decoder_blocks"]), + patch_size=int(video_decoder_config.get("patch_size", patch_size)), + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=bool( + video_decoder_config.get("causal_decoder", decoder_causal) + ), + timestep_conditioning=bool( + video_decoder_config.get( + "timestep_conditioning", timestep_conditioning + ) + ), + base_channels=int( + video_decoder_config.get("decoder_base_channels", 128) + ), + spatial_padding_mode=str( + video_decoder_config.get( + "spatial_padding_mode", decoder_spatial_padding_mode + ) + ), + ) + else: + self.decoder = LTX2VideoDecoder3d( + in_channels=latent_channels, + out_channels=out_channels, + block_out_channels=decoder_block_out_channels, + spatio_temporal_scaling=decoder_spatio_temporal_scaling, + layers_per_block=decoder_layers_per_block, + patch_size=patch_size, + patch_size_t=patch_size_t, + resnet_norm_eps=resnet_norm_eps, + is_causal=decoder_causal, + inject_noise=decoder_inject_noise, + timestep_conditioning=timestep_conditioning, + upsample_residual=upsample_residual, + upsample_factor=upsample_factor, + spatial_padding_mode=decoder_spatial_padding_mode, + ) latents_mean = torch.zeros((latent_channels,), requires_grad=False) latents_std = torch.ones((latent_channels,), requires_grad=False) diff --git a/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py b/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py index 82ad20d2a..efd5e56ac 100644 --- a/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py +++ b/python/sglang/multimodal_gen/runtime/models/vocoder/ltx_2_vocoder.py @@ -1,13 +1,237 @@ import math from abc import ABC +from contextlib import nullcontext from typing import Tuple +import einops import torch import torch.nn as nn import torch.nn.functional as F from sglang.multimodal_gen.configs.models.vocoder.ltx_vocoder import LTXVocoderConfig +LRELU_SLOPE = 0.1 + + +def get_padding(kernel_size: int, dilation: int = 1) -> int: + return int((kernel_size * dilation - dilation) / 2) + + +def _sinc(x: torch.Tensor) -> torch.Tensor: + return torch.where( + x == 0, + torch.tensor(1.0, device=x.device, dtype=x.dtype), + torch.sin(math.pi * x) / math.pi / x, + ) + + +def kaiser_sinc_filter1d( + cutoff: float, half_width: float, kernel_size: int +) -> torch.Tensor: + even = kernel_size % 2 == 0 + half_size = kernel_size // 2 + delta_f = 4 * half_width + amplitude = 2.285 * (half_size - 1) * math.pi * delta_f + 7.95 + if amplitude > 50.0: + beta = 0.1102 * (amplitude - 8.7) + elif amplitude >= 21.0: + beta = 0.5842 * (amplitude - 21) ** 0.4 + 0.07886 * (amplitude - 21.0) + else: + beta = 0.0 + window = torch.kaiser_window(kernel_size, beta=beta, periodic=False) + time = ( + torch.arange(-half_size, half_size) + 0.5 + if even + else torch.arange(kernel_size) - half_size + ) + if cutoff == 0: + filter_ = torch.zeros_like(time) + else: + filter_ = 2 * cutoff * window * _sinc(2 * cutoff * time) + filter_ /= filter_.sum() + return filter_.view(1, 1, kernel_size) + + +class LowPassFilter1d(nn.Module): + def __init__( + self, + cutoff: float = 0.5, + half_width: float = 0.6, + stride: int = 1, + padding: bool = True, + padding_mode: str = "replicate", + kernel_size: int = 12, + ): + super().__init__() + self.kernel_size = kernel_size + self.even = kernel_size % 2 == 0 + self.pad_left = kernel_size // 2 - int(self.even) + self.pad_right = kernel_size // 2 + self.stride = stride + self.padding = padding + self.padding_mode = padding_mode + self.register_buffer( + "filter", kaiser_sinc_filter1d(cutoff, half_width, kernel_size) + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, channels, _ = x.shape + if self.padding: + x = F.pad(x, (self.pad_left, self.pad_right), mode=self.padding_mode) + return F.conv1d( + x, + self.filter.expand(channels, -1, -1), + stride=self.stride, + groups=channels, + ) + + +class UpSample1d(nn.Module): + def __init__( + self, + ratio: int = 2, + kernel_size: int | None = None, + persistent: bool = True, + window_type: str = "kaiser", + ): + super().__init__() + self.ratio = ratio + self.stride = ratio + + if window_type == "hann": + rolloff = 0.99 + lowpass_filter_width = 6 + width = math.ceil(lowpass_filter_width / rolloff) + self.kernel_size = 2 * width * ratio + 1 + self.pad = width + self.pad_left = 2 * width * ratio + self.pad_right = self.kernel_size - ratio + time_axis = (torch.arange(self.kernel_size) / ratio - width) * rolloff + time_clamped = time_axis.clamp(-lowpass_filter_width, lowpass_filter_width) + window = torch.cos(time_clamped * math.pi / lowpass_filter_width / 2) ** 2 + sinc_filter = (torch.sinc(time_axis) * window * rolloff / ratio).view( + 1, 1, -1 + ) + else: + self.kernel_size = ( + int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size + ) + self.pad = self.kernel_size // ratio - 1 + self.pad_left = ( + self.pad * self.stride + (self.kernel_size - self.stride) // 2 + ) + self.pad_right = ( + self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2 + ) + sinc_filter = kaiser_sinc_filter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + kernel_size=self.kernel_size, + ) + + self.register_buffer("filter", sinc_filter, persistent=persistent) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + _, channels, _ = x.shape + x = F.pad(x, (self.pad, self.pad), mode="replicate") + filt = self.filter.to(dtype=x.dtype, device=x.device).expand(channels, -1, -1) + x = self.ratio * F.conv_transpose1d( + x, filt, stride=self.stride, groups=channels + ) + return x[..., self.pad_left : -self.pad_right] + + +class DownSample1d(nn.Module): + def __init__(self, ratio: int = 2, kernel_size: int | None = None): + super().__init__() + self.lowpass = LowPassFilter1d( + cutoff=0.5 / ratio, + half_width=0.6 / ratio, + stride=ratio, + kernel_size=int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.lowpass(x) + + +class Activation1d(nn.Module): + def __init__( + self, + activation: nn.Module, + up_ratio: int = 2, + down_ratio: int = 2, + up_kernel_size: int = 12, + down_kernel_size: int = 12, + ): + super().__init__() + self.act = activation + self.upsample = UpSample1d(up_ratio, up_kernel_size) + self.downsample = DownSample1d(down_ratio, down_kernel_size) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.upsample(x) + x = self.act(x) + return self.downsample(x) + + +class Snake(nn.Module): + def __init__( + self, + in_features: int, + alpha: float = 1.0, + alpha_trainable: bool = True, + alpha_logscale: bool = True, + ): + super().__init__() + self.alpha_logscale = alpha_logscale + self.alpha = nn.Parameter( + torch.zeros(in_features) + if alpha_logscale + else torch.ones(in_features) * alpha + ) + self.alpha.requires_grad = alpha_trainable + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + return x + (1.0 / (alpha + self.eps)) * torch.sin(x * alpha).pow(2) + + +class SnakeBeta(nn.Module): + def __init__( + self, + in_features: int, + alpha: float = 1.0, + alpha_trainable: bool = True, + alpha_logscale: bool = True, + ): + super().__init__() + self.alpha_logscale = alpha_logscale + self.alpha = nn.Parameter( + torch.zeros(in_features) + if alpha_logscale + else torch.ones(in_features) * alpha + ) + self.alpha.requires_grad = alpha_trainable + self.beta = nn.Parameter( + torch.zeros(in_features) + if alpha_logscale + else torch.ones(in_features) * alpha + ) + self.beta.requires_grad = alpha_trainable + self.eps = 1e-9 + + def forward(self, x: torch.Tensor) -> torch.Tensor: + alpha = self.alpha.unsqueeze(0).unsqueeze(-1) + beta = self.beta.unsqueeze(0).unsqueeze(-1) + if self.alpha_logscale: + alpha = torch.exp(alpha) + beta = torch.exp(beta) + return x + (1.0 / (beta + self.eps)) * torch.sin(x * alpha).pow(2) + class ResBlock(nn.Module): def __init__( @@ -61,6 +285,252 @@ class ResBlock(nn.Module): return x +class AMPBlock1(nn.Module): + def __init__( + self, + channels: int, + kernel_size: int = 3, + dilation: tuple[int, int, int] = (1, 3, 5), + activation: str = "snake", + ): + super().__init__() + act_cls = SnakeBeta if activation == "snakebeta" else Snake + self.convs1 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[0], + padding=get_padding(kernel_size, dilation[0]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[1], + padding=get_padding(kernel_size, dilation[1]), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=dilation[2], + padding=get_padding(kernel_size, dilation[2]), + ), + ] + ) + self.convs2 = nn.ModuleList( + [ + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + nn.Conv1d( + channels, + channels, + kernel_size, + 1, + dilation=1, + padding=get_padding(kernel_size, 1), + ), + ] + ) + self.acts1 = nn.ModuleList( + [Activation1d(act_cls(channels)) for _ in range(len(self.convs1))] + ) + self.acts2 = nn.ModuleList( + [Activation1d(act_cls(channels)) for _ in range(len(self.convs2))] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + for conv1, conv2, act1, act2 in zip( + self.convs1, self.convs2, self.acts1, self.acts2 + ): + xt = act1(x) + xt = conv1(xt) + xt = act2(xt) + xt = conv2(xt) + x = x + xt + return x + + +class LTX23MelSTFT(nn.Module): + class STFTFn(nn.Module): + def __init__(self, filter_length: int, hop_length: int, win_length: int): + super().__init__() + self.hop_length = hop_length + self.win_length = win_length + n_freqs = filter_length // 2 + 1 + self.register_buffer( + "forward_basis", torch.zeros(n_freqs * 2, 1, filter_length) + ) + self.register_buffer( + "inverse_basis", torch.zeros(n_freqs * 2, 1, filter_length) + ) + + def forward(self, y: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + if y.dim() == 2: + y = y.unsqueeze(1) + left_pad = max(0, self.win_length - self.hop_length) + y = F.pad(y, (left_pad, 0)) + spec = F.conv1d(y, self.forward_basis, stride=self.hop_length, padding=0) + n_freqs = spec.shape[1] // 2 + real, imag = spec[:, :n_freqs], spec[:, n_freqs:] + magnitude = torch.sqrt(real**2 + imag**2) + phase = torch.atan2(imag.float(), real.float()).to(real.dtype) + return magnitude, phase + + def __init__( + self, filter_length: int, hop_length: int, win_length: int, n_mel_channels: int + ): + super().__init__() + self.stft_fn = self.STFTFn(filter_length, hop_length, win_length) + n_freqs = filter_length // 2 + 1 + self.register_buffer("mel_basis", torch.zeros(n_mel_channels, n_freqs)) + + def mel_spectrogram( + self, y: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + magnitude, phase = self.stft_fn(y) + energy = torch.norm(magnitude, dim=1) + mel = torch.matmul(self.mel_basis.to(magnitude.dtype), magnitude) + log_mel = torch.log(torch.clamp(mel, min=1e-5)) + return log_mel, magnitude, phase, energy + + +class LTX23VocoderCore(nn.Module): + def __init__( # noqa: PLR0913 + self, + resblock_kernel_sizes: list[int] | None = None, + upsample_rates: list[int] | None = None, + upsample_kernel_sizes: list[int] | None = None, + resblock_dilation_sizes: list[list[int]] | None = None, + upsample_initial_channel: int = 1024, + resblock: str = "1", + output_sampling_rate: int = 24000, + activation: str = "snake", + use_tanh_at_final: bool = True, + apply_final_activation: bool = True, + use_bias_at_final: bool = True, + ): + super().__init__() + if resblock_kernel_sizes is None: + resblock_kernel_sizes = [3, 7, 11] + if upsample_rates is None: + upsample_rates = [6, 5, 2, 2, 2] + if upsample_kernel_sizes is None: + upsample_kernel_sizes = [16, 15, 8, 4, 4] + if resblock_dilation_sizes is None: + resblock_dilation_sizes = [[1, 3, 5], [1, 3, 5], [1, 3, 5]] + + self.output_sampling_rate = output_sampling_rate + self.num_kernels = len(resblock_kernel_sizes) + self.num_upsamples = len(upsample_rates) + self.use_tanh_at_final = use_tanh_at_final + self.apply_final_activation = apply_final_activation + self.is_amp = resblock == "AMP1" + + self.conv_pre = nn.Conv1d( + in_channels=128, + out_channels=upsample_initial_channel, + kernel_size=7, + stride=1, + padding=3, + ) + self.ups = nn.ModuleList( + nn.ConvTranspose1d( + upsample_initial_channel // (2**i), + upsample_initial_channel // (2 ** (i + 1)), + kernel_size, + stride, + padding=(kernel_size - stride) // 2, + ) + for i, (stride, kernel_size) in enumerate( + zip(upsample_rates, upsample_kernel_sizes, strict=True) + ) + ) + + final_channels = upsample_initial_channel // (2 ** len(upsample_rates)) + self.resblocks = nn.ModuleList() + for i in range(len(upsample_rates)): + channels = upsample_initial_channel // (2 ** (i + 1)) + for kernel_size, dilations in zip( + resblock_kernel_sizes, resblock_dilation_sizes, strict=True + ): + if self.is_amp: + self.resblocks.append( + AMPBlock1( + channels, + kernel_size, + tuple(dilations), + activation=activation, + ) + ) + else: + self.resblocks.append( + ResBlock( + channels, + kernel_size=kernel_size, + dilations=tuple(dilations), + leaky_relu_negative_slope=LRELU_SLOPE, + padding_mode=get_padding(kernel_size, 1), + ) + ) + + self.act_post = ( + Activation1d(SnakeBeta(final_channels)) if self.is_amp else nn.LeakyReLU() + ) + self.conv_post = nn.Conv1d( + in_channels=final_channels, + out_channels=2, + kernel_size=7, + stride=1, + padding=3, + bias=use_bias_at_final, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.transpose(2, 3) + if x.dim() == 4: + assert x.shape[1] == 2, "Input must have 2 channels for stereo" + x = einops.rearrange(x, "b s c t -> b (s c) t") + + x = self.conv_pre(x) + for i in range(self.num_upsamples): + if not self.is_amp: + x = F.leaky_relu(x, LRELU_SLOPE) + x = self.ups[i](x) + start = i * self.num_kernels + end = start + self.num_kernels + block_outputs = torch.stack( + [self.resblocks[idx](x) for idx in range(start, end)], + dim=0, + ) + x = block_outputs.mean(dim=0) + + x = self.act_post(x) + x = self.conv_post(x) + if self.apply_final_activation: + x = torch.tanh(x) if self.use_tanh_at_final else torch.clamp(x, -1, 1) + return x + + class LTX2Vocoder(ABC, nn.Module): r""" LTX 2.0 vocoder for converting generated mel spectrograms back to audio waveforms. @@ -72,10 +542,61 @@ class LTX2Vocoder(ABC, nn.Module): ): super().__init__() self.config = config + nested_vocoder_cfg = getattr(config.arch_config, "vocoder", None) + if isinstance(nested_vocoder_cfg, dict) and "bwe" in nested_vocoder_cfg: + vocoder_cfg = nested_vocoder_cfg.get("vocoder", {}) + bwe_cfg = nested_vocoder_cfg["bwe"] + self.vocoder = LTX23VocoderCore( + resblock_kernel_sizes=vocoder_cfg.get("resblock_kernel_sizes"), + upsample_rates=vocoder_cfg.get("upsample_rates"), + upsample_kernel_sizes=vocoder_cfg.get("upsample_kernel_sizes"), + resblock_dilation_sizes=vocoder_cfg.get("resblock_dilation_sizes"), + upsample_initial_channel=vocoder_cfg.get( + "upsample_initial_channel", 1024 + ), + resblock=vocoder_cfg.get("resblock", "1"), + output_sampling_rate=bwe_cfg["input_sampling_rate"], + activation=vocoder_cfg.get("activation", "snake"), + use_tanh_at_final=vocoder_cfg.get("use_tanh_at_final", True), + apply_final_activation=vocoder_cfg.get("apply_final_activation", True), + use_bias_at_final=vocoder_cfg.get("use_bias_at_final", True), + ) + self.bwe_generator = LTX23VocoderCore( + resblock_kernel_sizes=bwe_cfg.get("resblock_kernel_sizes"), + upsample_rates=bwe_cfg.get("upsample_rates"), + upsample_kernel_sizes=bwe_cfg.get("upsample_kernel_sizes"), + resblock_dilation_sizes=bwe_cfg.get("resblock_dilation_sizes"), + upsample_initial_channel=bwe_cfg.get("upsample_initial_channel", 1024), + resblock=bwe_cfg.get("resblock", "1"), + output_sampling_rate=bwe_cfg["output_sampling_rate"], + activation=bwe_cfg.get("activation", "snake"), + use_tanh_at_final=bwe_cfg.get("use_tanh_at_final", True), + apply_final_activation=bwe_cfg.get("apply_final_activation", True), + use_bias_at_final=bwe_cfg.get("use_bias_at_final", True), + ) + self.mel_stft = LTX23MelSTFT( + filter_length=bwe_cfg["n_fft"], + hop_length=bwe_cfg["hop_length"], + win_length=bwe_cfg.get("win_size", bwe_cfg["n_fft"]), + n_mel_channels=bwe_cfg["num_mels"], + ) + self.input_sampling_rate = bwe_cfg["input_sampling_rate"] + self.output_sampling_rate = bwe_cfg["output_sampling_rate"] + self.hop_length = bwe_cfg["hop_length"] + with torch.device("cpu"): + self.resampler = UpSample1d( + ratio=self.output_sampling_rate // self.input_sampling_rate, + persistent=False, + window_type="hann", + ) + self.sample_rate = self.output_sampling_rate + return + self.sample_rate = ( getattr(config.arch_config, "sample_rate", None) or getattr(config.arch_config, "sampling_rate", None) or getattr(config.arch_config, "audio_sample_rate", None) + or getattr(config.arch_config, "output_sampling_rate", None) ) in_channels = config.arch_config.in_channels @@ -139,6 +660,12 @@ class LTX2Vocoder(ABC, nn.Module): self.conv_out = nn.Conv1d(output_channels, out_channels, 7, stride=1, padding=3) + def _compute_ltx23_mel(self, audio: torch.Tensor) -> torch.Tensor: + batch, channels, _ = audio.shape + flat = audio.reshape(batch * channels, -1) + mel, _, _, _ = self.mel_stft.mel_spectrogram(flat) + return mel.reshape(batch, channels, mel.shape[1], mel.shape[2]) + def forward( self, hidden_states: torch.Tensor, time_last: bool = False ) -> torch.Tensor: @@ -157,6 +684,32 @@ class LTX2Vocoder(ABC, nn.Module): `torch.Tensor`: Audio waveform tensor of shape (batch_size, out_channels, audio_length) """ + if hasattr(self, "bwe_generator"): + input_dtype = hidden_states.dtype + autocast_ctx = ( + torch.autocast( + device_type=hidden_states.device.type, dtype=torch.float32 + ) + if hidden_states.device.type != "cpu" + else nullcontext() + ) + with autocast_ctx: + waveform = self.vocoder(hidden_states.float()) + length_low_rate = waveform.shape[-1] + output_length = ( + length_low_rate + * self.output_sampling_rate + // self.input_sampling_rate + ) + remainder = length_low_rate % self.hop_length + if remainder != 0: + waveform = F.pad(waveform, (0, self.hop_length - remainder)) + mel = self._compute_ltx23_mel(waveform) + residual = self.bwe_generator(mel.transpose(2, 3)) + skip = self.resampler(waveform) + assert residual.shape == skip.shape + waveform = torch.clamp(residual + skip, -1, 1)[..., :output_length] + return waveform.to(input_dtype) # Ensure that the time/frame dimension is last if not time_last: diff --git a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py index 311a18c79..80173fe13 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -5,6 +5,9 @@ import numpy as np import torch from diffusers import FlowMatchEulerDiscreteScheduler +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, ) @@ -44,6 +47,8 @@ def _resolve_ltx2_two_stage_component_paths( if "spatial_upsampler" not in resolved: spatial_candidates = [ os.path.join(model_path, "latent_upsampler"), + os.path.join(model_path, "ltx-2.3-spatial-upscaler-x2-1.1.safetensors"), + os.path.join(model_path, "ltx-2.3-spatial-upscaler-x2-1.0.safetensors"), os.path.join(model_path, "ltx-2-spatial-upscaler-x2-1.0.safetensors"), ] for candidate in spatial_candidates: @@ -53,12 +58,15 @@ def _resolve_ltx2_two_stage_component_paths( break if "distilled_lora" not in resolved: - distilled_lora = os.path.join( - model_path, "ltx-2-19b-distilled-lora-384.safetensors" - ) - if os.path.exists(distilled_lora): - resolved["distilled_lora"] = distilled_lora - auto_resolved.append(f"distilled_lora={distilled_lora}") + distilled_lora_candidates = [ + os.path.join(model_path, "ltx-2.3-22b-distilled-lora-384.safetensors"), + os.path.join(model_path, "ltx-2-19b-distilled-lora-384.safetensors"), + ] + for distilled_lora in distilled_lora_candidates: + if os.path.exists(distilled_lora): + resolved["distilled_lora"] = distilled_lora + auto_resolved.append(f"distilled_lora={distilled_lora}") + break if auto_resolved: logger.info( @@ -81,6 +89,8 @@ def calculate_ltx2_shift( def prepare_ltx2_mu(batch: Req, server_args: ServerArgs): + if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config): + return "mu", None latent_num_frames = (int(batch.num_frames) - 1) // int( server_args.pipeline_config.vae_temporal_compression ) + 1 @@ -92,16 +102,49 @@ def prepare_ltx2_mu(batch: Req, server_args: ServerArgs): return "mu", calculate_ltx2_shift(video_sequence_length) +def build_official_ltx2_sigmas( + steps: int, + *, + max_shift: float = 2.05, + base_shift: float = 0.95, + stretch: bool = True, + terminal: float = 0.1, + default_number_of_tokens: int = MAX_SHIFT_ANCHOR, +) -> list[float]: + sigmas = torch.linspace(1.0, 0.0, steps + 1, dtype=torch.float32) + + mm = (max_shift - base_shift) / (MAX_SHIFT_ANCHOR - BASE_SHIFT_ANCHOR) + b = base_shift - mm * BASE_SHIFT_ANCHOR + sigma_shift = float(default_number_of_tokens) * mm + b + + non_zero_mask = sigmas != 0 + shifted = torch.where( + non_zero_mask, + math.exp(sigma_shift) / (math.exp(sigma_shift) + (1.0 / sigmas - 1.0)), + torch.zeros_like(sigmas), + ) + + if stretch: + one_minus_z = 1.0 - shifted[non_zero_mask] + scale_factor = one_minus_z[-1] / (1.0 - terminal) + shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor) + + return shifted[:-1].tolist() + + class LTX2SigmaPreparationStage(PipelineStage): """Prepare native LTX-2 sigma schedule before timestep setup.""" def forward(self, batch: Req, server_args: ServerArgs) -> Req: batch.extra["ltx2_phase"] = "stage1" - batch.sigmas = np.linspace( - 1.0, - 1.0 / int(batch.num_inference_steps), - int(batch.num_inference_steps), - ).tolist() + if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config): + batch.sigmas = build_official_ltx2_sigmas(int(batch.num_inference_steps)) + else: + batch.sigmas = np.linspace( + 1.0, + 1.0 / int(batch.num_inference_steps), + int(batch.num_inference_steps), + ).tolist() return batch diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py index fbc371b35..099543c76 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/__init__.py @@ -66,6 +66,7 @@ def build_pipeline( ) else: logger.info("No pipeline_class_name specified, using model_index.json") + model_info = get_model_info( model_path, backend=server_args.backend, diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py index bb30bd10b..f8494f3f5 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/schedule_batch.py @@ -135,7 +135,7 @@ class Req: trajectory_latents: torch.Tensor | None = None trajectory_audio_latents: torch.Tensor | None = None - # Extra parameters that might be needed by specific pipeline implementations + # Extra parameters that might be needed by specific pipeline implementations (e.g., LTX2.3 DenoisingAVStage) extra: dict[str, Any] = field(default_factory=dict) is_warmup: bool = False diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py index e56622189..362d81097 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/decoding_av.py @@ -25,6 +25,11 @@ class LTX2AVDecodingStage(DecodingStage): self.video_processor = VideoProcessor(vae_scale_factor=32) + @staticmethod + def _ltx2_should_externally_denorm_video_latents(server_args: ServerArgs) -> bool: + arch_config = server_args.pipeline_config.vae_config.arch_config + return str(getattr(arch_config, "video_decoder_variant", "ltx_2")) != "ltx_2_3" + def forward(self, batch: Req, server_args: ServerArgs) -> OutputBatch: self.load_model() @@ -40,9 +45,10 @@ class LTX2AVDecodingStage(DecodingStage): original_dtype = vae_dtype self.vae.to(torch.bfloat16) latents = latents.to(torch.bfloat16) - std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) - mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) - latents = latents * std + mean + if self._ltx2_should_externally_denorm_video_latents(server_args): + std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latents) + mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latents) + latents = latents * std + mean latents = server_args.pipeline_config.preprocess_decoding( latents, server_args, vae=self.vae ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py index 50118f6d0..e42c0ae9d 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/denoising_av.py @@ -1,5 +1,7 @@ import copy +import json import math +import os import time from io import BytesIO @@ -10,8 +12,15 @@ import torch from diffusers.models.autoencoders.vae import DiagonalGaussianDistribution from diffusers.models.modeling_outputs import AutoencoderKLOutput from diffusers.utils.torch_utils import randn_tensor +from safetensors.torch import load_file as safetensors_load_file +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) from sglang.multimodal_gen.runtime.managers.forward_context import set_forward_context +from sglang.multimodal_gen.runtime.models.vaes.ltx_2_3_condition_encoder import ( + LTX23VideoConditionEncoder, +) from sglang.multimodal_gen.runtime.models.vision_utils import ( load_image, normalize, @@ -46,6 +55,8 @@ class LTX2AVDenoisingStage(DenoisingStage): transformer=transformer, scheduler=scheduler, vae=vae, **kwargs ) self.audio_vae = audio_vae + self._condition_image_encoder = None + self._condition_image_encoder_dir = None @staticmethod def _get_video_latent_num_frames_for_model( @@ -116,13 +127,6 @@ class LTX2AVDenoisingStage(DenoisingStage): ) -> dict[str, object] | None: if stage != "stage1": return None - - pipeline_ref = getattr(self, "pipeline", None) - pipeline = pipeline_ref() if callable(pipeline_ref) else pipeline_ref - pipeline_name = getattr(pipeline, "pipeline_name", None) - if pipeline_name != "LTX2TwoStagePipeline": - return None - return batch.extra.get("ltx2_stage1_guider_params") @staticmethod @@ -141,6 +145,56 @@ class LTX2AVDenoisingStage(DenoisingStage): factor = rescale_scale * factor + (1.0 - rescale_scale) return pred * factor + @staticmethod + def _prepare_ltx2_ti2v_clean_state( + latents: torch.Tensor, + image_latent: torch.Tensor, + num_img_tokens: int, + zero_clean_latent: bool, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + latents = latents.clone() + conditioned = image_latent[:, :num_img_tokens, :].to( + device=latents.device, dtype=latents.dtype + ) + latents[:, :num_img_tokens, :] = conditioned + denoise_mask = torch.ones( + (latents.shape[0], latents.shape[1], 1), + device=latents.device, + dtype=torch.float32, + ) + denoise_mask[:, :num_img_tokens, :] = 0.0 + if zero_clean_latent: + clean_latent = torch.zeros_like(latents) + else: + clean_latent = latents.detach().clone() + clean_latent[:, :num_img_tokens, :] = conditioned + return latents, denoise_mask, clean_latent + + @staticmethod + def _ltx2_velocity_to_x0( + sample: torch.Tensor, + velocity: torch.Tensor, + sigma: float | torch.Tensor, + ) -> torch.Tensor: + if isinstance(sigma, torch.Tensor): + sigma = sigma.to(device=sample.device, dtype=torch.float32) + while sigma.ndim < sample.ndim: + sigma = sigma.unsqueeze(-1) + return (sample.float() - sigma * velocity.float()).to(sample.dtype) + return (sample.float() - float(sigma) * velocity.float()).to(sample.dtype) + + @staticmethod + def _repeat_batch_dim(tensor: torch.Tensor, target_batch_size: int) -> torch.Tensor: + """Repeat along batch dim while preserving any tokenwise timestep layout.""" + if tensor.shape[0] == int(target_batch_size): + return tensor + if tensor.shape[0] <= 0 or int(target_batch_size) % int(tensor.shape[0]) != 0: + raise ValueError( + f"Cannot repeat tensor with batch={tensor.shape[0]} to target_batch_size={target_batch_size}" + ) + repeat_factor = int(target_batch_size) // int(tensor.shape[0]) + return tensor.repeat(repeat_factor, *([1] * (tensor.ndim - 1))) + @classmethod def _ltx2_calculate_guided_x0( cls, @@ -252,6 +306,40 @@ class LTX2AVDenoisingStage(DenoisingStage): return True return int(getattr(batch, "sp_video_start_frame", 0)) == 0 + def _get_condition_image_encoder( + self, + server_args: ServerArgs, + *, + device: torch.device, + dtype: torch.dtype, + ) -> LTX23VideoConditionEncoder | None: + arch_config = server_args.pipeline_config.vae_config.arch_config + encoder_subdir = str(getattr(arch_config, "condition_encoder_subdir", "")) + if not encoder_subdir: + return None + + vae_model_path = server_args.model_paths["vae"] + encoder_dir = os.path.join(vae_model_path, encoder_subdir) + config_path = os.path.join(encoder_dir, "config.json") + weights_path = os.path.join(encoder_dir, "model.safetensors") + if not os.path.exists(config_path) or not os.path.exists(weights_path): + raise ValueError( + f"LTX-2 condition encoder files not found under {encoder_dir}" + ) + + cached_dir = self._condition_image_encoder_dir + encoder = self._condition_image_encoder + if encoder is None or cached_dir != encoder_dir: + with open(config_path, encoding="utf-8") as f: + config = json.load(f) + encoder = LTX23VideoConditionEncoder(config) + encoder.load_state_dict(safetensors_load_file(weights_path), strict=True) + self._condition_image_encoder = encoder + self._condition_image_encoder_dir = encoder_dir + + encoder = encoder.to(device=device, dtype=dtype) + return encoder + def _prepare_ltx2_image_latent(self, batch: Req, server_args: ServerArgs) -> None: """Encode `batch.image_path` into packed token latents for LTX-2 TI2V.""" if ( @@ -276,8 +364,11 @@ class LTX2AVDenoisingStage(DenoisingStage): ) img = load_image(image_path) + img_array = np.array(img).astype(np.uint8)[..., :3] + img_array = self._apply_video_codec_compression(img_array, crf=33) + conditioned_img = PIL.Image.fromarray(img_array) batch.condition_image = self._resize_center_crop( - img, width=int(batch.width), height=int(batch.height) + conditioned_img, width=int(batch.width), height=int(batch.height) ) latents_device = ( @@ -287,17 +378,22 @@ class LTX2AVDenoisingStage(DenoisingStage): ) encode_dtype = batch.latents.dtype original_dtype = PRECISION_TO_TYPE[server_args.pipeline_config.vae_precision] - self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) vae_autocast_enabled = ( original_dtype != torch.float32 ) and not server_args.disable_autocast + condition_image_encoder = self._get_condition_image_encoder( + server_args, device=latents_device, dtype=encode_dtype + ) + if condition_image_encoder is None: + self.vae = self.vae.to(device=latents_device, dtype=encode_dtype) video_condition = self._resize_center_crop_tensor( - img, + conditioned_img, width=int(batch.width), height=int(batch.height), device=latents_device, dtype=encode_dtype, + apply_codec_compression=False, ) with torch.autocast( @@ -306,31 +402,42 @@ class LTX2AVDenoisingStage(DenoisingStage): enabled=vae_autocast_enabled, ): try: - if server_args.pipeline_config.vae_tiling: + if ( + condition_image_encoder is None + and server_args.pipeline_config.vae_tiling + ): self.vae.enable_tiling() except Exception: pass if not vae_autocast_enabled: video_condition = video_condition.to(encode_dtype) - latent_dist: DiagonalGaussianDistribution = self.vae.encode(video_condition) - if isinstance(latent_dist, AutoencoderKLOutput): - latent_dist = latent_dist.latent_dist + if condition_image_encoder is not None: + latent = condition_image_encoder(video_condition) + else: + latent_dist: DiagonalGaussianDistribution = self.vae.encode( + video_condition + ) + if isinstance(latent_dist, AutoencoderKLOutput): + latent_dist = latent_dist.latent_dist - mode = server_args.pipeline_config.vae_config.encode_sample_mode() - if mode == "argmax": - latent = latent_dist.mode() - elif mode == "sample": - if batch.generator is None: - raise ValueError("Generator must be provided for VAE sampling.") - latent = latent_dist.sample(batch.generator) + if condition_image_encoder is None: + mode = server_args.pipeline_config.vae_config.encode_sample_mode() + if mode == "argmax": + latent = latent_dist.mode() + elif mode == "sample": + if batch.generator is None: + raise ValueError("Generator must be provided for VAE sampling.") + latent = latent_dist.sample(batch.generator) + else: + raise ValueError(f"Unsupported encode_sample_mode: {mode}") + + # Per-channel normalization: normalized = (x - mean) / std + mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) + std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) + latent = (latent - mean) / std else: - raise ValueError(f"Unsupported encode_sample_mode: {mode}") - - # Per-channel normalization: normalized = (x - mean) / std - mean = self.vae.latents_mean.view(1, -1, 1, 1, 1).to(latent) - std = self.vae.latents_std.view(1, -1, 1, 1, 1).to(latent) - latent = (latent - mean) / std + latent = latent.to(dtype=encode_dtype) packed = server_args.pipeline_config.maybe_pack_latents( latent, latent.shape[0], batch @@ -362,9 +469,12 @@ class LTX2AVDenoisingStage(DenoisingStage): batch.height, ) - self.vae.to(original_dtype) + if condition_image_encoder is None: + self.vae.to(original_dtype) if server_args.vae_cpu_offload: self.vae = self.vae.to("cpu") + if condition_image_encoder is not None: + self._condition_image_encoder = condition_image_encoder.to("cpu") @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: @@ -437,19 +547,15 @@ class LTX2AVDenoisingStage(DenoisingStage): if do_ti2v: if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") - latents[:, :num_img_tokens, :] = batch.image_latent[ - :, :num_img_tokens, : - ].to(device=latents.device, dtype=latents.dtype) - denoise_mask = torch.ones( - (latents.shape[0], latents.shape[1], 1), - device=latents.device, - dtype=torch.float32, + use_zero_clean_latent = is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ) + latents, denoise_mask, clean_latent = self._prepare_ltx2_ti2v_clean_state( + latents=latents, + image_latent=batch.image_latent, + num_img_tokens=num_img_tokens, + zero_clean_latent=use_zero_clean_latent, ) - denoise_mask[:, :num_img_tokens, :] = 0.0 - clean_latent = latents.detach().clone() - clean_latent[:, :num_img_tokens, :] = batch.image_latent[ - :, :num_img_tokens, : - ].to(device=latents.device, dtype=latents.dtype) with torch.autocast( device_type=current_platform.device_type, dtype=target_dtype, @@ -558,11 +664,12 @@ class LTX2AVDenoisingStage(DenoisingStage): ], dim=0, ) - timestep_video = timestep_video.expand( - int(latent_model_input.shape[0]) + cfg_batch_size = int(latent_model_input.shape[0]) + timestep_video = self._repeat_batch_dim( + timestep_video, cfg_batch_size ) - timestep_audio = timestep_audio.expand( - int(latent_model_input.shape[0]) + timestep_audio = self._repeat_batch_dim( + timestep_audio, cfg_batch_size ) with set_forward_context( @@ -616,11 +723,6 @@ class LTX2AVDenoisingStage(DenoisingStage): audio_latents = audio_scheduler.step( a_v_pos, t_device, audio_latents, return_dict=False )[0] - if do_ti2v: - latents[:, :num_img_tokens, :] = batch.image_latent[ - :, :num_img_tokens, : - ].to(device=latents.device, dtype=latents.dtype) - latents = self.post_forward_for_ti2v_task( batch, server_args, reserved_frames_mask, latents, z ) @@ -714,14 +816,20 @@ class LTX2AVDenoisingStage(DenoisingStage): if a_v_neg is not None: a_v_neg = a_v_neg.float() - # Velocity -> denoised (x0): x0 = x - sigma * v sigma_val = float(sigma.item()) - denoised_video = (latents.float() - sigma_val * v_pos).to( - latents.dtype + video_sigma_for_x0: float | torch.Tensor = sigma_val + if do_ti2v and denoise_mask is not None: + video_sigma_for_x0 = sigma.to( + device=latents.device, dtype=torch.float32 + ) * denoise_mask.squeeze(-1) + denoised_video = self._ltx2_velocity_to_x0( + latents, v_pos, video_sigma_for_x0 ) - denoised_audio = ( - audio_latents.float() - sigma_val * a_v_pos - ).to(audio_latents.dtype) + denoised_audio = self._ltx2_velocity_to_x0( + audio_latents, a_v_pos, sigma_val + ) + denoised_video_cond = denoised_video + denoised_audio_cond = denoised_audio denoised_video_neg = None denoised_audio_neg = None denoised_video_perturbed = None @@ -737,12 +845,12 @@ class LTX2AVDenoisingStage(DenoisingStage): and v_neg is not None and a_v_neg is not None ): - denoised_video_neg = ( - latents.float() - sigma_val * v_neg - ).to(latents.dtype) - denoised_audio_neg = ( - audio_latents.float() - sigma_val * a_v_neg - ).to(audio_latents.dtype) + denoised_video_neg = self._ltx2_velocity_to_x0( + latents, v_neg, video_sigma_for_x0 + ) + denoised_audio_neg = self._ltx2_velocity_to_x0( + audio_latents, a_v_neg, sigma_val + ) if stage1_guider_params is not None: video_skip = self._ltx2_should_skip_step( i, int(stage1_guider_params["video_skip_step"]) @@ -784,12 +892,12 @@ class LTX2AVDenoisingStage(DenoisingStage): stage1_guider_params["audio_stg_blocks"] ), ) - denoised_video_perturbed = ( - latents.float() - sigma_val * v_ptb.float() - ).to(latents.dtype) - denoised_audio_perturbed = ( - audio_latents.float() - sigma_val * a_v_ptb.float() - ).to(audio_latents.dtype) + denoised_video_perturbed = self._ltx2_velocity_to_x0( + latents, v_ptb.float(), video_sigma_for_x0 + ) + denoised_audio_perturbed = self._ltx2_velocity_to_x0( + audio_latents, a_v_ptb.float(), sigma_val + ) need_modality = ( float(stage1_guider_params["video_modality_scale"]) @@ -822,12 +930,12 @@ class LTX2AVDenoisingStage(DenoisingStage): disable_a2v_cross_attn=True, disable_v2a_cross_attn=True, ) - denoised_video_modality = ( - latents.float() - sigma_val * v_mod.float() - ).to(latents.dtype) - denoised_audio_modality = ( - audio_latents.float() - sigma_val * a_v_mod.float() - ).to(audio_latents.dtype) + denoised_video_modality = self._ltx2_velocity_to_x0( + latents, v_mod.float(), video_sigma_for_x0 + ) + denoised_audio_modality = self._ltx2_velocity_to_x0( + audio_latents, a_v_mod.float(), sigma_val + ) if not video_skip: denoised_video = self._ltx2_calculate_guided_x0( @@ -934,11 +1042,6 @@ class LTX2AVDenoisingStage(DenoisingStage): audio_latents.float() + v_audio.float() * dt ).to(dtype=audio_latents.dtype) - if do_ti2v: - latents[:, :num_img_tokens, :] = batch.image_latent[ - :, :num_img_tokens, : - ].to(device=latents.device, dtype=latents.dtype) - latents = self.post_forward_for_ti2v_task( batch, server_args, reserved_frames_mask, latents, z ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py index 130953346..6bd26554c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/latent_preparation_av.py @@ -1,6 +1,9 @@ import torch from diffusers.utils.torch_utils import randn_tensor +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req from sglang.multimodal_gen.runtime.pipelines_core.stages.latent_preparation import ( @@ -60,10 +63,112 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): ): return torch.float32 + @staticmethod + def _packed_video_latent_shape( + latent_shape: tuple[int, int, int, int, int], + pipeline_config, + ) -> tuple[int, int, int]: + batch_size, channels, num_frames, height, width = latent_shape + patch_size_t = int(pipeline_config.patch_size_t) + patch_size = int(pipeline_config.patch_size) + return ( + batch_size, + (num_frames // patch_size_t) + * (height // patch_size) + * (width // patch_size), + channels * patch_size_t * patch_size * patch_size, + ) + + @staticmethod + def _packed_audio_latent_shape( + latent_shape: tuple[int, int, int, int], + ) -> tuple[int, int, int]: + batch_size, channels, latent_length, mel_bins = latent_shape + return (batch_size, latent_length, channels * mel_bins) + def forward(self, batch: Req, server_args: ServerArgs) -> Req: - # 1. Prepare Video Latents using base class logic - # This sets batch.latents and batch.raw_latent_shape - batch = super().forward(batch, server_args) + if not is_ltx23_native_variant( + server_args.pipeline_config.vae_config.arch_config + ): + batch = super().forward(batch, server_args) + + try: + generate_audio = batch.generate_audio + except AttributeError: + generate_audio = True + if not generate_audio: + batch.audio_latents = None + batch.raw_audio_latent_shape = None + return batch + + device = get_local_torch_device() + dtype = self._get_latent_dtype(batch, server_args) + generator = batch.generator + + audio_latents = batch.audio_latents + batch_size = batch.batch_size + num_frames = batch.num_frames + + if audio_latents is None: + shape = server_args.pipeline_config.prepare_audio_latent_shape( + batch, batch_size, num_frames + ) + + audio_latents = randn_tensor( + shape, generator=generator, device=device, dtype=dtype + ) + else: + audio_latents = audio_latents.to(device) + + audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( + audio_latents, batch_size, batch + ) + + batch.audio_latents = audio_latents + batch.raw_audio_latent_shape = audio_latents.shape + return batch + + # 1. Prepare video latents directly in packed token space. + # Official LTX-2.3 pipelines sample noise after patchify; generating unpacked + # [B, C, F, H, W] noise and packing afterwards changes token ordering. + latent_num_frames = self.adjust_video_length(batch, server_args) + batch_size = batch.batch_size + dtype = self._get_latent_dtype(batch, server_args) + device = get_local_torch_device() + generator = batch.generator + + latents = batch.latents + num_frames = ( + latent_num_frames if latent_num_frames is not None else batch.num_frames + ) + + if latents is None: + latent_shape = server_args.pipeline_config.prepare_latent_shape( + batch, batch_size, num_frames + ) + latents = randn_tensor( + self._packed_video_latent_shape( + latent_shape, server_args.pipeline_config + ), + generator=generator, + device=device, + dtype=dtype, + ) + + latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents) + if latent_ids is not None: + batch.latent_ids = latent_ids.to(device=device) + else: + latents = latents.to(device) + latents = server_args.pipeline_config.maybe_pack_latents( + latents, batch_size, batch + ) + + if hasattr(self.scheduler, "init_noise_sigma"): + latents = latents * self.scheduler.init_noise_sigma + + batch.latents = latents + batch.raw_latent_shape = latents.shape # 2. Prepare Audio Latents (optional) # Default to True if not specified @@ -76,28 +181,24 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): batch.raw_audio_latent_shape = None return batch - device = get_local_torch_device() - dtype = self._get_latent_dtype(batch, server_args) - generator = batch.generator - audio_latents = batch.audio_latents - batch_size = batch.batch_size - num_frames = batch.num_frames if audio_latents is None: - shape = server_args.pipeline_config.prepare_audio_latent_shape( - batch, batch_size, num_frames + latent_shape = server_args.pipeline_config.prepare_audio_latent_shape( + batch, batch_size, batch.num_frames ) audio_latents = randn_tensor( - shape, generator=generator, device=device, dtype=dtype + self._packed_audio_latent_shape(latent_shape), + generator=generator, + device=device, + dtype=dtype, ) else: audio_latents = audio_latents.to(device) - - audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( - audio_latents, batch_size, batch - ) + audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( + audio_latents, batch_size, batch + ) # Store in batch batch.audio_latents = audio_latents diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index eabc2e9dd..53176b351 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -24,8 +24,13 @@ from sglang.utils import load_diffusion_overlay_registry_from_env logger = init_logger(__name__) # Built-in diffusion model overlay registry. -# Keep this empty until concrete overlay repos are ready to ship. -BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = {} +BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = { + "Lightricks/LTX-2.3": { + "overlay_repo_id": "MickJ/LTX-2.3-overlay", + "overlay_revision": "main", + "bundled_overlay_subdir": "ltx_2_3", + }, +} MODEL_OVERLAY_METADATA_PATTERNS = [ @@ -42,6 +47,43 @@ MODEL_OVERLAY_METADATA_PATTERNS = [ _MODEL_OVERLAY_REGISTRY_CACHE: dict[str, dict[str, Any]] | None = None +def _compute_overlay_fingerprint(overlay_dir: str) -> str: + hasher = hashlib.sha256() + for root, dir_names, file_names in os.walk(overlay_dir): + dir_names[:] = sorted( + d for d in dir_names if d != "__pycache__" and not d.endswith(".egg-info") + ) + for file_name in sorted(file_names): + if file_name.endswith((".safetensors", ".bin", ".pth", ".pt")): + continue + file_path = os.path.join(root, file_name) + rel_path = os.path.relpath(file_path, overlay_dir).replace(os.sep, "/") + hasher.update(rel_path.encode("utf-8")) + with open(file_path, "rb") as f: + hasher.update(hashlib.sha256(f.read()).digest()) + return hasher.hexdigest() + + +def _resolve_bundled_overlay_dir(overlay_spec: dict[str, Any]) -> str | None: + bundled_overlay_subdir = overlay_spec.get("bundled_overlay_subdir") + if not bundled_overlay_subdir: + return None + bundled_overlay_dir = os.path.abspath( + os.path.join( + os.path.dirname(__file__), + "..", + "..", + "model_overlays", + str(bundled_overlay_subdir), + ) + ) + if not os.path.isdir(bundled_overlay_dir): + return None + if load_overlay_manifest_if_present(bundled_overlay_dir) is None: + return None + return bundled_overlay_dir + + def get_diffusion_cache_root() -> str: return os.path.expanduser( os.getenv("SGLANG_DIFFUSION_CACHE_ROOT", "~/.cache/sgl_diffusion") @@ -300,6 +342,15 @@ def download_overlay_metadata( *, snapshot_download_fn: Callable[..., str], ) -> str: + bundled_overlay_dir = _resolve_bundled_overlay_dir(overlay_spec) + if bundled_overlay_dir is not None: + logger.info( + "Using bundled overlay metadata for %s from %s", + source_model_id, + bundled_overlay_dir, + ) + return bundled_overlay_dir + overlay_repo_id = str(overlay_spec["overlay_repo_id"]) if os.path.exists(overlay_repo_id): logger.info( @@ -419,6 +470,7 @@ def materialize_overlay_model( materializer_version = str(manifest.get("materializer_version", "v1")) overlay_repo_id = str(overlay_spec["overlay_repo_id"]) overlay_revision = str(overlay_spec.get("overlay_revision", "main")) + overlay_fingerprint = _compute_overlay_fingerprint(overlay_dir) cache_key = hashlib.sha256( json.dumps( { @@ -426,6 +478,7 @@ def materialize_overlay_model( "overlay_repo_id": overlay_repo_id, "overlay_revision": overlay_revision, "materializer_version": materializer_version, + "overlay_fingerprint": overlay_fingerprint, }, sort_keys=True, ).encode("utf-8") @@ -502,6 +555,7 @@ def materialize_overlay_model( "overlay_repo_id": overlay_repo_id, "overlay_revision": overlay_revision, "materializer_version": materializer_version, + "overlay_fingerprint": overlay_fingerprint, }, f, indent=2, diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 99926bd78..2ff322ddb 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2504,6 +2504,54 @@ "expected_e2e_ms": 8091.46, "expected_avg_denoise_ms": 141.37, "expected_median_denoise_ms": 142.63 + }, + "ltx_2.3_one_stage_ti2v": { + "stages_ms": { + "InputValidationStage": 3.27, + "TextEncodingStage": 1766.71, + "LTX2TextConnectorStage": 27.34, + "LTX2SigmaPreparationStage": 0.14, + "TimestepPreparationStage": 15.78, + "LTX2AVLatentPreparationStage": 0.25, + "LTX2AVDenoisingStage": 23757.65, + "LTX2AVDecodingStage": 1171.65, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 709.91, + "1": 704.93, + "2": 707.02, + "3": 702.54, + "4": 702.07, + "5": 746.58, + "6": 774.04, + "7": 766.37, + "8": 751.37, + "9": 734.67, + "10": 710.79, + "11": 689.84, + "12": 688.12, + "13": 689.74, + "14": 687.5, + "15": 699.65, + "16": 697.09, + "17": 686.87, + "18": 691.43, + "19": 712.52, + "20": 705.43, + "21": 727.21, + "22": 706.73, + "23": 705.63, + "24": 714.95, + "25": 707.69, + "26": 748.35, + "27": 738.02, + "28": 738.05, + "29": 726.96 + }, + "expected_e2e_ms": 26916.58, + "expected_avg_denoise_ms": 715.73, + "expected_median_denoise_ms": 707.35 } } } diff --git a/python/sglang/multimodal_gen/test/server/testcase_configs.py b/python/sglang/multimodal_gen/test/server/testcase_configs.py index 151653ee5..07e2e0f21 100644 --- a/python/sglang/multimodal_gen/test/server/testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/testcase_configs.py @@ -935,7 +935,6 @@ TWO_GPU_CASES_A = [ model_path="Lightricks/LTX-2", modality="video", num_gpus=2, - dit_layerwise_offload=True, extras=["--pipeline-class-name LTX2TwoStagePipeline"], ), T2V_sampling_params, @@ -1040,6 +1039,15 @@ TWO_GPU_CASES_B = [ ), TI2I_sampling_params, ), + DiffusionTestCase( + "ltx_2.3_one_stage_ti2v", + DiffusionServerArgs( + model_path="Lightricks/LTX-2.3", + modality="video", + num_gpus=2, + ), + TI2V_sampling_params, + ), ] if not current_platform.is_hip(): diff --git a/python/sglang/multimodal_gen/test/unit/test_model_overlay_ltx23.py b/python/sglang/multimodal_gen/test/unit/test_model_overlay_ltx23.py new file mode 100644 index 000000000..62bd1f4fb --- /dev/null +++ b/python/sglang/multimodal_gen/test/unit/test_model_overlay_ltx23.py @@ -0,0 +1,392 @@ +import json +import os +import tempfile +from types import SimpleNamespace + +import pytest +import torch +from safetensors import safe_open +from safetensors.torch import save_file + +pytest.importorskip("triton.compiler") + +from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + is_ltx23_native_variant, +) +from sglang.multimodal_gen.configs.sample.sampling_params import SamplingParams +from sglang.multimodal_gen.model_overlays.ltx_2_3._overlay.materialize import ( + _build_transformer_config, + _build_vae_config, + _rename_connector_key, + _repack_ltx23_image_encoder_weights, + _repack_ltx23_video_decoder_weights, +) +from sglang.multimodal_gen.registry import get_model_info +from sglang.multimodal_gen.runtime.pipelines.ltx_2_pipeline import ( + _resolve_ltx2_two_stage_component_paths, + build_official_ltx2_sigmas, + prepare_ltx2_mu, +) +from sglang.multimodal_gen.runtime.pipelines_core.schedule_batch import Req +from sglang.multimodal_gen.runtime.pipelines_core.stages.decoding_av import ( + LTX2AVDecodingStage, +) +from sglang.multimodal_gen.runtime.pipelines_core.stages.denoising_av import ( + LTX2AVDenoisingStage, +) +from sglang.multimodal_gen.runtime.utils.model_overlay import ( + resolve_model_overlay_target, +) + + +def _make_req(**sampling_kwargs) -> Req: + return Req( + sampling_params=SamplingParams(**sampling_kwargs), + prompt="prompt", + prompt_embeds=[torch.zeros(1, 1, 1)], + ) + + +def test_ltx23_builtin_overlay_target_is_hf_repo(): + target = resolve_model_overlay_target("Lightricks/LTX-2.3") + assert target is not None + + source_model_id, overlay_spec = target + assert source_model_id == "Lightricks/LTX-2.3" + assert str(overlay_spec["overlay_repo_id"]) == "MickJ/LTX-2.3-overlay" + assert str(overlay_spec["overlay_revision"]) == "main" + assert str(overlay_spec["bundled_overlay_subdir"]) == "ltx_2_3" + + +def test_ltx23_model_info_resolves_to_native_pipeline_and_sampling_params(): + model_info = get_model_info("Lightricks/LTX-2.3", backend="sglang") + + assert model_info is not None + assert model_info.pipeline_cls.__name__ == "LTX2Pipeline" + assert model_info.sampling_param_cls.__name__ == "LTX23SamplingParams" + + +def test_ltx23_sampling_defaults_use_cuda_generator(): + sampling_params = SamplingParams.from_pretrained( + "Lightricks/LTX-2.3", + backend="sglang", + ) + + assert sampling_params.generator_device == "cuda" + assert sampling_params.guidance_scale == 3.0 + assert sampling_params.num_inference_steps == 30 + + +def test_ltx2_sampling_defaults_keep_cpu_generator(): + sampling_params = SamplingParams.from_pretrained( + "Lightricks/LTX-2", + backend="sglang", + ) + + assert sampling_params.generator_device == "cpu" + + +def test_ltx23_build_request_extra_sets_stage1_guider_defaults(): + sampling_params = SamplingParams.from_pretrained( + "Lightricks/LTX-2.3", + backend="sglang", + ) + + assert sampling_params.build_request_extra()["ltx2_stage1_guider_params"] == { + "video_cfg_scale": 3.0, + "video_stg_scale": 1.0, + "video_rescale_scale": 0.7, + "video_modality_scale": 3.0, + "video_skip_step": 0, + "video_stg_blocks": [28], + "audio_cfg_scale": 7.0, + "audio_stg_scale": 1.0, + "audio_rescale_scale": 0.7, + "audio_modality_scale": 3.0, + "audio_skip_step": 0, + "audio_stg_blocks": [28], + } + + +def test_sampling_params_apply_request_extra_populates_req_extra(): + sampling_params = SamplingParams.from_pretrained( + "Lightricks/LTX-2.3", + backend="sglang", + ) + req = Req(sampling_params=sampling_params, prompt="prompt") + + sampling_params.apply_request_extra(req) + + assert req.extra["ltx2_stage1_guider_params"]["video_cfg_scale"] == 3.0 + assert req.extra["ltx2_stage1_guider_params"]["audio_cfg_scale"] == 7.0 + + +def test_ltx23_uses_official_sigma_schedule(): + sigmas = build_official_ltx2_sigmas(30) + + assert len(sigmas) == 30 + assert sigmas[0] == pytest.approx(1.0) + assert sigmas[1] == pytest.approx(0.99495703, abs=1e-6) + assert sigmas[-1] == pytest.approx(0.1, abs=1e-6) + + +def test_ltx23_native_variant_uses_explicit_marker_only(): + assert is_ltx23_native_variant(SimpleNamespace(ltx_variant="ltx_2_3")) is True + assert is_ltx23_native_variant(SimpleNamespace(ltx_variant="ltx_2")) is False + + +def test_prepare_ltx2_mu_respects_variant_marker(): + ltx23_server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(ltx_variant="ltx_2_3") + ) + ) + ) + legacy_server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(ltx_variant="ltx_2") + ), + vae_temporal_compression=8, + vae_scale_factor=32, + ) + ) + + assert prepare_ltx2_mu( + _make_req(num_frames=121, height=512, width=768), + ltx23_server_args, + ) == ("mu", None) + + key, mu = prepare_ltx2_mu( + _make_req(num_frames=121, height=512, width=768), + legacy_server_args, + ) + assert key == "mu" + assert isinstance(mu, float) + assert mu > 0.0 + + +def test_ltx23_ti2v_clean_latent_uses_zero_background(): + latents = torch.arange(24, dtype=torch.float32).view(1, 6, 4) + image_latent = torch.full((1, 2, 4), 99.0) + + conditioned, denoise_mask, clean_latent = ( + LTX2AVDenoisingStage._prepare_ltx2_ti2v_clean_state( + latents=latents, + image_latent=image_latent, + num_img_tokens=2, + zero_clean_latent=True, + ) + ) + + assert torch.equal(conditioned[:, :2], image_latent) + assert torch.equal(clean_latent[:, :2], image_latent) + assert torch.equal(clean_latent[:, 2:], torch.zeros_like(clean_latent[:, 2:])) + assert torch.equal(denoise_mask[:, :2], torch.zeros_like(denoise_mask[:, :2])) + assert torch.equal(denoise_mask[:, 2:], torch.ones_like(denoise_mask[:, 2:])) + + +def test_ltx2_ti2v_clean_latent_keeps_legacy_background_when_requested(): + latents = torch.arange(24, dtype=torch.float32).view(1, 6, 4) + image_latent = torch.full((1, 2, 4), 99.0) + + conditioned, _, clean_latent = LTX2AVDenoisingStage._prepare_ltx2_ti2v_clean_state( + latents=latents, + image_latent=image_latent, + num_img_tokens=2, + zero_clean_latent=False, + ) + + assert torch.equal(conditioned[:, :2], image_latent) + assert torch.equal(clean_latent[:, :2], image_latent) + assert torch.equal(clean_latent[:, 2:], latents[:, 2:]) + + +def test_ltx23_velocity_to_x0_supports_tokenwise_sigma(): + sample = torch.tensor([[[1.0, 2.0], [3.0, 4.0]]], dtype=torch.float32) + velocity = torch.tensor([[[0.5, 0.5], [1.0, 1.0]]], dtype=torch.float32) + sigma = torch.tensor([[0.0, 0.5]], dtype=torch.float32) + + denoised = LTX2AVDenoisingStage._ltx2_velocity_to_x0(sample, velocity, sigma) + + expected = torch.tensor([[[1.0, 2.0], [2.5, 3.5]]], dtype=torch.float32) + assert torch.allclose(denoised, expected) + + +def test_ltx23_connector_repack_renames_qk_norm_keys(): + assert ( + _rename_connector_key( + "model.diffusion_model.video_embeddings_connector.transformer_1d_blocks.0.attn1.q_norm.weight" + ) + == "video_connector.transformer_blocks.0.attn1.norm_q.weight" + ) + assert ( + _rename_connector_key( + "model.diffusion_model.audio_embeddings_connector.transformer_1d_blocks.1.attn1.k_norm.weight" + ) + == "audio_connector.transformer_blocks.1.attn1.norm_k.weight" + ) + + +def test_ltx23_transformer_config_forces_sdpa_for_v2a_cross_attention(): + with tempfile.TemporaryDirectory() as tmpdir: + donor_dir = os.path.join(tmpdir, "donor") + os.makedirs(os.path.join(donor_dir, "transformer"), exist_ok=True) + with open(os.path.join(donor_dir, "transformer", "config.json"), "w") as f: + json.dump({"_class_name": "OldClass", "num_layers": 1}, f) + + config = _build_transformer_config(donor_dir) + + assert config["_class_name"] == "LTX2VideoTransformer3DModel" + assert config["force_sdpa_v2a_cross_attention"] is True + + +def test_ltx23_vae_config_adds_required_markers(): + with tempfile.TemporaryDirectory() as tmpdir: + auxiliary_dir = os.path.join(tmpdir, "aux") + config_donor_dir = os.path.join(tmpdir, "donor") + os.makedirs(os.path.join(auxiliary_dir, "vae"), exist_ok=True) + os.makedirs(os.path.join(config_donor_dir, "vae"), exist_ok=True) + + with open(os.path.join(auxiliary_dir, "vae", "config.json"), "w") as f: + json.dump( + { + "_class_name": "AutoencoderKLLTX2Video", + "scaling_factor": 1.0, + "patch_size": 4, + "decoder_causal": False, + "timestep_conditioning": False, + "encoder_spatial_padding_mode": "zeros", + "decoder_spatial_padding_mode": "reflect", + }, + f, + ) + with open(os.path.join(config_donor_dir, "vae", "config.json"), "w") as f: + json.dump( + { + "vae": { + "decoder_blocks": [["res_x", {"num_layers": 2}]], + "decoder_base_channels": 128, + "patch_size": 4, + "spatial_padding_mode": "zeros", + } + }, + f, + ) + + config = _build_vae_config(auxiliary_dir, config_donor_dir) + + assert config["ltx_variant"] == "ltx_2_3" + assert config["condition_encoder_subdir"] == "ltx23_image_encoder" + assert config["video_decoder_variant"] == "ltx_2_3" + assert config["video_decoder_config"]["decoder_base_channels"] == 128 + + +def test_ltx23_repack_image_encoder_keeps_only_encoder_tensors(): + with tempfile.TemporaryDirectory() as tmpdir: + source_path = os.path.join(tmpdir, "source.safetensors") + output_path = os.path.join(tmpdir, "output.safetensors") + save_file( + { + "encoder.conv_in.conv.weight": torch.ones(1), + "decoder.conv_in.conv.weight": torch.full((1,), 2.0), + "per_channel_statistics.mean-of-means": torch.full((2,), 3.0), + }, + source_path, + ) + + _repack_ltx23_image_encoder_weights(source_path, output_path) + + with safe_open(output_path, framework="pt") as f: + assert sorted(f.keys()) == [ + "conv_in.conv.weight", + "per_channel_statistics.mean-of-means", + ] + + +def test_ltx23_repack_video_decoder_keeps_decoder_and_stats(): + with tempfile.TemporaryDirectory() as tmpdir: + auxiliary_path = os.path.join(tmpdir, "aux.safetensors") + donor_path = os.path.join(tmpdir, "donor.safetensors") + output_path = os.path.join(tmpdir, "output.safetensors") + save_file( + { + "encoder.conv_in.conv.weight": torch.full((1,), 5.0), + }, + auxiliary_path, + ) + save_file( + { + "decoder.conv_in.conv.weight": torch.ones(1), + "per_channel_statistics.mean-of-means": torch.full((2,), 3.0), + "per_channel_statistics.std-of-means": torch.full((2,), 4.0), + }, + donor_path, + ) + + _repack_ltx23_video_decoder_weights(auxiliary_path, donor_path, output_path) + + with safe_open(output_path, framework="pt") as f: + assert sorted(f.keys()) == [ + "decoder.conv_in.conv.weight", + "decoder.per_channel_statistics.mean_of_means", + "decoder.per_channel_statistics.std_of_means", + "encoder.conv_in.conv.weight", + "latents_mean", + "latents_std", + ] + + +def test_ltx23_decode_skips_external_denorm(): + ltx23_server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(video_decoder_variant="ltx_2_3") + ) + ) + ) + legacy_server_args = SimpleNamespace( + pipeline_config=SimpleNamespace( + vae_config=SimpleNamespace( + arch_config=SimpleNamespace(video_decoder_variant="ltx_2") + ) + ) + ) + + assert ( + LTX2AVDecodingStage._ltx2_should_externally_denorm_video_latents( + ltx23_server_args + ) + is False + ) + assert ( + LTX2AVDecodingStage._ltx2_should_externally_denorm_video_latents( + legacy_server_args + ) + is True + ) + + +def test_ltx2_two_stage_component_auto_resolution_preserves_legacy_candidates(tmp_path): + legacy_spatial = tmp_path / "ltx-2-spatial-upscaler-x2-1.0.safetensors" + legacy_lora = tmp_path / "ltx-2-19b-distilled-lora-384.safetensors" + legacy_spatial.touch() + legacy_lora.touch() + + resolved = _resolve_ltx2_two_stage_component_paths(str(tmp_path), {}) + + assert resolved["spatial_upsampler"] == str(legacy_spatial) + assert resolved["distilled_lora"] == str(legacy_lora) + + +def test_ltx23_two_stage_component_auto_resolution_prefers_23_assets(tmp_path): + spatial = tmp_path / "ltx-2.3-spatial-upscaler-x2-1.1.safetensors" + lora = tmp_path / "ltx-2.3-22b-distilled-lora-384.safetensors" + spatial.touch() + lora.touch() + + resolved = _resolve_ltx2_two_stage_component_paths(str(tmp_path), {}) + + assert resolved["spatial_upsampler"] == str(spatial) + assert resolved["distilled_lora"] == str(lora)