diff --git a/docs/diffusion/compatibility_matrix.md b/docs/diffusion/compatibility_matrix.md index 5520f35f2..37b95acfa 100644 --- a/docs/diffusion/compatibility_matrix.md +++ b/docs/diffusion/compatibility_matrix.md @@ -33,23 +33,29 @@ 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 | ✅ | ❌ | ❌ | ❌ | ✅ | ✅ | ⭕ | -| Wan2.1 Fun 1.3B InP | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | -| Helios Base | `BestWishYsh/Helios-Base` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| Helios Mid | `BestWishYsh/Helios-Mid` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| Helios Distilled | `BestWishYsh/Helios-Distilled` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| LTX-2 (one and two stages) | `Lightricks/LTX-2` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | -| LTX-2.3 (one and two stages) | `Lightricks/LTX-2.3` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| Wan2.1 Fun 1.3B InP | `weizhou03/Wan2.1-Fun-1.3B-InP-Diffusers` | 480p | ✅ | ✅ | ✅ | ⭕ | ❌ | ❌ | ✅ | +| Helios Base | `BestWishYsh/Helios-Base` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| Helios Mid | `BestWishYsh/Helios-Mid` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| Helios Distilled | `BestWishYsh/Helios-Distilled` | 720p | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| LTX-2 (one/two-stage/TI2V) | `Lightricks/LTX-2` | 768×512
1536×1024 | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +| LTX-2.3 (one/two-stage/TI2V/HQ) | `Lightricks/LTX-2.3` | 768×512
1536×1024
1920×1088 (HQ default) | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | **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 and LTX-2.3 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`. +3. LTX pipeline selection: + - One-stage: `--pipeline-class-name LTX2Pipeline` + - Two-stage: `--pipeline-class-name LTX2TwoStagePipeline` + - Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`) + - LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ). + - 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`. - For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`. -4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {legacy,snapshot,resident}`: +4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`: - `snapshot` is the default and recommended mode. - `resident` usually provides the best latency/throughput but uses much more VRAM. - - `legacy` preserves the historical switching path for fallback/debug. + - `original` keeps official two-stage semantics without the premerged stage-2 transformer path. + - Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`. ### Image Generation Models diff --git a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx index b5bf1d6be..0c9e038b6 100644 --- a/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx +++ b/docs_new/docs/sglang-diffusion/compatibility_matrix.mdx @@ -285,7 +285,7 @@ Optimization columns are abbreviated to keep the matrix readable: ❌ - LTX-2 (one and two stages) + LTX-2 (one/two-stage/TI2V) Lightricks/LTX-2 768×512
1536×1024 ❌ @@ -297,9 +297,9 @@ Optimization columns are abbreviated to keep the matrix readable: ❌ - LTX-2.3 (one and two stages) + LTX-2.3 (one/two-stage/TI2V/HQ) Lightricks/LTX-2.3 - 768×512
1536×1024 + 768×512
1536×1024
1920×1088 (HQ default) ❌ ❌ ❌ @@ -315,12 +315,18 @@ Optimization columns are abbreviated to keep the matrix readable: 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 and LTX-2.3 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`. - - For LTX models, the `Resolution` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`. -4. LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {legacy,snapshot,resident}`: +3. LTX pipeline selection: + - One-stage: `--pipeline-class-name LTX2Pipeline` + - Two-stage: `--pipeline-class-name LTX2TwoStagePipeline` + - Two-stage HQ: `--pipeline-class-name LTX2TwoStageHQPipeline` (HQ defaults to 1920×1088; you can still override `--width/--height`) + - LTX-2 and LTX-2.3 support both T2V and TI2V (`--image-path`) on one-stage and two-stage pipelines (including HQ). + - 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`. + - For LTX models, the `Resolutions` column uses output video `width×height` semantics, matching `sglang generate --width ... --height ...`. +4. LTX-2 / LTX-2.3 two-stage also supports `--ltx2-two-stage-device-mode {original,snapshot,resident}`: - `snapshot` is the default and recommended mode. - `resident` usually provides the best latency/throughput but uses much more VRAM. - - `legacy` preserves the historical switching path for fallback/debug. + - `original` keeps official two-stage semantics without the premerged stage-2 transformer path. + - Example (one prior run): `original` `154.67s`, `snapshot` `114.05s`, `resident` `75.71s`; peak VRAM trend is `original < snapshot < resident`. ### Image Generation Models 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 03478f9ae..c32d4bd80 100644 --- a/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/pipeline_configs/ltx_2.py @@ -244,6 +244,9 @@ class LTX2PipelineConfig(PipelineConfig): def tokenize_prompt(self, prompt: list[str], tokenizer, tok_kwargs) -> dict: # Adapted from diffusers_pipeline.py _get_gemma_prompt_embeds # But we only need tokenization here, the embedding happens in TextEncodingStage + # Official LTX Gemma tokenizer trims surrounding whitespace before + # tokenization. + prompt = [text.strip() for text in prompt] # Gemma expects left padding for chat-style prompts tokenizer.padding_side = "left" diff --git a/python/sglang/multimodal_gen/configs/sample/ltx_2.py b/python/sglang/multimodal_gen/configs/sample/ltx_2.py index dec4f918d..86837abb7 100644 --- a/python/sglang/multimodal_gen/configs/sample/ltx_2.py +++ b/python/sglang/multimodal_gen/configs/sample/ltx_2.py @@ -82,3 +82,38 @@ class LTX23SamplingParams(LTX2SamplingParams): "audio_stg_blocks": self.audio_stg_blocks, } return extra + + +@dataclasses.dataclass +class LTX23HQSamplingParams(LTX23SamplingParams): + """Sampling parameters matching official LTX-2.3 HQ two-stage defaults.""" + + height: int = 1088 + width: int = 1920 + num_inference_steps: int = 15 + distilled_lora_strength_stage_1: float = 0.25 + distilled_lora_strength_stage_2: float = 0.5 + + video_cfg_scale: float = 3.0 + video_stg_scale: float = 0.0 + video_rescale_scale: float = 0.45 + video_modality_scale: float = 3.0 + video_skip_step: int = 0 + video_stg_blocks: list[int] = field(default_factory=list) + + audio_cfg_scale: float = 7.0 + audio_stg_scale: float = 0.0 + audio_rescale_scale: float = 1.0 + audio_modality_scale: float = 3.0 + audio_skip_step: int = 0 + audio_stg_blocks: list[int] = field(default_factory=list) + + def build_request_extra(self) -> dict[str, Any]: + extra = super().build_request_extra() + extra["ltx2_distilled_lora_strength_stage_1"] = float( + self.distilled_lora_strength_stage_1 + ) + extra["ltx2_distilled_lora_strength_stage_2"] = float( + self.distilled_lora_strength_stage_2 + ) + 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 48434c209..179f2215c 100644 --- a/python/sglang/multimodal_gen/configs/sample/sampling_params.py +++ b/python/sglang/multimodal_gen/configs/sample/sampling_params.py @@ -571,18 +571,28 @@ class SamplingParams: def from_user_sampling_params_args( model_path: str, server_args: "ServerArgs", *args, **kwargs ): + pipeline_class_name = getattr(server_args, "pipeline_class_name", None) try: - sampling_params = SamplingParams.from_pretrained( - model_path, backend=server_args.backend, model_id=server_args.model_id - ) - except (AttributeError, ValueError) as e: + sampling_params = None + if pipeline_class_name: + from sglang.multimodal_gen.registry import get_pipeline_config_classes + + config_classes = get_pipeline_config_classes(pipeline_class_name) + if config_classes is not None: + _, sampling_params_cls = config_classes + sampling_params = sampling_params_cls() + + if sampling_params is None: + sampling_params = SamplingParams.from_pretrained( + model_path, + backend=server_args.backend, + model_id=server_args.model_id, + ) + except (AttributeError, ValueError): # Handle safetensors files or other cases where model_index.json is not available # Use appropriate SamplingParams based on pipeline_class_name from registry if os.path.isfile(model_path) and model_path.endswith(".safetensors"): # Determine which sampling params to use based on pipeline_class_name - pipeline_class_name = getattr(server_args, "pipeline_class_name", None) - - # Try to get SamplingParams from registry from sglang.multimodal_gen.registry import get_pipeline_config_classes config_classes = ( @@ -616,7 +626,7 @@ class SamplingParams: user_kwargs = dict(kwargs) user_kwargs.pop("diffusers_kwargs", None) - user_sampling_params = SamplingParams(*args, **user_kwargs) + user_sampling_params = type(sampling_params)(*args, **user_kwargs) # TODO: refactor sampling_params._merge_with_user_params( user_sampling_params, explicit_fields=set(user_kwargs.keys()) @@ -982,7 +992,14 @@ class SamplingParams: for field in dataclasses.fields(user_params): field_name = field.name user_value = getattr(user_params, field_name) - default_class_value = getattr(SamplingParams, field_name) + if hasattr(SamplingParams, field_name): + default_class_value = getattr(SamplingParams, field_name) + elif field.default is not dataclasses.MISSING: + default_class_value = field.default + elif field.default_factory is not dataclasses.MISSING: + default_class_value = field.default_factory() + else: + default_class_value = dataclasses.MISSING is_user_modified = user_value != default_class_value or ( explicit_fields is not None and field_name in explicit_fields diff --git a/python/sglang/multimodal_gen/registry.py b/python/sglang/multimodal_gen/registry.py index aaa49b08e..2b4cb52ea 100644 --- a/python/sglang/multimodal_gen/registry.py +++ b/python/sglang/multimodal_gen/registry.py @@ -99,6 +99,7 @@ from sglang.multimodal_gen.configs.sample.hunyuan import ( from sglang.multimodal_gen.configs.sample.hunyuan3d import Hunyuan3DSamplingParams from sglang.multimodal_gen.configs.sample.ltx_2 import ( LTX2SamplingParams, + LTX23HQSamplingParams, LTX23SamplingParams, ) from sglang.multimodal_gen.configs.sample.mova import ( @@ -632,6 +633,11 @@ def _register_configs(): lambda path: "ltx-2.3" in path.lower(), ], ) + # register dedicated sampling params for LTX2TwoStageHQPipeline + _PIPELINE_CONFIG_REGISTRY.setdefault( + "LTX2TwoStageHQPipeline", + (LTX2PipelineConfig, LTX23HQSamplingParams), + ) # Hunyuan register_configs( diff --git a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py b/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py index 8f4a216ca..3ef45b65c 100644 --- a/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py +++ b/python/sglang/multimodal_gen/runtime/entrypoints/cli/generate.py @@ -31,6 +31,32 @@ from sglang.multimodal_gen.utils import FlexibleArgumentParser logger = init_logger(__name__) +def _resolve_cli_sampling_params_cls(server_args: ServerArgs) -> type[SamplingParams]: + pipeline_class_name = getattr(server_args, "pipeline_class_name", None) + if pipeline_class_name: + from sglang.multimodal_gen.registry import get_pipeline_config_classes + + config_classes = get_pipeline_config_classes(pipeline_class_name) + if config_classes is not None: + _, sampling_params_cls = config_classes + return sampling_params_cls + + try: + from sglang.multimodal_gen.registry import get_model_info + + model_info = get_model_info( + server_args.model_path, + backend=server_args.backend, + model_id=server_args.model_id, + ) + if model_info is not None: + return model_info.sampling_param_cls + except Exception as exc: + logger.debug("Falling back to base SamplingParams for CLI parsing: %s", exc) + + return SamplingParams + + def add_multimodal_gen_generate_args(parser: argparse.ArgumentParser): """Add the arguments for the generate command.""" parser.add_argument( @@ -130,6 +156,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None args.request_id = "mocked_fake_id_for_offline_generate" server_args = ServerArgs.from_cli_args(args, unknown_args) + sampling_params_cls = _resolve_cli_sampling_params_cls(server_args) sampling_params_kwargs = {} config_file = getattr(args, "config", None) @@ -137,7 +164,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None if config_file: config_args = ServerArgs.load_config_file(config_file) or {} sampling_param_fields = { - field.name for field in dataclasses.fields(SamplingParams) + field.name for field in dataclasses.fields(sampling_params_cls) } sampling_params_kwargs.update( { @@ -147,7 +174,7 @@ def generate_cmd(args: argparse.Namespace, unknown_args: list[str] | None = None } ) - sampling_params_kwargs.update(SamplingParams.get_cli_args(args)) + sampling_params_kwargs.update(sampling_params_cls.get_cli_args(args)) _apply_output_file_path_override(args, sampling_params_kwargs) sampling_params_kwargs["request_id"] = generate_request_id() diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index ed318ee69..961000208 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -36,6 +36,9 @@ from sglang.multimodal_gen.utils import get_mixed_precision_state torch._dynamo.config.recompile_limit = 64 +LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024 + + class BaseLayerWithLoRA(nn.Module): def __init__( @@ -176,16 +179,48 @@ class BaseLayerWithLoRA(nn.Module): """ # Merge all LoRA adapters in order for lora_A, lora_B, _, lora_strength in lora_list: - lora_delta = self.slice_lora_b_weights( - lora_B.to(data) - ) @ self.slice_lora_a_weights(lora_A.to(data)) - # Apply lora_alpha / lora_rank scaling for consistency with forward() - if self.lora_alpha is not None and self.lora_rank is not None: - if self.lora_alpha != self.lora_rank: - lora_delta = lora_delta * (self.lora_alpha / self.lora_rank) - if lora_delta.dim() > 2: - lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1]) - data += lora_strength * lora_delta + lora_A_sliced = self.slice_lora_a_weights(lora_A.to(data)) + lora_B_sliced = self.slice_lora_b_weights(lora_B.to(data)) + + scale = lora_strength + if ( + self.lora_alpha is not None + and self.lora_rank is not None + and self.lora_alpha != self.lora_rank + ): + scale *= self.lora_alpha / self.lora_rank + + if not isinstance(lora_B_sliced, torch.Tensor): + lora_delta = lora_B_sliced @ lora_A_sliced + if isinstance(lora_delta, torch.Tensor) and lora_delta.dim() > 2: + lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1]) + data.add_(lora_delta, alpha=scale) + continue + + if lora_A_sliced.dim() > 2 or lora_B_sliced.dim() > 2: + lora_delta = lora_B_sliced @ lora_A_sliced + if lora_delta.dim() > 2: + lora_delta = lora_delta.reshape(-1, lora_delta.shape[-1]) + data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data + data_2d.add_(lora_delta, alpha=scale) + continue + + data_2d = data.reshape(-1, data.shape[-1]) if data.dim() > 2 else data + lora_B_2d = ( + lora_B_sliced.reshape(-1, lora_B_sliced.shape[-1]) + if lora_B_sliced.dim() > 2 + else lora_B_sliced + ) + + chunk_rows = max( + 1, + LORA_MERGE_CHUNK_BYTES + // (data_2d.shape[-1] * max(1, data_2d.element_size())), + ) + for start in range(0, lora_B_2d.shape[0], chunk_rows): + end = min(start + chunk_rows, lora_B_2d.shape[0]) + chunk_delta = lora_B_2d[start:end] @ lora_A_sliced + data_2d[start:end].add_(chunk_delta, alpha=scale) @torch.no_grad() def merge_lora_weights(self, strength: float | None = None) -> None: 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 ec255d63b..f34a0f8db 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py +++ b/python/sglang/multimodal_gen/runtime/models/dits/ltx_2.py @@ -1219,6 +1219,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): return timestep return timestep.amax(dim=tuple(range(1, timestep.ndim))) + def _scale_timestep_for_adaln(self, timestep: torch.Tensor) -> torch.Tensor: + ltx_variant = str(getattr(self.config.arch_config, "ltx_variant", "ltx_2")) + if ltx_variant == "ltx_2_3" and bool( + getattr(self, "_sglang_use_ltx23_hq_timestep_semantics", False) + ): + return timestep * float(self.timestep_scale_multiplier) + return timestep + def _validate_tp_config(self, *, arch: LTX2ArchConfig, tp_size: int) -> None: """Validate TP-related dimension constraints (fail-fast).""" if tp_size < 1: @@ -1663,8 +1671,10 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): audio_hidden_states, _ = self.audio_patchify_proj(audio_hidden_states) # 3. Prepare timestep embeddings # 3.1. Prepare global modality (video and audio) timestep embedding and modulation parameters + timestep_for_adaln = self._scale_timestep_for_adaln(timestep) + audio_timestep_for_adaln = self._scale_timestep_for_adaln(audio_timestep) temb, embedded_timestep = self.adaln_single( - timestep.flatten(), + timestep_for_adaln.flatten(), hidden_dtype=hidden_states.dtype, ) temb = temb.view(batch_size, -1, temb.size(-1)) @@ -1673,7 +1683,7 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): ) temb_audio, audio_embedded_timestep = self.audio_adaln_single( - audio_timestep.flatten(), + audio_timestep_for_adaln.flatten(), hidden_dtype=audio_hidden_states.dtype, ) temb_audio = temb_audio.view(batch_size, -1, temb_audio.size(-1)) @@ -1688,8 +1698,9 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): if prompt_timestep is None else prompt_timestep ) + prompt_timestep_for_adaln = self._scale_timestep_for_adaln(prompt_timestep) temb_prompt, _ = self.prompt_adaln_single( - prompt_timestep.flatten(), hidden_dtype=hidden_states.dtype + prompt_timestep_for_adaln.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: @@ -1698,8 +1709,11 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): if audio_prompt_timestep is None else audio_prompt_timestep ) + audio_prompt_timestep_for_adaln = self._scale_timestep_for_adaln( + audio_prompt_timestep + ) temb_audio_prompt, _ = self.audio_prompt_adaln_single( - audio_prompt_timestep.flatten(), + audio_prompt_timestep_for_adaln.flatten(), hidden_dtype=audio_hidden_states.dtype, ) temb_audio_prompt = temb_audio_prompt.view( @@ -1714,8 +1728,14 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): prompt_timestep, audio_prompt_timestep, ) + av_ca_video_timestep_for_adaln = self._scale_timestep_for_adaln( + av_ca_video_timestep + ) + av_ca_audio_timestep_for_adaln = self._scale_timestep_for_adaln( + av_ca_audio_timestep + ) temb_ca_scale_shift, _ = self.av_ca_video_scale_shift_adaln_single( - av_ca_video_timestep.flatten(), hidden_dtype=hidden_dtype + av_ca_video_timestep_for_adaln.flatten(), hidden_dtype=hidden_dtype ) temb_ca_scale_shift = temb_ca_scale_shift.view( batch_size, -1, temb_ca_scale_shift.shape[-1] @@ -1723,20 +1743,21 @@ class LTX2VideoTransformer3DModel(CachableDiT, OffloadableDiTMixin): av_ca_gate_factor = self._get_av_ca_gate_timestep_factor() temb_ca_gate, _ = self.av_ca_a2v_gate_adaln_single( - av_ca_video_timestep.flatten() * av_ca_gate_factor, + av_ca_video_timestep_for_adaln.flatten() * av_ca_gate_factor, hidden_dtype=hidden_dtype, ) temb_ca_gate = temb_ca_gate.view(batch_size, -1, temb_ca_gate.shape[-1]) temb_ca_audio_scale_shift, _ = self.av_ca_audio_scale_shift_adaln_single( - av_ca_audio_timestep.flatten(), hidden_dtype=audio_hidden_states.dtype + av_ca_audio_timestep_for_adaln.flatten(), + hidden_dtype=audio_hidden_states.dtype, ) temb_ca_audio_scale_shift = temb_ca_audio_scale_shift.view( batch_size, -1, temb_ca_audio_scale_shift.shape[-1] ) temb_ca_audio_gate, _ = self.av_ca_v2a_gate_adaln_single( - av_ca_audio_timestep.flatten() * av_ca_gate_factor, + av_ca_audio_timestep_for_adaln.flatten() * av_ca_gate_factor, hidden_dtype=audio_hidden_states.dtype, ) temb_ca_audio_gate = temb_ca_audio_gate.view( 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 0d189adba..0ad0e767b 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -6,9 +6,11 @@ import torch from diffusers import FlowMatchEulerDiscreteScheduler from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( + LTX2PipelineConfig, is_ltx23_native_variant, sync_ltx23_runtime_vae_markers, ) +from sglang.multimodal_gen.configs.sample.ltx_2 import LTX23HQSamplingParams from sglang.multimodal_gen.runtime.distributed import get_local_torch_device from sglang.multimodal_gen.runtime.loader.component_loaders.component_loader import ( PipelineComponentLoader, @@ -120,12 +122,18 @@ def build_official_ltx2_sigmas( stretch: bool = True, terminal: float = 0.1, default_number_of_tokens: int = MAX_SHIFT_ANCHOR, + number_of_tokens: int | None = None, ) -> 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 + tokens = ( + int(number_of_tokens) + if number_of_tokens is not None + else int(default_number_of_tokens) + ) + sigma_shift = float(tokens) * mm + b non_zero_mask = sigmas != 0 shifted = torch.where( @@ -136,8 +144,9 @@ def build_official_ltx2_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) + if bool(torch.any(one_minus_z != 0)): + scale_factor = one_minus_z[-1] / (1.0 - terminal) + shifted[non_zero_mask] = 1.0 - (one_minus_z / scale_factor) return shifted[:-1].tolist() @@ -148,7 +157,28 @@ class LTX2SigmaPreparationStage(PipelineStage): def forward(self, batch: Req, server_args: ServerArgs) -> Req: batch.extra["ltx2_phase"] = "stage1" if is_ltx23_native_variant(server_args.pipeline_config.vae_config.arch_config): - batch.sigmas = build_official_ltx2_sigmas(int(batch.num_inference_steps)) + # Resolution-aware sigma shift is only required for the HQ pipeline + # (which targets 1080p+ resolutions and was aligned against official + # LTX-2.3 HQ sigmas). Legacy one-stage and two-stage LTX-2.3 paths + # were baselined against the constant-anchor schedule. + if server_args.pipeline_class_name == "LTX2TwoStageHQPipeline": + latent_num_frames = (int(batch.num_frames) - 1) // int( + server_args.pipeline_config.vae_temporal_compression + ) + 1 + latent_height = int(batch.height) // int( + server_args.pipeline_config.vae_scale_factor + ) + latent_width = int(batch.width) // int( + server_args.pipeline_config.vae_scale_factor + ) + batch.sigmas = build_official_ltx2_sigmas( + int(batch.num_inference_steps), + number_of_tokens=latent_num_frames * latent_height * latent_width, + ) + else: + batch.sigmas = build_official_ltx2_sigmas( + int(batch.num_inference_steps) + ) else: batch.sigmas = np.linspace( 1.0, @@ -171,7 +201,11 @@ def _add_ltx2_front_stages(pipeline: ComposedPipelineBase): ) -def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase): +def _add_ltx2_stage1_generation_stages( + pipeline: ComposedPipelineBase, + *, + denoising_sampler_name: str = "euler", +): pipeline.add_stage(LTX2SigmaPreparationStage()) pipeline.add_standard_timestep_preparation_stage( prepare_extra_kwargs=[prepare_ltx2_mu] @@ -191,6 +225,7 @@ def _add_ltx2_stage1_generation_stages(pipeline: ComposedPipelineBase): scheduler=pipeline.get_module("scheduler"), vae=pipeline.get_module("vae"), audio_vae=pipeline.get_module("audio_vae"), + sampler_name=denoising_sampler_name, pipeline=pipeline, ), ] @@ -542,18 +577,25 @@ class LTX2TwoStageDeviceManager: module.to("cpu") return + pin_memory = bool( + self.server_args.pin_cpu_memory and torch.get_device_module().is_available() + ) for name, param in module.named_parameters(): snapshot = param_snapshots.get(name) if snapshot is None: - raise KeyError( - f"Missing CPU parameter snapshot for {module_name}.{name}" + snapshot = self._clone_cpu_tensor_snapshot( + param.data, pin_memory=pin_memory ) + param_snapshots[name] = snapshot param.data = snapshot for name, buffer in module.named_buffers(): snapshot = buffer_snapshots.get(name) if snapshot is None: - raise KeyError(f"Missing CPU buffer snapshot for {module_name}.{name}") + snapshot = self._clone_cpu_tensor_snapshot( + buffer.data, pin_memory=pin_memory + ) + buffer_snapshots[name] = snapshot # Preserve runtime-updated buffers (e.g., lazily built caches) when # releasing back to CPU snapshots. if buffer.device.type == "cuda": @@ -666,6 +708,10 @@ class LTX2TwoStageDeviceManager: class LTX2TwoStagePipeline(_BaseLTX2Pipeline): pipeline_name = "LTX2TwoStagePipeline" STAGE_2_DISTILLED_SIGMA_VALUES = [0.909375, 0.725, 0.421875, 0.0] + STAGE_1_DISTILLED_LORA_STRENGTH = 0.0 + STAGE_2_DISTILLED_LORA_STRENGTH = 1.0 + STAGE_1_DENOISING_SAMPLER_NAME = "euler" + STAGE_2_DENOISING_SAMPLER_NAME = "euler" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -712,6 +758,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): self._stage1_lora_path = server_args.lora_path self._stage1_lora_scale = float(server_args.lora_scale) self._active_lora_phase = None + self._active_lora_signature = None self._use_premerged_stage2_transformer = False def _initialize_premerged_stage2_transformer(self, server_args: ServerArgs) -> None: @@ -733,7 +780,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): lora_nickname="ltx2_stage2_distilled", lora_path=self._distilled_lora_path, target="transformer_2", - strength=1.0, + strength=self.STAGE_2_DISTILLED_LORA_STRENGTH, merge_weights=True, ) @@ -757,16 +804,48 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): "resident", ) - def _can_short_circuit_lora_switch(self, phase: str) -> bool: - return ( - phase in ("stage1", "stage2") - and self._use_premerged_stage2_transformer - and self._stage1_lora_path is None - ) + def _get_stage_distilled_lora_strength( + self, phase: str, batch: Req | None + ) -> float: + if phase == "stage1": + default_strength = self.STAGE_1_DISTILLED_LORA_STRENGTH + extra_key = "ltx2_distilled_lora_strength_stage_1" + elif phase == "stage2": + default_strength = self.STAGE_2_DISTILLED_LORA_STRENGTH + extra_key = "ltx2_distilled_lora_strength_stage_2" + else: + raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}") + + if batch is None: + return float(default_strength) + + request_strength = batch.extra.get(extra_key) + if request_strength is None: + return float(default_strength) + return float(request_strength) + + def _can_short_circuit_lora_switch( + self, phase: str, batch: Req | None = None + ) -> bool: + distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch) + if phase == "stage1": + return ( + self._use_premerged_stage2_transformer + and self._stage1_lora_path is None + and distilled_lora_strength == 0.0 + ) + if phase == "stage2": + return ( + self._use_premerged_stage2_transformer + and self._stage1_lora_path is None + and distilled_lora_strength == self.STAGE_2_DISTILLED_LORA_STRENGTH + ) + return False def _build_lora_switch_spec( - self, phase: str + self, phase: str, batch: Req | None = None ) -> tuple[list[str], list[str], list[float], list[str]]: + distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch) lora_nicknames: list[str] = [] lora_paths: list[str] = [] lora_strengths: list[float] = [] @@ -778,33 +857,42 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): lora_paths.append(self._stage1_lora_path) lora_strengths.append(self._stage1_lora_scale) lora_targets.append("transformer") + if distilled_lora_strength != 0.0: + lora_nicknames.append("ltx2_stage1_distilled") + lora_paths.append(self._distilled_lora_path) + lora_strengths.append(distilled_lora_strength) + lora_targets.append("transformer") elif phase == "stage2": if self._stage1_lora_path: lora_nicknames.append("ltx2_stage1_base") lora_paths.append(self._stage1_lora_path) lora_strengths.append(self._stage1_lora_scale) lora_targets.append("transformer") - lora_nicknames.append("ltx2_stage2_distilled") - lora_paths.append(self._distilled_lora_path) - lora_strengths.append(1.0) - lora_targets.append("transformer") + if distilled_lora_strength != 0.0: + lora_nicknames.append("ltx2_stage2_distilled") + lora_paths.append(self._distilled_lora_path) + lora_strengths.append(distilled_lora_strength) + lora_targets.append("transformer") else: raise ValueError(f"Unknown LTX2 two-stage LoRA phase: {phase}") return lora_nicknames, lora_paths, lora_strengths, lora_targets - def switch_lora_phase(self, phase: str) -> None: - if phase == self._active_lora_phase: + def switch_lora_phase(self, phase: str, batch: Req | None = None) -> None: + distilled_lora_strength = self._get_stage_distilled_lora_strength(phase, batch) + phase_signature = (phase, distilled_lora_strength) + if phase_signature == self._active_lora_signature: return if self._device_manager.switch_phase( phase - ) and self._can_short_circuit_lora_switch(phase): + ) and self._can_short_circuit_lora_switch(phase, batch): self._active_lora_phase = phase + self._active_lora_signature = phase_signature return lora_nicknames, lora_paths, lora_strengths, lora_targets = ( - self._build_lora_switch_spec(phase) + self._build_lora_switch_spec(phase, batch) ) if lora_nicknames: set_lora_kwargs = dict( @@ -830,6 +918,7 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): self.deactivate_lora_weights(target="transformer") self._active_lora_phase = phase + self._active_lora_signature = phase_signature def create_pipeline_stages(self, server_args: ServerArgs): _add_ltx2_front_stages(self) @@ -837,7 +926,10 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): self.add_stage( LTX2LoRASwitchStage(pipeline=self, phase="stage1"), ) - _add_ltx2_stage1_generation_stages(self) + _add_ltx2_stage1_generation_stages( + self, + denoising_sampler_name=self.STAGE_1_DENOISING_SAMPLER_NAME, + ) self.add_stages( [ LTX2UpsampleStage( @@ -863,10 +955,21 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): vae=self.get_module("vae"), audio_vae=self.get_module("audio_vae"), pipeline=self, + sampler_name=self.STAGE_2_DENOISING_SAMPLER_NAME, ), ] ) _add_ltx2_decoding_stage(self) -EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline] +class LTX2TwoStageHQPipeline(LTX2TwoStagePipeline): + pipeline_name = "LTX2TwoStageHQPipeline" + pipeline_config_cls = LTX2PipelineConfig + sampling_params_cls = LTX23HQSamplingParams + STAGE_1_DISTILLED_LORA_STRENGTH = 0.25 + STAGE_2_DISTILLED_LORA_STRENGTH = 0.5 + STAGE_1_DENOISING_SAMPLER_NAME = "res2s" + STAGE_2_DENOISING_SAMPLER_NAME = "res2s" + + +EntryClass = [LTX2Pipeline, LTX2TwoStagePipeline, LTX2TwoStageHQPipeline] 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 bfabc2c4c..313aafd36 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 @@ -112,8 +112,16 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): vae=None, audio_vae=None, pipeline=None, + sampler_name: str = "euler", ): - super().__init__(transformer, scheduler, vae, audio_vae, pipeline=pipeline) + super().__init__( + transformer, + scheduler, + vae, + audio_vae, + pipeline=pipeline, + sampler_name=sampler_name, + ) self.distilled_sigmas = torch.tensor(distilled_sigmas) @staticmethod @@ -171,6 +179,41 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): return False return "LTX-2.3" not in str(getattr(server_args, "model_path", "")) + @staticmethod + def _build_stage2_renoise_generator( + batch: Req, reference_tensor: torch.Tensor + ) -> torch.Generator: + seeds = getattr(batch, "seeds", None) + if seeds: + seed = int(seeds[0]) + else: + seed = int(getattr(batch, "seed", 10)) + device = reference_tensor.device + dtype = reference_tensor.dtype + generator = torch.Generator(device=device).manual_seed(seed) + video_shape = batch.extra.get("ltx2_stage1_packed_video_shape") + audio_shape = batch.extra.get("ltx2_stage1_packed_audio_shape") + if video_shape is not None: + _ = torch.randn( + tuple(video_shape), device=device, dtype=dtype, generator=generator + ) + if audio_shape is not None: + _ = torch.randn( + tuple(audio_shape), device=device, dtype=dtype, generator=generator + ) + return generator + + @staticmethod + def _ltx2_renoise_like( + reference_tensor: torch.Tensor, generator: torch.Generator + ) -> torch.Tensor: + return torch.randn( + reference_tensor.shape, + device=reference_tensor.device, + dtype=reference_tensor.dtype, + generator=generator, + ) + def forward(self, batch: Req, server_args: ServerArgs) -> Req: """Run the distilled refinement schedule on top of the shared AV denoiser.""" batch.extra["ltx2_phase"] = "stage2" @@ -191,6 +234,23 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): if self._should_reset_stage2_generators(server_args): self._reset_stage2_generators(batch) noise_scale = float(self.distilled_sigmas[0].item()) + # HQ pipeline uses a dedicated, deterministic renoise generator seeded + # from the request seed and advanced by stage-1 packed shapes to match + # official LTX-2.3 HQ output. Legacy two-stage paths were baselined + # against `batch.generator`'s natural advance through stage-1, so keep + # them on the original `_randn_like_with_batch_generators` sampling. + is_hq_pipeline = server_args.pipeline_class_name == "LTX2TwoStageHQPipeline" + if is_hq_pipeline: + video_reference_for_gen = ( + batch.latents if isinstance(batch.latents, torch.Tensor) else None + ) + if video_reference_for_gen is None: + video_reference_for_gen = batch.audio_latents + renoise_generator = self._build_stage2_renoise_generator( + batch, video_reference_for_gen + ) + else: + renoise_generator = None if is_native_ti2v: prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state( latents=batch.latents, @@ -199,34 +259,63 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): zero_clean_latent=True, clean_latent_background=batch.ltx2_ti2v_clean_latent_background, ) - video_noise = self._randn_like_with_batch_generators( - prepared_latents, batch - ) + if is_hq_pipeline: + video_noise = self._ltx2_renoise_like( + prepared_latents, renoise_generator + ) + else: + video_noise = self._randn_like_with_batch_generators( + prepared_latents, batch + ) scaled_mask = ( denoise_mask.to(device=prepared_latents.device, dtype=torch.float32) * noise_scale ) - batch.latents = ( - video_noise * scaled_mask + prepared_latents * (1 - scaled_mask) - ).to(prepared_latents.dtype) + if is_hq_pipeline: + batch.latents = ( + video_noise.float() * scaled_mask + + prepared_latents.float() * (1.0 - scaled_mask) + ).to(prepared_latents.dtype) + else: + batch.latents = ( + video_noise * scaled_mask + prepared_latents * (1 - scaled_mask) + ).to(prepared_latents.dtype) else: - video_noise = self._randn_like_with_batch_generators(batch.latents, batch) - batch.latents = ( - video_noise * noise_scale + batch.latents * (1 - noise_scale) - ).to(batch.latents.dtype) + if is_hq_pipeline: + video_noise = self._ltx2_renoise_like(batch.latents, renoise_generator) + batch.latents = ( + video_noise.float() * noise_scale + + batch.latents.float() * (1.0 - noise_scale) + ).to(batch.latents.dtype) + else: + video_noise = self._randn_like_with_batch_generators( + batch.latents, batch + ) + batch.latents = ( + video_noise * noise_scale + batch.latents * (1 - noise_scale) + ).to(batch.latents.dtype) if isinstance(batch.audio_latents, torch.Tensor): - audio_noise = self._randn_like_with_batch_generators( - batch.audio_latents, batch - ) - audio_scaled_mask = ( - torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32) - * noise_scale - ) - batch.audio_latents = ( - audio_noise * audio_scaled_mask - + batch.audio_latents * (1 - audio_scaled_mask) - ).to(batch.audio_latents.dtype) + if is_hq_pipeline: + audio_noise = self._ltx2_renoise_like( + batch.audio_latents, renoise_generator + ) + batch.audio_latents = ( + audio_noise.float() * noise_scale + + batch.audio_latents.float() * (1.0 - noise_scale) + ).to(batch.audio_latents.dtype) + else: + audio_noise = self._randn_like_with_batch_generators( + batch.audio_latents, batch + ) + audio_scaled_mask = ( + torch.ones_like(batch.audio_latents[..., :1], dtype=torch.float32) + * noise_scale + ) + batch.audio_latents = ( + audio_noise * audio_scaled_mask + + batch.audio_latents * (1 - audio_scaled_mask) + ).to(batch.audio_latents.dtype) if not is_ltx23_native_variant( server_args.pipeline_config.vae_config.arch_config ): @@ -244,8 +333,25 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): self.scheduler = copy.deepcopy(original_scheduler) distilled_device = self.scheduler.sigmas.device - self.scheduler.sigmas = self.distilled_sigmas.to(distilled_device) num_steps = len(self.distilled_sigmas) - 1 + # HQ pipeline extends the sigma schedule so the final step targets a + # small non-zero sigma (0.0011) instead of 0.0, matching official + # LTX-2.3 HQ's last-step behavior. Legacy two-stage baselines used the + # un-extended schedule (final step goes to 0.0). + if ( + server_args.pipeline_class_name == "LTX2TwoStageHQPipeline" + and self.distilled_sigmas[-1].item() == 0.0 + ): + scheduler_sigmas = torch.cat( + [ + self.distilled_sigmas[:-1], + torch.tensor([0.0011, 0.0], dtype=self.distilled_sigmas.dtype), + ], + dim=0, + ) + else: + scheduler_sigmas = self.distilled_sigmas + self.scheduler.sigmas = scheduler_sigmas.to(distilled_device) self.scheduler.num_inference_steps = num_steps self.scheduler.timesteps = (self.distilled_sigmas[:num_steps] * 1000).to( distilled_device 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 70c993c19..ce390de30 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 @@ -155,14 +155,16 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): latent_shape = server_args.pipeline_config.prepare_latent_shape( batch, batch_size, num_frames ) + packed_video_shape = self._packed_video_latent_shape( + latent_shape, server_args.pipeline_config + ) latents = randn_tensor( - self._packed_video_latent_shape( - latent_shape, server_args.pipeline_config - ), + packed_video_shape, generator=generator, device=device, dtype=dtype, ) + batch.extra["ltx2_stage1_packed_video_shape"] = tuple(packed_video_shape) latent_ids = server_args.pipeline_config.maybe_prepare_latent_ids(latents) if latent_ids is not None: @@ -196,13 +198,14 @@ class LTX2AVLatentPreparationStage(LatentPreparationStage): latent_shape = server_args.pipeline_config.prepare_audio_latent_shape( batch, batch_size, batch.num_frames ) - + packed_audio_shape = self._packed_audio_latent_shape(latent_shape) audio_latents = randn_tensor( - self._packed_audio_latent_shape(latent_shape), + packed_audio_shape, generator=generator, device=device, dtype=dtype, ) + batch.extra["ltx2_stage1_packed_audio_shape"] = tuple(packed_audio_shape) else: audio_latents = audio_latents.to(device) audio_latents = server_args.pipeline_config.maybe_pack_audio_latents( diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py index e24dc4ccb..492ba154f 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/ltx_2_denoising.py @@ -1,7 +1,9 @@ import copy +from contextlib import contextmanager from dataclasses import dataclass, field import torch +from diffusers.utils.torch_utils import randn_tensor from sglang.multimodal_gen.configs.pipeline_configs.ltx_2 import ( is_ltx23_native_variant, @@ -22,6 +24,9 @@ from sglang.multimodal_gen.runtime.server_args import ( is_ltx2_two_stage_pipeline_name, ) +LTX23_RES2S_STEP_NOISE_SEED = -1 +LTX23_RES2S_SUBSTEP_NOISE_SEED = 9999 + @dataclass(slots=True) class LTX2DenoisingContext(DenoisingContext): @@ -41,6 +46,10 @@ class LTX2DenoisingContext(DenoisingContext): last_denoised_video: torch.Tensor | None = None last_denoised_audio: torch.Tensor | None = None trajectory_audio_latents: list[torch.Tensor] = field(default_factory=list) + use_native_hq_res2s_sde_noise: bool = False + use_ltx23_hq_timestep_semantics: bool = False + res2s_step_noise_generator: torch.Generator | None = None + res2s_substep_noise_generator: torch.Generator | None = None @dataclass(slots=True) @@ -94,10 +103,43 @@ class LTX2DenoisingStage(DenoisingStage): "audio_encoder_attention_mask", ) - def __init__(self, transformer, scheduler, vae=None, **kwargs): + def __init__( + self, + transformer, + scheduler, + vae=None, + *, + sampler_name: str = "euler", + **kwargs, + ): super().__init__( transformer=transformer, scheduler=scheduler, vae=vae, **kwargs ) + self.sampler_name = sampler_name + + @staticmethod + def _randn_like_with_batch_generators( + reference_tensor: torch.Tensor, batch: Req + ) -> torch.Tensor: + generator = getattr(batch, "generator", None) + if isinstance(generator, list): + bsz = int(reference_tensor.shape[0]) + valid_generators = [g for g in generator if isinstance(g, torch.Generator)] + if len(valid_generators) == 1: + generator = valid_generators[0] + elif len(valid_generators) >= bsz: + generator = valid_generators[:bsz] + else: + generator = None + elif not isinstance(generator, torch.Generator): + generator = None + + return randn_tensor( + reference_tensor.shape, + generator=generator, + device=reference_tensor.device, + dtype=reference_tensor.dtype, + ) @staticmethod def _get_video_latent_num_frames_for_model( @@ -186,6 +228,293 @@ class LTX2DenoisingStage(DenoisingStage): factor = rescale_scale * factor + (1.0 - rescale_scale) return pred * factor + @staticmethod + def _ltx2_channelwise_normalize(noise: torch.Tensor) -> torch.Tensor: + return noise.sub_(noise.mean(dim=(-2, -1), keepdim=True)).div_( + noise.std(dim=(-2, -1), keepdim=True) + ) + + @classmethod + def _ltx2_res2s_new_noise( + cls, + reference_tensor: torch.Tensor, + generator: torch.Generator, + ) -> torch.Tensor: + noise = torch.randn( + reference_tensor.shape, + generator=generator, + dtype=torch.float64, + device=reference_tensor.device, + ) + noise = (noise - noise.mean()) / noise.std() + return cls._ltx2_channelwise_normalize(noise) + + @staticmethod + def _ltx2_init_res2s_noise_generators(ctx: LTX2DenoisingContext) -> None: + reference_tensor = ( + ctx.latents if isinstance(ctx.latents, torch.Tensor) else ctx.audio_latents + ) + if reference_tensor is None: + raise ValueError("LTX-2 res2s requires video or audio latents.") + device = reference_tensor.device + ctx.res2s_step_noise_generator = torch.Generator(device=device).manual_seed( + LTX23_RES2S_STEP_NOISE_SEED + ) + ctx.res2s_substep_noise_generator = torch.Generator(device=device).manual_seed( + LTX23_RES2S_SUBSTEP_NOISE_SEED + ) + + @classmethod + def _ltx2_res2s_noise_like( + cls, + reference_tensor: torch.Tensor, + ctx: LTX2DenoisingContext, + *, + substep: bool, + ) -> torch.Tensor: + generator = ( + ctx.res2s_substep_noise_generator + if substep + else ctx.res2s_step_noise_generator + ) + if generator is None: + raise ValueError("LTX-2 res2s noise generator was not initialized.") + return cls._ltx2_res2s_new_noise(reference_tensor, generator).to( + dtype=reference_tensor.dtype + ) + + @staticmethod + def _ltx2_apply_clean_latent_mask( + latents: torch.Tensor, + ctx: LTX2DenoisingContext, + ) -> torch.Tensor: + if ctx.denoise_mask is None or ctx.clean_latent is None: + return latents + return ( + latents.float() * ctx.denoise_mask + + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) + ).to(dtype=latents.dtype) + + @staticmethod + def _ltx2_phi_1(neg_h: torch.Tensor) -> torch.Tensor: + small = neg_h.abs() < 1e-4 + series = 1.0 + 0.5 * neg_h + (neg_h * neg_h) / 6.0 + return torch.where(small, series, torch.expm1(neg_h) / neg_h) + + @classmethod + def _ltx2_phi_2(cls, neg_h: torch.Tensor) -> torch.Tensor: + small = neg_h.abs() < 1e-4 + series = 0.5 + neg_h / 6.0 + (neg_h * neg_h) / 24.0 + exact = (torch.expm1(neg_h) - neg_h) / (neg_h * neg_h) + return torch.where(small, series, exact) + + @classmethod + def _ltx2_get_res2s_coefficients( + cls, h: torch.Tensor, c2: float = 0.5 + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + a21 = c2 * cls._ltx2_phi_1(-h * c2) + b2 = cls._ltx2_phi_2(-h) / c2 + b1 = cls._ltx2_phi_1(-h) - b2 + return a21, b1, b2 + + @staticmethod + def _ltx2_get_sde_coeff( + sigma_next: torch.Tensor, + *, + sigma_up: torch.Tensor | None = None, + sigma_down: torch.Tensor | None = None, + sigma_max: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if sigma_down is not None: + alpha_ratio = (1.0 - sigma_next) / (1.0 - sigma_down) + sigma_up = torch.sqrt( + torch.clamp( + sigma_next.square() - sigma_down.square() * alpha_ratio.square(), + min=0.0, + ) + ) + elif sigma_up is not None: + sigma_up = torch.minimum(sigma_up, sigma_next * 0.9999) + sigmax = sigma_max if sigma_max is not None else torch.ones_like(sigma_next) + sigma_signal = sigmax - sigma_next + sigma_residual = torch.sqrt( + torch.clamp(sigma_next.square() - sigma_up.square(), min=0.0) + ) + alpha_ratio = sigma_signal + sigma_residual + sigma_down = sigma_residual / alpha_ratio + else: + alpha_ratio = torch.ones_like(sigma_next) + sigma_down = sigma_next + sigma_up = torch.zeros_like(sigma_next) + return ( + torch.nan_to_num(alpha_ratio), + torch.nan_to_num(sigma_down), + torch.nan_to_num(sigma_up), + ) + + @classmethod + def _ltx2_res2s_sde_step( + cls, + *, + sample: torch.Tensor, + denoised_sample: torch.Tensor, + sigma: torch.Tensor, + sigma_next: torch.Tensor, + noise: torch.Tensor, + eta: float = 0.5, + ) -> torch.Tensor: + alpha_ratio, sigma_down, sigma_up = cls._ltx2_get_sde_coeff( + sigma_next, + sigma_up=sigma_next * eta, + ) + if bool((sigma_up == 0).any()) or bool((sigma_next == 0).any()): + return denoised_sample.to(dtype=sample.dtype) + eps_next = (sample - denoised_sample) / (sigma - sigma_next) + denoised_next = sample - sigma * eps_next + x_noised = ( + alpha_ratio * (denoised_next + sigma_down * eps_next) + sigma_up * noise + ) + return x_noised.to(dtype=sample.dtype) + + def _ltx2_stage2_res2s_step( + self, + *, + ctx: "LTX2DenoisingContext", + batch: Req, + sigma: torch.Tensor, + sigma_next: torch.Tensor, + model_video_velocity: torch.Tensor, + model_audio_velocity: torch.Tensor, + midpoint_model_call, + ) -> tuple[torch.Tensor, torch.Tensor]: + """res2s RK2 step for unguided stage-2 refinement (HQ pipeline). + + Converts velocity -> x_0 denoised estimates, runs the official res2s + update (midpoint SDE, bongmath anchor refinement, midpoint re-eval, + final RK2 combination with SDE noise). Mirrors the guided stage-1 res2s + math but without CFG/STG (stage-2 HQ uses the simple CFG path). + """ + sigma_val = float(sigma.item()) + sigma_next_val = float(sigma_next.item()) + + if sigma_val == 0.0: + denoised_video = ctx.latents.float() + denoised_audio = ctx.audio_latents.float() + else: + denoised_video = ctx.latents.float() - sigma * model_video_velocity.float() + denoised_audio = ( + ctx.audio_latents.float() - sigma * model_audio_velocity.float() + ) + + if sigma_val == 0.0 or sigma_next_val == 0.0: + next_video = denoised_video.to(dtype=ctx.latents.dtype) + next_audio = denoised_audio.to(dtype=ctx.audio_latents.dtype) + next_video = self._ltx2_apply_clean_latent_mask(next_video, ctx) + return next_video, next_audio + + sigma_d = sigma.double() + sigma_next_d = sigma_next.double() + h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) + a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0)) + + anchor_video = ctx.latents.double() + anchor_audio = ctx.audio_latents.double() + eps1_video = denoised_video.double() - anchor_video + eps1_audio = denoised_audio.double() - anchor_audio + + midpoint_video_det = anchor_video + h * a21 * eps1_video + midpoint_audio_det = anchor_audio + h * a21 * eps1_audio + + sub_noise_video = ( + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators(ctx.latents, batch).float() + ) + sub_noise_audio = ( + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=True).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.audio_latents, batch + ).float() + ) + midpoint_video_latents = self._ltx2_res2s_sde_step( + sample=anchor_video, + denoised_sample=midpoint_video_det, + sigma=sigma_d, + sigma_next=sub_sigma, + noise=sub_noise_video, + ) + midpoint_audio_latents = self._ltx2_res2s_sde_step( + sample=anchor_audio, + denoised_sample=midpoint_audio_det, + sigma=sigma_d, + sigma_next=sub_sigma, + noise=sub_noise_audio, + ) + midpoint_video_latents = self._ltx2_apply_clean_latent_mask( + midpoint_video_latents.to(dtype=ctx.latents.dtype), ctx + ) + midpoint_audio_latents = midpoint_audio_latents.to( + dtype=ctx.audio_latents.dtype + ) + + # Bongmath anchor refinement for the first stage-2 step. + if float(h.item()) < 0.5 and sigma_val > 0.03: + x_mid_v = midpoint_video_latents.double() + x_mid_a = midpoint_audio_latents.double() + for _ in range(100): + anchor_video = x_mid_v - h * a21 * eps1_video + eps1_video = denoised_video.double() - anchor_video + anchor_audio = x_mid_a - h * a21 * eps1_audio + eps1_audio = denoised_audio.double() - anchor_audio + + mid_v, mid_a = midpoint_model_call( + midpoint_video_latents, midpoint_audio_latents, sub_sigma + ) + + midpoint_denoised_video = midpoint_video_latents.float() - sub_sigma * mid_v + midpoint_denoised_audio = midpoint_audio_latents.float() - sub_sigma * mid_a + + eps2_video = midpoint_denoised_video.double() - anchor_video + eps2_audio = midpoint_denoised_audio.double() - anchor_audio + + next_video_det = anchor_video + h * (b1 * eps1_video + b2 * eps2_video) + next_audio_det = anchor_audio + h * (b1 * eps1_audio + b2 * eps2_audio) + + step_noise_video = ( + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators(ctx.latents, batch).float() + ) + step_noise_audio = ( + self._ltx2_res2s_noise_like(ctx.audio_latents, ctx, substep=False).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.audio_latents, batch + ).float() + ) + next_video = self._ltx2_res2s_sde_step( + sample=anchor_video, + denoised_sample=next_video_det, + sigma=sigma_d, + sigma_next=sigma_next_d, + noise=step_noise_video, + ) + next_audio = self._ltx2_res2s_sde_step( + sample=anchor_audio, + denoised_sample=next_audio_det, + sigma=sigma_d, + sigma_next=sigma_next_d, + noise=step_noise_audio, + ) + + next_video = self._ltx2_apply_clean_latent_mask( + next_video.to(dtype=ctx.latents.dtype), ctx + ) + next_audio = next_audio.to(dtype=ctx.audio_latents.dtype) + return next_video, next_audio + @staticmethod def _prepare_ltx2_ti2v_clean_state( latents: torch.Tensor, @@ -369,21 +698,41 @@ class LTX2DenoisingStage(DenoisingStage): ) batch_size = int(latent_model_input.shape[0]) - timestep = step.t_device.expand(batch_size) + use_raw_sigma_timestep = ctx.use_ltx23_hq_timestep_semantics + use_ltx23_two_stage_prompt_timestep = ( + ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage + ) + timestep = ( + sigma.to(device=ctx.latents.device, dtype=torch.float32).expand(batch_size) + if use_raw_sigma_timestep + else step.t_device.to( + device=ctx.latents.device, dtype=torch.float32 + ).expand(batch_size) + ) if ctx.denoise_mask is not None: - timestep_video = timestep.unsqueeze(-1) * ctx.denoise_mask.squeeze(-1) - elif ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + if use_raw_sigma_timestep: + timestep_video = ( + timestep.view(batch_size, *([1] * (ctx.denoise_mask.ndim - 1))) + * ctx.denoise_mask + ) + else: + timestep_video = timestep.unsqueeze(-1) * ctx.denoise_mask.squeeze(-1) + elif use_raw_sigma_timestep: + timestep_video = timestep.view(batch_size, 1, 1).expand( + batch_size, int(latent_model_input.shape[1]), 1 + ) + elif use_ltx23_two_stage_prompt_timestep: timestep_video = timestep.view(batch_size, 1).expand( batch_size, int(latent_model_input.shape[1]) ) else: timestep_video = timestep - if ( - ctx.is_ltx23_variant - and not ctx.use_ltx23_legacy_one_stage - and audio_latent_model_input.ndim == 3 - ): + if use_raw_sigma_timestep and audio_latent_model_input.ndim == 3: + timestep_audio = timestep.view(batch_size, 1, 1).expand( + batch_size, int(audio_latent_model_input.shape[1]), 1 + ) + elif use_ltx23_two_stage_prompt_timestep and audio_latent_model_input.ndim == 3: timestep_audio = timestep.view(batch_size, 1).expand( batch_size, int(audio_latent_model_input.shape[1]) ) @@ -392,7 +741,14 @@ class LTX2DenoisingStage(DenoisingStage): prompt_timestep_video = None prompt_timestep_audio = None - if ctx.is_ltx23_variant and not ctx.use_ltx23_legacy_one_stage: + if ctx.use_ltx23_hq_timestep_semantics: + prompt_timestep_video = sigma.to( + device=ctx.latents.device, dtype=torch.float32 + ).expand(batch_size) + prompt_timestep_audio = sigma.to( + device=ctx.audio_latents.device, dtype=torch.float32 + ).expand(batch_size) + elif use_ltx23_two_stage_prompt_timestep: timestep_scale_multiplier = float( getattr(step.current_model, "timestep_scale_multiplier", 1000) ) @@ -576,6 +932,39 @@ class LTX2DenoisingStage(DenoisingStage): return True return int(getattr(batch, "sp_video_start_frame", 0)) == 0 + @staticmethod + def _should_use_native_hq_res2s_sde_noise(server_args: ServerArgs) -> bool: + return server_args.pipeline_class_name == "LTX2TwoStageHQPipeline" + + @staticmethod + def _should_use_ltx23_hq_timestep_semantics(server_args: ServerArgs) -> bool: + return server_args.pipeline_class_name == "LTX2TwoStageHQPipeline" + + @staticmethod + @contextmanager + def _temporary_ltx23_hq_timestep_semantics(model, enabled: bool): + attr = "_sglang_use_ltx23_hq_timestep_semantics" + previous = bool(getattr(model, attr, False)) + setattr(model, attr, enabled) + try: + yield + finally: + setattr(model, attr, previous) + + @contextmanager + def _ltx2_model_forward_context( + self, + ctx: LTX2DenoisingContext, + step: DenoisingStepState, + ): + with self._temporary_ltx23_hq_timestep_semantics( + step.current_model, ctx.use_ltx23_hq_timestep_semantics + ): + with set_forward_context( + current_timestep=step.step_index, attn_metadata=step.attn_metadata + ): + yield + def _prepare_denoising_loop( self, batch: Req, @@ -592,6 +981,14 @@ class LTX2DenoisingStage(DenoisingStage): ctx.use_ltx23_legacy_one_stage = self._should_use_ltx23_legacy_one_stage( server_args ) + ctx.use_native_hq_res2s_sde_noise = ( + ctx.is_ltx23_variant + and self._should_use_native_hq_res2s_sde_noise(server_args) + ) + ctx.use_ltx23_hq_timestep_semantics = ( + ctx.is_ltx23_variant + and self._should_use_ltx23_hq_timestep_semantics(server_args) + ) ctx.stage = ( phase if phase is not None @@ -676,7 +1073,7 @@ class LTX2DenoisingStage(DenoisingStage): else None ) if callable(switch_lora_phase): - switch_lora_phase(ctx.stage) + switch_lora_phase(ctx.stage, batch=batch) ensure_phase_ready = ( getattr(pipeline, "ensure_ltx2_phase_ready", None) if pipeline is not None @@ -688,6 +1085,8 @@ class LTX2DenoisingStage(DenoisingStage): if ctx.audio_scheduler is None: raise ValueError("LTX-2 audio scheduler was not prepared.") ctx.audio_scheduler.set_begin_index(0) + if self.sampler_name == "res2s" and ctx.use_native_hq_res2s_sde_noise: + self._ltx2_init_res2s_noise_generators(ctx) def _prepare_step_attn_metadata( self, @@ -787,9 +1186,7 @@ class LTX2DenoisingStage(DenoisingStage): repeated_attention_mask ) - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): + with self._ltx2_model_forward_context(ctx, step): model_video, model_audio = step.current_model(**model_kwargs) model_video = model_video.float() @@ -804,17 +1201,114 @@ class LTX2DenoisingStage(DenoisingStage): batch.guidance_scale * (model_audio_text - model_audio_uncond) ) - ctx.latents = self.scheduler.step( - model_video, step.t_device, ctx.latents, return_dict=False - )[0] - ctx.audio_latents = ctx.audio_scheduler.step( - model_audio, step.t_device, ctx.audio_latents, return_dict=False - )[0] - if ctx.denoise_mask is not None and ctx.clean_latent is not None: - ctx.latents = ( - ctx.latents.float() * ctx.denoise_mask - + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) - ).to(dtype=ctx.latents.dtype) + if self.sampler_name == "res2s": + # HQ stage-2 uses RK2 res2s here to match official LTX-2.3 HQ + # output. Without this path the scheduler falls back to Euler + # and loses ~3.7 dB against the official canonical. + def _stage2_midpoint_model_call( + video_latents: torch.Tensor, + audio_latents: torch.Tensor, + sigma_value: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + original_video_latents = ctx.latents + original_audio_latents = ctx.audio_latents + ctx.latents = video_latents + ctx.audio_latents = audio_latents + try: + model_inputs_local = self._prepare_ltx2_model_inputs( + ctx, step, batch, server_args, sigma_value + ) + batch_size_local = int( + model_inputs_local.latent_model_input.shape[0] + ) + base_model_kwargs_local = self._build_ltx2_base_model_kwargs( + ctx, batch, model_inputs_local + ) + model_kwargs_local = self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs_local, + encoder_hidden_states=batch.prompt_embeds[0], + audio_encoder_hidden_states=batch.audio_prompt_embeds[0], + encoder_attention_mask=prompt_attention_mask, + ) + if batch.do_classifier_free_guidance: + cfg_batch_size = batch_size_local * 2 + model_kwargs_local = self._repeat_ltx2_model_kwargs_batch( + model_kwargs_local, cfg_batch_size + ) + model_kwargs_local["encoder_hidden_states"] = torch.cat( + [ + batch.negative_prompt_embeds[0], + batch.prompt_embeds[0], + ], + dim=0, + ) + model_kwargs_local["audio_encoder_hidden_states"] = ( + torch.cat( + [ + batch.negative_audio_prompt_embeds[0], + batch.audio_prompt_embeds[0], + ], + dim=0, + ) + ) + if self._should_pass_ltx2_text_attention_mask(ctx): + repeated_attention_mask = self._cat_or_none( + [ + self._get_ltx_prompt_attention_mask( + batch, + is_ltx23_variant=( + ctx.is_ltx23_variant + and not ctx.use_ltx23_legacy_one_stage + ), + negative=True, + ), + prompt_attention_mask, + ] + ) + model_kwargs_local["encoder_attention_mask"] = ( + repeated_attention_mask + ) + model_kwargs_local["audio_encoder_attention_mask"] = ( + repeated_attention_mask + ) + + with self._ltx2_model_forward_context(ctx, step): + mid_v, mid_a = step.current_model(**model_kwargs_local) + + mid_v = mid_v.float() + mid_a = mid_a.float() + if batch.do_classifier_free_guidance: + mid_v_u, mid_v_t = mid_v.chunk(2) + mid_a_u, mid_a_t = mid_a.chunk(2) + mid_v = mid_v_u + batch.guidance_scale * (mid_v_t - mid_v_u) + mid_a = mid_a_u + batch.guidance_scale * (mid_a_t - mid_a_u) + return mid_v, mid_a + finally: + ctx.latents = original_video_latents + ctx.audio_latents = original_audio_latents + + ctx.latents, ctx.audio_latents = self._ltx2_stage2_res2s_step( + ctx=ctx, + batch=batch, + sigma=sigma, + sigma_next=sigma_next, + model_video_velocity=model_video, + model_audio_velocity=model_audio, + midpoint_model_call=_stage2_midpoint_model_call, + ) + else: + ctx.latents = self.scheduler.step( + model_video, step.t_device, ctx.latents, return_dict=False + )[0] + ctx.audio_latents = ctx.audio_scheduler.step( + model_audio, step.t_device, ctx.audio_latents, return_dict=False + )[0] + if ctx.denoise_mask is not None and ctx.clean_latent is not None: + ctx.latents = ( + ctx.latents.float() * ctx.denoise_mask + + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) + ).to(dtype=ctx.latents.dtype) ctx.latents = self.post_forward_for_ti2v_task( batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z ) @@ -847,324 +1341,491 @@ class LTX2DenoisingStage(DenoisingStage): float(stage1_guider_params["video_modality_scale"]) != 1.0 or float(stage1_guider_params["audio_modality_scale"]) != 1.0 ) - - if ctx.use_ltx23_legacy_one_stage: - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): - v_pos, a_v_pos = step.current_model( - **self._build_ltx2_model_kwargs( - ctx, - base_model_kwargs, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ) - ) - v_neg, a_v_neg = step.current_model( - **self._build_ltx2_model_kwargs( - ctx, - base_model_kwargs, - encoder_hidden_states=negative_encoder_hidden_states, - audio_encoder_hidden_states=negative_audio_encoder_hidden_states, - encoder_attention_mask=negative_encoder_attention_mask, - ) - ) - - v_pos = v_pos.float() - a_v_pos = a_v_pos.float() - v_neg = v_neg.float() - a_v_neg = a_v_neg.float() - - v_ptb = None - a_v_ptb = None - if need_perturbed: - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): - v_ptb, a_v_ptb = step.current_model( - **self._build_ltx2_model_kwargs( - ctx, - base_model_kwargs, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) - ) - v_ptb = v_ptb.float() - a_v_ptb = a_v_ptb.float() - - v_mod = None - a_v_mod = None - if need_modality: - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): - v_mod, a_v_mod = step.current_model( - **self._build_ltx2_model_kwargs( - ctx, - base_model_kwargs, - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) - ) - v_mod = v_mod.float() - a_v_mod = a_v_mod.float() - else: - # NOTE: this flag must be identical across all SP ranks so that - # every rank executes the same number of model-forward calls (each - # of which contains NCCL collectives). - # _should_apply_ltx2_ti2v() is SP-rank-dependent (only the rank owning the first latent - # frame returns True), so we must NOT use it here. - # Instead we check the rank-invariant attribute that is always set on every - # rank when the request is a TI2V request. - use_split_two_stage_ti2v_guider = ( + # NOTE: this flag must be identical across all SP ranks so that every + # rank executes the same number of model-forward calls (each of which + # contains NCCL collectives). + use_split_stage1_guided_passes = ( + server_args.pipeline_class_name == "LTX2TwoStageHQPipeline" + or ( is_ltx2_two_stage_pipeline_name(server_args.pipeline_class_name) and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 ) + ) - pass_specs: list[LTX2GuidancePassSpec] = [ - LTX2GuidancePassSpec( - name="cond", - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - ), - LTX2GuidancePassSpec( - name="neg", - encoder_hidden_states=negative_encoder_hidden_states, - audio_encoder_hidden_states=negative_audio_encoder_hidden_states, - encoder_attention_mask=negative_encoder_attention_mask, - ), - ] - if need_perturbed: - pass_specs.append( - LTX2GuidancePassSpec( - name="perturbed", - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - skip_video_self_attn_blocks=tuple( - stage1_guider_params["video_stg_blocks"] - ), - skip_audio_self_attn_blocks=tuple( - stage1_guider_params["audio_stg_blocks"] - ), - ) + def evaluate_stage1_guided_x0( + *, + video_latents: torch.Tensor, + audio_latents: torch.Tensor, + sigma_value: torch.Tensor, + update_skip_cache: bool, + ) -> tuple[torch.Tensor, torch.Tensor]: + original_video_latents = ctx.latents + original_audio_latents = ctx.audio_latents + ctx.latents = video_latents + ctx.audio_latents = audio_latents + try: + model_inputs_local = self._prepare_ltx2_model_inputs( + ctx, step, batch, server_args, sigma_value ) - if need_modality: - pass_specs.append( - LTX2GuidancePassSpec( - name="modality", - encoder_hidden_states=encoder_hidden_states, - audio_encoder_hidden_states=audio_encoder_hidden_states, - encoder_attention_mask=encoder_attention_mask, - disable_a2v_cross_attn=True, - disable_v2a_cross_attn=True, - ) + batch_size_local = int(model_inputs_local.latent_model_input.shape[0]) + base_model_kwargs_local = self._build_ltx2_base_model_kwargs( + ctx, batch, model_inputs_local ) - num_passes = len(pass_specs) - expanded_batch_size = batch_size * num_passes - perturbation_configs = tuple( - { - "skip_video_self_attn_blocks": pass_spec.skip_video_self_attn_blocks, - "skip_audio_self_attn_blocks": pass_spec.skip_audio_self_attn_blocks, - "skip_a2v_cross_attn": pass_spec.disable_a2v_cross_attn, - "skip_v2a_cross_attn": pass_spec.disable_v2a_cross_attn, - } - for pass_spec in pass_specs - for _ in range(batch_size) - ) - batched_model_kwargs = self._repeat_ltx2_model_kwargs_batch( - base_model_kwargs, expanded_batch_size - ) - batched_model_kwargs = self._build_ltx2_model_kwargs( - ctx, - batched_model_kwargs, - encoder_hidden_states=torch.cat( - [pass_spec.encoder_hidden_states for pass_spec in pass_specs], dim=0 - ), - audio_encoder_hidden_states=torch.cat( - [pass_spec.audio_encoder_hidden_states for pass_spec in pass_specs], - dim=0, - ), - encoder_attention_mask=self._cat_or_none( - [pass_spec.encoder_attention_mask for pass_spec in pass_specs] - ), - ) - if use_split_two_stage_ti2v_guider: - split_sizes = [1] * expanded_batch_size - batched_video_chunks = [] - batched_audio_chunks = [] - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): - for model_kwargs_chunk, perturbation_config in zip( - self._split_ltx2_model_kwargs( - batched_model_kwargs, split_sizes - ), - perturbation_configs, - strict=True, - ): - model_kwargs_chunk["perturbation_configs"] = ( - perturbation_config, + if ctx.use_ltx23_legacy_one_stage: + with self._ltx2_model_forward_context(ctx, step): + v_pos, a_v_pos = step.current_model( + **self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs_local, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ) ) - video_chunk, audio_chunk = step.current_model( - **model_kwargs_chunk + v_neg, a_v_neg = step.current_model( + **self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs_local, + encoder_hidden_states=negative_encoder_hidden_states, + audio_encoder_hidden_states=negative_audio_encoder_hidden_states, + encoder_attention_mask=negative_encoder_attention_mask, + ) ) - batched_video_chunks.append(video_chunk) - batched_audio_chunks.append(audio_chunk) - batched_video = torch.cat(batched_video_chunks, dim=0) - batched_audio = torch.cat(batched_audio_chunks, dim=0) - else: - with set_forward_context( - current_timestep=step.step_index, attn_metadata=step.attn_metadata - ): - batched_video, batched_audio = step.current_model( - **batched_model_kwargs, - perturbation_configs=perturbation_configs, + v_pos = v_pos.float() + a_v_pos = a_v_pos.float() + v_neg = v_neg.float() + a_v_neg = a_v_neg.float() + + v_ptb = None + a_v_ptb = None + if need_perturbed: + with self._ltx2_model_forward_context(ctx, step): + v_ptb, a_v_ptb = step.current_model( + **self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs_local, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + skip_video_self_attn_blocks=tuple( + stage1_guider_params["video_stg_blocks"] + ), + skip_audio_self_attn_blocks=tuple( + stage1_guider_params["audio_stg_blocks"] + ), + ) + ) + v_ptb = v_ptb.float() + a_v_ptb = a_v_ptb.float() + + v_mod = None + a_v_mod = None + if need_modality: + with self._ltx2_model_forward_context(ctx, step): + v_mod, a_v_mod = step.current_model( + **self._build_ltx2_model_kwargs( + ctx, + base_model_kwargs_local, + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + disable_a2v_cross_attn=True, + disable_v2a_cross_attn=True, + ) + ) + v_mod = v_mod.float() + a_v_mod = a_v_mod.float() + else: + pass_specs: list[LTX2GuidancePassSpec] = [ + LTX2GuidancePassSpec( + name="cond", + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + ), + LTX2GuidancePassSpec( + name="neg", + encoder_hidden_states=negative_encoder_hidden_states, + audio_encoder_hidden_states=negative_audio_encoder_hidden_states, + encoder_attention_mask=negative_encoder_attention_mask, + ), + ] + if need_perturbed: + pass_specs.append( + LTX2GuidancePassSpec( + name="perturbed", + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + skip_video_self_attn_blocks=tuple( + stage1_guider_params["video_stg_blocks"] + ), + skip_audio_self_attn_blocks=tuple( + stage1_guider_params["audio_stg_blocks"] + ), + ) + ) + if need_modality: + pass_specs.append( + LTX2GuidancePassSpec( + name="modality", + encoder_hidden_states=encoder_hidden_states, + audio_encoder_hidden_states=audio_encoder_hidden_states, + encoder_attention_mask=encoder_attention_mask, + disable_a2v_cross_attn=True, + disable_v2a_cross_attn=True, + ) + ) + + num_passes = len(pass_specs) + expanded_batch_size = batch_size_local * num_passes + perturbation_configs = tuple( + { + "skip_video_self_attn_blocks": pass_spec.skip_video_self_attn_blocks, + "skip_audio_self_attn_blocks": pass_spec.skip_audio_self_attn_blocks, + "skip_a2v_cross_attn": pass_spec.disable_a2v_cross_attn, + "skip_v2a_cross_attn": pass_spec.disable_v2a_cross_attn, + } + for pass_spec in pass_specs + for _ in range(batch_size_local) ) + batched_model_kwargs = self._repeat_ltx2_model_kwargs_batch( + base_model_kwargs_local, expanded_batch_size + ) + batched_model_kwargs = self._build_ltx2_model_kwargs( + ctx, + batched_model_kwargs, + encoder_hidden_states=torch.cat( + [ + pass_spec.encoder_hidden_states + for pass_spec in pass_specs + ], + dim=0, + ), + audio_encoder_hidden_states=torch.cat( + [ + pass_spec.audio_encoder_hidden_states + for pass_spec in pass_specs + ], + dim=0, + ), + encoder_attention_mask=self._cat_or_none( + [ + pass_spec.encoder_attention_mask + for pass_spec in pass_specs + ] + ), + ) + if use_split_stage1_guided_passes: + split_sizes = [1] * expanded_batch_size + batched_video_chunks = [] + batched_audio_chunks = [] + with self._ltx2_model_forward_context(ctx, step): + for model_kwargs_chunk, perturbation_config in zip( + self._split_ltx2_model_kwargs( + batched_model_kwargs, split_sizes + ), + perturbation_configs, + strict=True, + ): + model_kwargs_chunk["perturbation_configs"] = ( + perturbation_config, + ) + video_chunk, audio_chunk = step.current_model( + **model_kwargs_chunk + ) + batched_video_chunks.append(video_chunk) + batched_audio_chunks.append(audio_chunk) - batched_video = batched_video.float() - batched_audio = batched_audio.float() - pass_outputs = { - pass_spec.name: ( - video_chunk, - audio_chunk, + batched_video = torch.cat(batched_video_chunks, dim=0) + batched_audio = torch.cat(batched_audio_chunks, dim=0) + else: + with self._ltx2_model_forward_context(ctx, step): + batched_video, batched_audio = step.current_model( + **batched_model_kwargs, + perturbation_configs=perturbation_configs, + ) + + batched_video = batched_video.float() + batched_audio = batched_audio.float() + pass_outputs = { + pass_spec.name: ( + video_chunk, + audio_chunk, + ) + for pass_spec, video_chunk, audio_chunk in zip( + pass_specs, + batched_video.chunk(num_passes, dim=0), + batched_audio.chunk(num_passes, dim=0), + strict=True, + ) + } + v_pos, a_v_pos = pass_outputs["cond"] + v_neg, a_v_neg = pass_outputs["neg"] + v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None)) + v_mod, a_v_mod = pass_outputs.get("modality", (None, None)) + + sigma_value_float = float(sigma_value.item()) + video_sigma_for_x0: float | torch.Tensor = sigma_value_float + if ctx.denoise_mask is not None: + video_sigma_for_x0 = sigma_value.to( + device=video_latents.device, dtype=torch.float32 + ) * ctx.denoise_mask.squeeze(-1) + + denoised_video_local = self._ltx2_velocity_to_x0( + video_latents, v_pos, video_sigma_for_x0 ) - for pass_spec, video_chunk, audio_chunk in zip( - pass_specs, - batched_video.chunk(num_passes, dim=0), - batched_audio.chunk(num_passes, dim=0), - strict=True, + denoised_audio_local = self._ltx2_velocity_to_x0( + audio_latents, a_v_pos, sigma_value_float ) - } - v_pos, a_v_pos = pass_outputs["cond"] - v_neg, a_v_neg = pass_outputs["neg"] - v_ptb, a_v_ptb = pass_outputs.get("perturbed", (None, None)) - v_mod, a_v_mod = pass_outputs.get("modality", (None, None)) + denoised_video_neg = self._ltx2_velocity_to_x0( + video_latents, v_neg, video_sigma_for_x0 + ) + denoised_audio_neg = self._ltx2_velocity_to_x0( + audio_latents, a_v_neg, sigma_value_float + ) + denoised_video_perturbed = ( + None + if v_ptb is None + else self._ltx2_velocity_to_x0( + video_latents, v_ptb, video_sigma_for_x0 + ) + ) + denoised_audio_perturbed = ( + None + if a_v_ptb is None + else self._ltx2_velocity_to_x0( + audio_latents, a_v_ptb, sigma_value_float + ) + ) + denoised_video_modality = ( + None + if v_mod is None + else self._ltx2_velocity_to_x0( + video_latents, v_mod, video_sigma_for_x0 + ) + ) + denoised_audio_modality = ( + None + if a_v_mod is None + else self._ltx2_velocity_to_x0( + audio_latents, a_v_mod, sigma_value_float + ) + ) + + guided_video = self._ltx2_calculate_guided_x0( + cond=denoised_video_local, + uncond_text=denoised_video_neg, + uncond_perturbed=( + denoised_video_perturbed + if denoised_video_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_video_modality + if denoised_video_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["video_cfg_scale"]), + stg_scale=float(stage1_guider_params["video_stg_scale"]), + rescale_scale=float(stage1_guider_params["video_rescale_scale"]), + modality_scale=float(stage1_guider_params["video_modality_scale"]), + ) + if video_skip and ctx.last_denoised_video is not None: + denoised_video_local = ctx.last_denoised_video + else: + denoised_video_local = guided_video + if update_skip_cache: + ctx.last_denoised_video = guided_video + + guided_audio = self._ltx2_calculate_guided_x0( + cond=denoised_audio_local, + uncond_text=denoised_audio_neg, + uncond_perturbed=( + denoised_audio_perturbed + if denoised_audio_perturbed is not None + else 0.0 + ), + uncond_modality=( + denoised_audio_modality + if denoised_audio_modality is not None + else 0.0 + ), + cfg_scale=float(stage1_guider_params["audio_cfg_scale"]), + stg_scale=float(stage1_guider_params["audio_stg_scale"]), + rescale_scale=float(stage1_guider_params["audio_rescale_scale"]), + modality_scale=float(stage1_guider_params["audio_modality_scale"]), + ) + if audio_skip and ctx.last_denoised_audio is not None: + denoised_audio_local = ctx.last_denoised_audio + else: + denoised_audio_local = guided_audio + if update_skip_cache: + ctx.last_denoised_audio = guided_audio + + denoised_video_local = self._ltx2_apply_clean_latent_mask( + denoised_video_local, ctx + ) + return denoised_video_local, denoised_audio_local + finally: + ctx.latents = original_video_latents + ctx.audio_latents = original_audio_latents sigma_val = float(sigma.item()) - video_sigma_for_x0: float | torch.Tensor = sigma_val - if ctx.denoise_mask is not None: - video_sigma_for_x0 = sigma.to( - device=ctx.latents.device, dtype=torch.float32 - ) * ctx.denoise_mask.squeeze(-1) - - denoised_video = self._ltx2_velocity_to_x0( - ctx.latents, v_pos, video_sigma_for_x0 - ) - denoised_audio = self._ltx2_velocity_to_x0( - ctx.audio_latents, a_v_pos, sigma_val - ) - denoised_video_neg = self._ltx2_velocity_to_x0( - ctx.latents, v_neg, video_sigma_for_x0 - ) - denoised_audio_neg = self._ltx2_velocity_to_x0( - ctx.audio_latents, a_v_neg, sigma_val - ) - denoised_video_perturbed = ( - None - if v_ptb is None - else self._ltx2_velocity_to_x0(ctx.latents, v_ptb, video_sigma_for_x0) - ) - denoised_audio_perturbed = ( - None - if a_v_ptb is None - else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_ptb, sigma_val) - ) - denoised_video_modality = ( - None - if v_mod is None - else self._ltx2_velocity_to_x0(ctx.latents, v_mod, video_sigma_for_x0) - ) - denoised_audio_modality = ( - None - if a_v_mod is None - else self._ltx2_velocity_to_x0(ctx.audio_latents, a_v_mod, sigma_val) + denoised_video, denoised_audio = evaluate_stage1_guided_x0( + video_latents=ctx.latents, + audio_latents=ctx.audio_latents, + sigma_value=sigma, + update_skip_cache=True, ) - if not video_skip: - denoised_video = self._ltx2_calculate_guided_x0( - cond=denoised_video, - uncond_text=denoised_video_neg, - uncond_perturbed=( - denoised_video_perturbed - if denoised_video_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_video_modality - if denoised_video_modality is not None - else 0.0 - ), - cfg_scale=float(stage1_guider_params["video_cfg_scale"]), - stg_scale=float(stage1_guider_params["video_stg_scale"]), - rescale_scale=float(stage1_guider_params["video_rescale_scale"]), - modality_scale=float(stage1_guider_params["video_modality_scale"]), - ) - ctx.last_denoised_video = denoised_video - elif ctx.last_denoised_video is not None: - denoised_video = ctx.last_denoised_video + if self.sampler_name == "res2s": + if sigma_val == 0.0 or float(sigma_next.item()) == 0.0: + next_video_latents = denoised_video.to(dtype=ctx.latents.dtype) + next_audio_latents = denoised_audio.to(dtype=ctx.audio_latents.dtype) + else: + sigma_d = sigma.double() + sigma_next_d = sigma_next.double() + h = -torch.log(torch.clamp(sigma_next_d / sigma_d, min=1e-12)) + a21, b1, b2 = self._ltx2_get_res2s_coefficients(h) + sub_sigma = torch.sqrt(torch.clamp(sigma_d * sigma_next_d, min=0.0)) - if not audio_skip: - denoised_audio = self._ltx2_calculate_guided_x0( - cond=denoised_audio, - uncond_text=denoised_audio_neg, - uncond_perturbed=( - denoised_audio_perturbed - if denoised_audio_perturbed is not None - else 0.0 - ), - uncond_modality=( - denoised_audio_modality - if denoised_audio_modality is not None - else 0.0 - ), - cfg_scale=float(stage1_guider_params["audio_cfg_scale"]), - stg_scale=float(stage1_guider_params["audio_stg_scale"]), - rescale_scale=float(stage1_guider_params["audio_rescale_scale"]), - modality_scale=float(stage1_guider_params["audio_modality_scale"]), - ) - ctx.last_denoised_audio = denoised_audio - elif ctx.last_denoised_audio is not None: - denoised_audio = ctx.last_denoised_audio + anchor_video = ctx.latents.double() + anchor_audio = ctx.audio_latents.double() + eps1_video = denoised_video.double() - anchor_video + eps1_audio = denoised_audio.double() - anchor_audio - if ctx.denoise_mask is not None and ctx.clean_latent is not None: - denoised_video = ( - denoised_video * ctx.denoise_mask - + ctx.clean_latent.float() * (1.0 - ctx.denoise_mask) - ).to(denoised_video.dtype) + midpoint_video_deterministic = anchor_video + h * a21 * eps1_video + midpoint_audio_deterministic = anchor_audio + h * a21 * eps1_audio - # 6. Convert x0 predictions back to velocity and update both latent streams. - if sigma_val == 0.0: - v_video = torch.zeros_like(denoised_video) - v_audio = torch.zeros_like(denoised_audio) + substep_video_noise = ( + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=True).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.latents, batch + ).float() + ) + substep_audio_noise = ( + self._ltx2_res2s_noise_like( + ctx.audio_latents, ctx, substep=True + ).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.audio_latents, batch + ).float() + ) + + midpoint_video_latents = self._ltx2_res2s_sde_step( + sample=anchor_video, + denoised_sample=midpoint_video_deterministic, + sigma=sigma_d, + sigma_next=sub_sigma, + noise=substep_video_noise, + ) + midpoint_audio_latents = self._ltx2_res2s_sde_step( + sample=anchor_audio, + denoised_sample=midpoint_audio_deterministic, + sigma=sigma_d, + sigma_next=sub_sigma, + noise=substep_audio_noise, + ) + + midpoint_video_latents = self._ltx2_apply_clean_latent_mask( + midpoint_video_latents.to(dtype=ctx.latents.dtype), + ctx, + ) + midpoint_audio_latents = midpoint_audio_latents.to( + dtype=ctx.audio_latents.dtype + ) + + if float(h.item()) < 0.5 and sigma_val > 0.03: + x_mid_v = midpoint_video_latents.double() + x_mid_a = midpoint_audio_latents.double() + for _ in range(100): + anchor_video = x_mid_v - h * a21 * eps1_video + eps1_video = denoised_video.double() - anchor_video + anchor_audio = x_mid_a - h * a21 * eps1_audio + eps1_audio = denoised_audio.double() - anchor_audio + + midpoint_denoised_video, midpoint_denoised_audio = ( + evaluate_stage1_guided_x0( + video_latents=midpoint_video_latents, + audio_latents=midpoint_audio_latents, + sigma_value=sub_sigma, + update_skip_cache=False, + ) + ) + eps2_video = midpoint_denoised_video.double() - anchor_video + eps2_audio = midpoint_denoised_audio.double() - anchor_audio + + next_video_deterministic = anchor_video + h * ( + b1 * eps1_video + b2 * eps2_video + ) + next_audio_deterministic = anchor_audio + h * ( + b1 * eps1_audio + b2 * eps2_audio + ) + + step_video_noise = ( + self._ltx2_res2s_noise_like(ctx.latents, ctx, substep=False).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.latents, batch + ).float() + ) + step_audio_noise = ( + self._ltx2_res2s_noise_like( + ctx.audio_latents, ctx, substep=False + ).float() + if ctx.use_native_hq_res2s_sde_noise + else self._randn_like_with_batch_generators( + ctx.audio_latents, batch + ).float() + ) + next_video_latents = self._ltx2_res2s_sde_step( + sample=anchor_video, + denoised_sample=next_video_deterministic, + sigma=sigma_d, + sigma_next=sigma_next_d, + noise=step_video_noise, + ) + next_audio_latents = self._ltx2_res2s_sde_step( + sample=anchor_audio, + denoised_sample=next_audio_deterministic, + sigma=sigma_d, + sigma_next=sigma_next_d, + noise=step_audio_noise, + ) + + next_video_latents = self._ltx2_apply_clean_latent_mask( + next_video_latents.to(dtype=ctx.latents.dtype), + ctx, + ) + next_audio_latents = next_audio_latents.to( + dtype=ctx.audio_latents.dtype + ) else: - v_video = ((ctx.latents.float() - denoised_video.float()) / sigma_val).to( - ctx.latents.dtype - ) - v_audio = ( - (ctx.audio_latents.float() - denoised_audio.float()) / sigma_val - ).to(ctx.audio_latents.dtype) + if sigma_val == 0.0: + v_video = torch.zeros_like(denoised_video) + v_audio = torch.zeros_like(denoised_audio) + else: + v_video = ( + (ctx.latents.float() - denoised_video.float()) / sigma_val + ).to(ctx.latents.dtype) + v_audio = ( + (ctx.audio_latents.float() - denoised_audio.float()) / sigma_val + ).to(ctx.audio_latents.dtype) - ctx.latents = (ctx.latents.float() + v_video.float() * dt).to( - dtype=ctx.latents.dtype - ) - ctx.audio_latents = (ctx.audio_latents.float() + v_audio.float() * dt).to( - dtype=ctx.audio_latents.dtype - ) + next_video_latents = (ctx.latents.float() + v_video.float() * dt).to( + dtype=ctx.latents.dtype + ) + next_audio_latents = (ctx.audio_latents.float() + v_audio.float() * dt).to( + dtype=ctx.audio_latents.dtype + ) + + ctx.latents = next_video_latents + ctx.audio_latents = next_audio_latents ctx.latents = self.post_forward_for_ti2v_task( batch, server_args, ctx.reserved_frames_mask, ctx.latents, ctx.z ) diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py index 93baeba6e..c2969e99c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/text_connector.py @@ -47,7 +47,6 @@ class LTX2TextConnectorStage(PipelineStage): # Handle CFG: Concatenate negative and positive inputs if batch.do_classifier_free_guidance: - # Concatenate: [Negative, Positive] prompt_embeds = torch.cat([neg_prompt_embeds, prompt_embeds], dim=0) prompt_attention_mask = torch.cat( @@ -57,7 +56,9 @@ class LTX2TextConnectorStage(PipelineStage): # Prepare additive mask for connectors (as per Diffusers implementation) dtype = prompt_embeds.dtype - additive_attention_mask = (1 - prompt_attention_mask.to(dtype)) * -1000000.0 + additive_attention_mask = (prompt_attention_mask.to(torch.int64) - 1).to( + dtype + ) * torch.finfo(dtype).max # Call connectors # Expects: prompt_embeds, attention_mask, additive_mask=True diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py index f62b46f23..b2d7d20f3 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py @@ -55,7 +55,7 @@ class LTX2LoRASwitchStage(PipelineStage): raise ValueError( "LTX2LoRASwitchStage requires pipeline.switch_lora_phase()" ) - switch_fn(self.phase) + switch_fn(self.phase, batch=batch) batch.extra["ltx2_phase"] = self.phase return batch diff --git a/python/sglang/multimodal_gen/runtime/server_args.py b/python/sglang/multimodal_gen/runtime/server_args.py index 07f7d3af0..cdc82c78f 100644 --- a/python/sglang/multimodal_gen/runtime/server_args.py +++ b/python/sglang/multimodal_gen/runtime/server_args.py @@ -73,7 +73,7 @@ logger = init_logger(__name__) # GPUs on the faster no-offload default while preserving some headroom. WAN_LAYERWISE_OFFLOAD_AUTO_DISABLE_MEM_GB = 130 LTX2_TWO_STAGE_DEVICE_MODES = ("original", "snapshot", "resident") -LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline",) +LTX2_TWO_STAGE_PIPELINE_NAMES = ("LTX2TwoStagePipeline", "LTX2TwoStageHQPipeline") # H200-class GPUs (>=130 GiB total) can usually keep both LTX2 DiTs resident. LTX2_RESIDENT_AUTO_ENABLE_MEM_GB = 130 diff --git a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py index 2decf5a30..ae8d6bf62 100644 --- a/python/sglang/multimodal_gen/runtime/utils/model_overlay.py +++ b/python/sglang/multimodal_gen/runtime/utils/model_overlay.py @@ -26,10 +26,8 @@ logger = init_logger(__name__) # Built-in diffusion model overlay registry. BUILTIN_MODEL_OVERLAY_REGISTRY: dict[str, dict[str, Any]] = { "Lightricks/LTX-2.3": { - # TODO: consider move to lmsys hf repo "overlay_repo_id": "MickJ/LTX-2.3-overlay", - "overlay_revision": "main", - "bundled_overlay_subdir": "ltx_2_3", + "overlay_revision": "e0cc94f279ec16bb87c230134d40319f6ce40c5e", }, } diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index fcac9b162..637800b60 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -332,6 +332,16 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [ # num_frames=33, # ), # ), + DiffusionTestCase( + "ltx_2_3_hq_pipeline", + DiffusionServerArgs( + model_path="Lightricks/LTX-2.3", + extras=[ + "--pipeline-class-name LTX2TwoStageHQPipeline --ltx2-two-stage-device-mode snapshot" + ], + ), + T2I_sampling_params, + ), ] # Skip hunyuan3d on AMD: marching_cubes surface extraction produces invalid SDF on ROCm. @@ -519,6 +529,7 @@ TWO_GPU_CASES = [ "ltx_2_3_two_stage_ti2v_2gpus", DiffusionServerArgs( model_path="Lightricks/LTX-2.3", + ulysses_degree=2, extras=[ "--pipeline-class-name LTX2TwoStagePipeline --ltx2-two-stage-device-mode original" ], @@ -537,6 +548,7 @@ TWO_GPU_CASES = [ "ltx_2.3_two_stage_t2v_2gpus", DiffusionServerArgs( model_path="Lightricks/LTX-2.3", + ulysses_degree=2, extras=[ "--pipeline-class-name LTX2TwoStagePipeline", "--ltx2-two-stage-device-mode original", @@ -619,6 +631,7 @@ TWO_GPU_CASES = [ "ltx_2.3_one_stage_ti2v", DiffusionServerArgs( model_path="Lightricks/LTX-2.3", + ulysses_degree=2, ), TI2V_sampling_params, ), diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 0b23d2e2c..304ee410b 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -2629,6 +2629,47 @@ "expected_avg_denoise_ms": 890.17, "expected_median_denoise_ms": 896.09, "estimated_full_test_time_s": 155.3 + }, + "ltx_2_3_hq_pipeline": { + "stages_ms": { + "InputValidationStage": 0.11, + "TextEncodingStage": 984.78, + "LTX2TextConnectorStage": 30.42, + "LTX2HalveResolutionStage": 0.1, + "LTX2LoRASwitchStage": 0.01, + "LTX2SigmaPreparationStage": 0.36, + "TimestepPreparationStage": 21.28, + "LTX2AVLatentPreparationStage": 0.13, + "LTX2ImageEncodingStage": 0.03, + "LTX2AVDenoisingStage": 20227.05, + "LTX2UpsampleStage": 157.73, + "LTX2RefinementStage": 1676.07, + "LTX2AVDecodingStage": 521.04, + "per_frame_generation": null + }, + "denoise_step_ms": { + "0": 1406.0, + "1": 1362.75, + "2": 1306.18, + "3": 1299.65, + "4": 1282.16, + "5": 1290.32, + "6": 1284.64, + "7": 1265.01, + "8": 1304.06, + "9": 1246.21, + "10": 1102.19, + "11": 1379.4, + "12": 1467.28, + "13": 1469.49, + "14": 734.96, + "15": 547.19, + "16": 543.27, + "17": 539.0 + }, + "expected_e2e_ms": 24150.97, + "expected_avg_denoise_ms": 1157.21, + "expected_median_denoise_ms": 1287.48 } } }