From 05696527eadff3dbdb78646120e70d1f22402230 Mon Sep 17 00:00:00 2001 From: Mick Date: Sat, 25 Apr 2026 01:52:41 +0800 Subject: [PATCH] [diffusion] feat: support LoRA for LTX2.3 (#23649) --- docs_new/cookbook/diffusion/LTX/LTX.mdx | 206 ++++++++++++++++ docs_new/cookbook/diffusion/intro.mdx | 6 + docs_new/docs.json | 6 + .../src/snippets/diffusion/ltx-deployment.jsx | 233 ++++++++++++++++++ .../runtime/layers/lora/linear.py | 92 +++++-- .../runtime/pipelines/ltx_2_pipeline.py | 35 ++- .../runtime/pipelines_core/lora_pipeline.py | 45 +++- .../pipelines_core/stages/denoising_av.py | 9 + .../pipelines_core/stages/image_encoding.py | 135 +++++----- .../pipelines_core/stages/ltx_2_denoising.py | 147 +++++++++-- .../pipelines_core/stages/upsampling.py | 19 +- .../test/server/accuracy_config.py | 11 - .../test/server/accuracy_testcase_configs.py | 1 - .../multimodal_gen/test/server/gpu_cases.py | 8 - .../test/server/perf_baselines.json | 66 ----- 15 files changed, 810 insertions(+), 209 deletions(-) create mode 100644 docs_new/cookbook/diffusion/LTX/LTX.mdx create mode 100644 docs_new/src/snippets/diffusion/ltx-deployment.jsx diff --git a/docs_new/cookbook/diffusion/LTX/LTX.mdx b/docs_new/cookbook/diffusion/LTX/LTX.mdx new file mode 100644 index 000000000..967843e2e --- /dev/null +++ b/docs_new/cookbook/diffusion/LTX/LTX.mdx @@ -0,0 +1,206 @@ +--- +title: LTX +description: Run LTX-2 and LTX-2.3 video generation pipelines with SGLang Diffusion. +metatags: + description: "Deploy and use LTX-2 and LTX-2.3 video generation models with SGLang Diffusion, including one-stage, two-stage, HQ, TI2V, and LoRA examples." +--- + +import { LTXDeployment } from '/src/snippets/diffusion/ltx-deployment.jsx'; + +## 1. Model Introduction + +[LTX-2](https://huggingface.co/Lightricks/LTX-2) and [LTX-2.3](https://huggingface.co/Lightricks/LTX-2.3) are video generation models from Lightricks. SGLang Diffusion supports the LTX series through native one-stage and two-stage pipelines for text-to-video and image-conditioned video generation. + +Use `Lightricks/LTX-2` or `Lightricks/LTX-2.3` as `--model-path`. For two-stage generation, SGLang uses the spatial upsampler and distilled LoRA components from the model snapshot by default. LTX-2.3 also supports the HQ two-stage variant. + + +**License notice:** LTX-2 and LTX-2.3 are released under the LTX-2 Community License Agreement, not Apache 2.0. The license includes commercial-use restrictions for some entities. Review the [official Lightricks license](https://huggingface.co/Lightricks/LTX-2.3/blob/main/LICENSE) before production or commercial use; SGLang support does not grant additional model usage rights. + + +## 2. SGLang-diffusion Installation + +Install SGLang with diffusion dependencies: + +```bash Command +uv pip install "sglang[diffusion]" --prerelease=allow +``` + +For platform-specific setup, see the [SGLang Diffusion installation guide](/docs/sglang-diffusion/installation). + +## 3. Model Deployment + +This section provides deployment configurations optimized for different LTX pipelines and hardware targets. + +### 3.1 Basic Configuration + +The LTX series supports one-stage and two-stage pipelines. LTX-2.3 also supports the HQ two-stage pipeline. The recommended launch configuration depends on whether the target GPU can keep both two-stage DiTs resident. + +**Interactive Command Generator**: Use the configuration selector below to generate a deployment command. The default selection targets a single NVIDIA H200 with `resident` two-stage mode, which is the fastest startup path for the specified high-memory environment. + + + +### 3.2 Configuration Tips + +Choose the pipeline class based on the quality and latency target: + +| Use case | Pipeline class | Notes | +| --- | --- | --- | +| One-stage generation | `LTX2Pipeline` | Fastest LTX native path. Supports T2V and TI2V. | +| Two-stage generation | `LTX2TwoStagePipeline` | Uses a base stage and a refinement stage. Supported by LTX-2 and LTX-2.3. | +| Two-stage High Quality (HQ) generation | `LTX2TwoStageHQPipeline` | LTX-2.3 HQ path; defaults to 1920x1088 unless you override `--width` and `--height`. | + +Feature compatibility: + +| Pipeline class | T2V | TI2V (`--image-path`) | LoRA (`--lora-path`) | Notes | +| --- | --- | --- | --- | --- | +| `LTX2Pipeline` | Yes | Yes | Yes | One-stage path. Cannot be combined with HQ because HQ is a separate two-stage pipeline class. | +| `LTX2TwoStagePipeline` | Yes | Yes | Yes | Standard two-stage path for LTX-2 and LTX-2.3. | +| `LTX2TwoStageHQPipeline` | Yes | Yes | Yes | High Quality two-stage path for LTX-2.3. Use this instead of `LTX2Pipeline`; it is not a one-stage mode flag. | + +For two-stage pipelines, `--ltx2-two-stage-device-mode` controls transformer residency: + +| Mode | When to use it | +| --- | --- | +| `snapshot` | Recommended default. Balances latency and VRAM. | +| `resident` | Best latency on high-VRAM GPUs because both DiTs can stay resident. | +| `original` | Closest to the original two-stage switching semantics. | + +Other deployment flags: + +- `--lora-path`: Preload a community LoRA adapter. +- `--lora-weight-name`: Select the exact safetensors file when the LoRA repository contains multiple weight files. + + +For native LTX-2.3 two-stage serving without a user LoRA, `resident` is the fastest high-VRAM path. When you pass `--lora-path`, SGLang still applies the user LoRA during the two-stage switch, so use `resident` on H200-class GPUs for enough VRAM, but do not expect the same premerged-stage2 benefit as the no-user-LoRA path. + + +## 4. Model Invocation + +### 4.1 Basic Usage + +The examples below spell out the current SGLang sampling defaults for reproducibility: + +| Model path | Default output | Default frames | Default steps | +| --- | --- | --- | --- | +| `Lightricks/LTX-2` | 768x512 | 121 | 40 | +| `Lightricks/LTX-2.3` | 768x512 | 121 | 30 | +| `Lightricks/LTX-2.3` with `LTX2TwoStageHQPipeline` | 1920x1088 | 121 | 15 | + +#### 4.1.1 LTX-2 one-stage text-to-video + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2 \ + --pipeline-class-name LTX2Pipeline \ + --prompt "A quiet coastal town at sunrise, fishing boats moving slowly through golden mist, cinematic camera movement" \ + --save-output +``` + +#### 4.1.2 LTX-2.3 one-stage text-to-video + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2Pipeline \ + --prompt "A quiet coastal town at sunrise, fishing boats moving slowly through golden mist, cinematic camera movement" \ + --save-output +``` + +#### 4.1.3 LTX-2 two-stage text-to-video + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --prompt "A handheld shot follows a red tram crossing a rainy city square at night, reflections on the pavement, cinematic lighting" \ + --save-output +``` + +#### 4.1.4 LTX-2.3 two-stage text-to-video + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --prompt "A handheld shot follows a red tram crossing a rainy city square at night, reflections on the pavement, cinematic lighting" \ + --save-output +``` + +#### 4.1.5 LTX-2.3 HQ text-to-video + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStageHQPipeline \ + --prompt "A wide cinematic shot of alpine clouds rolling over a mountain ridge, soft morning light, slow aerial camera movement" \ + --save-output +``` + +#### 4.1.6 Image-to-video with one reference image + +Pass one image to `--image-path` for image-conditioned generation: + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --image-path ./inputs/start.png \ + --prompt "The camera slowly pushes forward as the subject turns toward warm window light, subtle natural motion, cinematic" \ + --save-output +``` + +#### 4.1.7 First-to-last-frame transition with two reference images + +Pass two images to `--image-path` for transition-style TI2V. The first image is used as the starting condition and the second image is used as the ending condition. + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --image-path ./inputs/start.png ./inputs/end.png \ + --prompt "A smooth cinematic transition from the first scene into the final scene, dynamic camera motion, motion blur, zhuanchang" \ + --save-output +``` + +### 4.2 Advanced Usage + +#### 4.2.1 Use community LoRAs + +Use `--lora-path` to load a LoRA adapter. If the Hugging Face repo contains multiple safetensors files, use `--lora-weight-name` to select the exact file. `--lora-scale` maps to the standard LoRA merge scale and defaults to `1.0`. + +The following example uses [`valiantcat/LTX-2.3-Transition-LORA`](https://huggingface.co/valiantcat/LTX-2.3-Transition-LORA): + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --lora-path valiantcat/LTX-2.3-Transition-LORA \ + --lora-weight-name ltx2.3-transition.safetensors \ + --prompt "A low-angle tracking shot moves through a foggy forest road. The camera rises above the treetops and transitions into a clear view of a snowy mountain peak under bright sunlight, zhuanchang" \ + --save-output +``` + +You can combine the Transition LoRA with two reference images: + +```bash Command +sglang generate \ + --model-path Lightricks/LTX-2.3 \ + --pipeline-class-name LTX2TwoStagePipeline \ + --image-path ./inputs/start.png ./inputs/end.png \ + --lora-path valiantcat/LTX-2.3-Transition-LORA \ + --lora-weight-name ltx2.3-transition.safetensors \ + --prompt "A fast cinematic transition from the first image to the second image, whip-pan motion, atmospheric lighting, zhuanchang" \ + --save-output +``` + + +Some community LoRAs only include weights for transformer blocks. In that case, SGLang logs a concise coverage summary and leaves unmatched LoRA-capable layers on the base model weights. This is expected when the adapter format intentionally omits those layers. + + +## 5. Practical Tips + +- Use `--pipeline-class-name LTX2TwoStagePipeline` as the default LTX two-stage quality path. +- Use `--pipeline-class-name LTX2TwoStageHQPipeline` when you want the HQ path and have enough VRAM for larger outputs. +- Use `--ltx2-two-stage-device-mode resident` on high-VRAM GPUs if latency matters more than memory usage. +- Use `--ltx2-two-stage-device-mode original` when comparing against official two-stage behavior. +- Keep `--width` and `--height` aligned with the target model resolution; for LTX models, these are output video dimensions. diff --git a/docs_new/cookbook/diffusion/intro.mdx b/docs_new/cookbook/diffusion/intro.mdx index 540a23ccb..15aa7c648 100644 --- a/docs_new/cookbook/diffusion/intro.mdx +++ b/docs_new/cookbook/diffusion/intro.mdx @@ -19,6 +19,12 @@ metatags: href="/cookbook/diffusion/Wan/Wan2.2" img="/cards/logos/wan.png" /> + { + const options = { + hardware: { + name: 'hardware', + title: 'Hardware Platform', + items: [ + { id: 'h200', label: 'H200', subtitle: 'Fastest, resident', default: true }, + { id: 'standard', label: 'Standard CUDA', subtitle: 'Snapshot mode', default: false }, + { id: 'official', label: 'Official Match', subtitle: 'Original switching', default: false }, + ], + }, + model: { + name: 'model', + title: 'Model', + items: [ + { id: 'ltx23', label: 'LTX-2.3', default: true }, + { id: 'ltx2', label: 'LTX-2', default: false }, + ], + }, + pipeline: { + name: 'pipeline', + title: 'Pipeline', + items: [ + { id: 'two-stage', label: 'Two Stage', default: true, validModels: ['ltx2', 'ltx23'] }, + { id: 'two-stage-hq', label: 'Two Stage HQ', subtitle: 'High Quality', default: false, validModels: ['ltx23'] }, + { id: 'one-stage', label: 'One Stage', default: false, validModels: ['ltx2', 'ltx23'] }, + ], + }, + }; + + const modelConfigs = { + ltx2: { + repoId: 'Lightricks/LTX-2', + pipelines: { + 'one-stage': 'LTX2Pipeline', + 'two-stage': 'LTX2TwoStagePipeline', + }, + supportedLoras: [], + }, + ltx23: { + repoId: 'Lightricks/LTX-2.3', + pipelines: { + 'one-stage': 'LTX2Pipeline', + 'two-stage': 'LTX2TwoStagePipeline', + 'two-stage-hq': 'LTX2TwoStageHQPipeline', + }, + supportedLoras: [ + { + id: 'transition', + path: 'valiantcat/LTX-2.3-Transition-LORA', + weightName: 'ltx2.3-transition.safetensors', + validPipelines: ['two-stage', 'two-stage-hq'], + }, + ], + }, + }; + + const getInitialState = () => ({ + hardware: 'h200', + model: 'ltx23', + pipeline: 'two-stage', + selectedLoraPath: 'none', + }); + + const [values, setValues] = useState(getInitialState); + const [isDark, setIsDark] = useState(false); + + useEffect(() => { + const checkDarkMode = () => { + const html = document.documentElement; + const isDarkMode = html.classList.contains('dark') || + html.getAttribute('data-theme') === 'dark' || + html.style.colorScheme === 'dark'; + setIsDark(isDarkMode); + }; + checkDarkMode(); + const observer = new MutationObserver(checkDarkMode); + observer.observe(document.documentElement, { attributes: true, attributeFilter: ['class', 'data-theme', 'style'] }); + return () => observer.disconnect(); + }, []); + + const availableLoras = (() => { + const config = modelConfigs[values.model]; + return (config?.supportedLoras || []).filter((lora) => lora.validPipelines.includes(values.pipeline)); + })(); + + const handleRadioChange = (optionName, itemId) => { + setValues((prev) => { + const next = { ...prev, [optionName]: itemId }; + + const validPipeline = options.pipeline.items.some((item) => ( + item.id === next.pipeline && item.validModels.includes(next.model) + )); + if (!validPipeline) { + next.pipeline = 'two-stage'; + } + + const config = modelConfigs[next.model]; + const nextSupported = (config?.supportedLoras || []).filter((lora) => lora.validPipelines.includes(next.pipeline)); + const isValid = nextSupported.some((lora) => lora.path === prev.selectedLoraPath); + if (!isValid) { + next.selectedLoraPath = 'none'; + } + return next; + }); + }; + + const handleLoraToggle = (path) => { + setValues((prev) => ({ + ...prev, + selectedLoraPath: prev.selectedLoraPath === path ? 'none' : path, + })); + }; + + const getDeviceMode = () => { + if (values.hardware === 'h200') { + return 'resident'; + } + if (values.hardware === 'official') { + return 'original'; + } + return 'snapshot'; + }; + + const generateCommand = () => { + const config = modelConfigs[values.model]; + const pipelineClass = config.pipelines[values.pipeline]; + if (!pipelineClass) { + return '# Error: Invalid configuration'; + } + + let command = `sglang serve \\\n --model-path ${config.repoId} \\\n --pipeline-class-name ${pipelineClass}`; + if (values.pipeline !== 'one-stage') { + command += ` \\\n --ltx2-two-stage-device-mode ${getDeviceMode()}`; + } + + const selectedLora = availableLoras.find((lora) => lora.path === values.selectedLoraPath); + if (selectedLora) { + command += ` \\\n --lora-path ${selectedLora.path} \\\n --lora-weight-name ${selectedLora.weightName}`; + } + + command += ` \\\n --port 30000`; + return command; + }; + + const containerStyle = { maxWidth: '900px', margin: '0 auto', display: 'flex', flexDirection: 'column', gap: '4px' }; + const cardStyle = { padding: '8px 12px', border: `1px solid ${isDark ? '#374151' : '#e5e7eb'}`, borderLeft: `3px solid ${isDark ? '#E85D4D' : '#D45D44'}`, borderRadius: '4px', display: 'flex', alignItems: 'center', gap: '12px', background: isDark ? '#1f2937' : '#fff' }; + const titleStyle = { fontSize: '13px', fontWeight: '600', minWidth: '140px', flexShrink: 0, color: isDark ? '#e5e7eb' : 'inherit' }; + const itemsStyle = { display: 'flex', rowGap: '2px', columnGap: '6px', flexWrap: 'wrap', alignItems: 'center', flex: 1 }; + const labelBaseStyle = { padding: '4px 10px', border: `1px solid ${isDark ? '#9ca3af' : '#d1d5db'}`, borderRadius: '3px', cursor: 'pointer', display: 'inline-flex', flexDirection: 'column', alignItems: 'center', justifyContent: 'center', fontWeight: '500', fontSize: '13px', transition: 'all 0.2s', userSelect: 'none', minWidth: '45px', textAlign: 'center', flex: 1, background: isDark ? '#374151' : '#fff', color: isDark ? '#e5e7eb' : 'inherit' }; + const checkedStyle = { background: '#D45D44', color: 'white', borderColor: '#D45D44' }; + const subtitleStyle = { display: 'block', fontSize: '9px', marginTop: '1px', lineHeight: '1.1', opacity: 0.7 }; + const commandDisplayStyle = { flex: 1, padding: '12px 16px', background: isDark ? '#111827' : '#f5f5f5', borderRadius: '6px', fontFamily: "'Menlo', 'Monaco', 'Courier New', monospace", fontSize: '12px', lineHeight: '1.5', color: isDark ? '#e5e7eb' : '#374151', whiteSpace: 'pre-wrap', overflowX: 'auto', margin: 0, border: `1px solid ${isDark ? '#374151' : '#e5e7eb'}` }; + + return ( +
+ {Object.entries(options).map(([key, option]) => { + const itemsToDisplay = key === 'pipeline' + ? option.items.filter((item) => item.validModels.includes(values.model)) + : option.items; + + return ( +
+
{option.title}
+
+ {itemsToDisplay.map((item) => { + const isChecked = values[option.name] === item.id; + return ( + + ); + })} +
+
+ ); + })} + +
+
Select LoRA Model
+
+ {availableLoras.length === 0 && ( +
+ No LoRA models available for this configuration. +
+ )} + {availableLoras.map((lora) => { + const isSelected = values.selectedLoraPath === lora.path; + return ( + + ); + })} +
+
+ +
+
Run this Command:
+ 
{generateCommand()}
+
+
+ ); +}; diff --git a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py index 961000208..ca31f8998 100644 --- a/python/sglang/multimodal_gen/runtime/layers/lora/linear.py +++ b/python/sglang/multimodal_gen/runtime/layers/lora/linear.py @@ -2,7 +2,7 @@ # SPDX-License-Identifier: Apache-2.0 # Code adapted from SGLang https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/lora/layers.py - +import os import torch from torch import nn @@ -37,10 +37,17 @@ torch._dynamo.config.recompile_limit = 64 LORA_MERGE_CHUNK_BYTES = 32 * 1024 * 1024 +LoRAWeightEntry = tuple[ + torch.nn.Parameter, + torch.nn.Parameter, + str | None, + float, + int | None, + int | None, +] class BaseLayerWithLoRA(nn.Module): - def __init__( self, base_layer: nn.Module, @@ -60,9 +67,7 @@ class BaseLayerWithLoRA(nn.Module): self.disable_lora: bool = True self.lora_rank = lora_rank self.lora_alpha = lora_alpha - self.lora_weights_list: list[ - tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] - ] = [] + self.lora_weights_list: list[LoRAWeightEntry] = [] self.lora_path: str | None = None self.strength: float = 1.0 @@ -147,7 +152,16 @@ class BaseLayerWithLoRA(nn.Module): self.strength = 1.0 # Add to list for multi-LoRA support - self.lora_weights_list.append((lora_A_param, lora_B_param, lora_path, strength)) + self.lora_weights_list.append( + ( + lora_A_param, + lora_B_param, + lora_path, + strength, + self.lora_rank, + self.lora_alpha, + ) + ) # Set backward compatibility attributes to point to the last LoRA (for single LoRA case) # This ensures backward compatibility while supporting multiple LoRA @@ -166,29 +180,27 @@ class BaseLayerWithLoRA(nn.Module): def _merge_lora_into_data( self, data: torch.Tensor, - lora_list: list[ - tuple[torch.nn.Parameter, torch.nn.Parameter, str | None, float] - ], + lora_list: list[LoRAWeightEntry], ) -> None: """ Merge all LoRA adapters into the data tensor in-place. Args: data: The base weight tensor to merge LoRA into (modified in-place) - lora_list: List of (lora_A, lora_B, lora_path, lora_strength) tuples + lora_list: List of (lora_A, lora_B, lora_path, lora_strength, rank, alpha) tuples """ # Merge all LoRA adapters in order - for lora_A, lora_B, _, lora_strength in lora_list: + for lora_A, lora_B, _, lora_strength, lora_rank, lora_alpha in lora_list: 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 + lora_alpha is not None + and lora_rank is not None + and lora_alpha != lora_rank ): - scale *= self.lora_alpha / self.lora_rank + scale *= lora_alpha / lora_rank if not isinstance(lora_B_sliced, torch.Tensor): lora_delta = lora_B_sliced @ lora_A_sliced @@ -222,6 +234,17 @@ class BaseLayerWithLoRA(nn.Module): chunk_delta = lora_B_2d[start:end] @ lora_A_sliced data_2d[start:end].add_(chunk_delta, alpha=scale) + def _should_merge_in_fp32( + self, + lora_list: list[LoRAWeightEntry], + ) -> bool: + if os.getenv("SGLANG_DIFFUSION_LORA_MERGE_FP32", "0") != "1": + return False + for _, _, lora_path, _, _, _ in lora_list: + if lora_path and "distilled-lora" in lora_path.lower(): + return False + return True + @torch.no_grad() def merge_lora_weights(self, strength: float | None = None) -> None: if strength is not None: @@ -236,11 +259,22 @@ class BaseLayerWithLoRA(nn.Module): # Use lora_weights_list if available, otherwise fall back to single LoRA for backward compatibility lora_list = self.lora_weights_list if self.lora_weights_list else [] if not lora_list and self.lora_A is not None and self.lora_B is not None: - lora_list = [(self.lora_A, self.lora_B, self.lora_path, self.strength)] + lora_list = [ + ( + self.lora_A, + self.lora_B, + self.lora_path, + self.strength, + self.lora_rank, + self.lora_alpha, + ) + ] if not lora_list: raise ValueError("LoRA weights not set. Please set them first.") + merge_in_fp32 = self._should_merge_in_fp32(lora_list) + if isinstance(self.base_layer.weight, DTensor): mesh = self.base_layer.weight.data.device_mesh unsharded_base_layer = ReplicatedLinear( @@ -257,10 +291,19 @@ class BaseLayerWithLoRA(nn.Module): data = self.base_layer.weight.data.to( get_local_torch_device() ).full_tensor() + target_dtype = data.dtype + if ( + merge_in_fp32 + and data.is_floating_point() + and data.dtype != torch.float32 + ): + data = data.to(torch.float32) self._merge_lora_into_data(data, lora_list) - unsharded_base_layer.weight = nn.Parameter(data.to(current_device)) + unsharded_base_layer.weight = nn.Parameter( + data.to(current_device, dtype=target_dtype) + ) if isinstance(getattr(self.base_layer, "bias", None), DTensor): unsharded_base_layer.bias = nn.Parameter( self.base_layer.bias.to(get_local_torch_device(), non_blocking=True) @@ -282,10 +325,19 @@ class BaseLayerWithLoRA(nn.Module): else: current_device = self.base_layer.weight.data.device data = self.base_layer.weight.data.to(get_local_torch_device()) + target_dtype = data.dtype + if ( + merge_in_fp32 + and data.is_floating_point() + and data.dtype != torch.float32 + ): + data = data.to(torch.float32) self._merge_lora_into_data(data, lora_list) - self.base_layer.weight.data = data.to(current_device, non_blocking=True) + self.base_layer.weight.data = data.to( + current_device, dtype=target_dtype, non_blocking=True + ) self.merged = True @@ -342,7 +394,6 @@ class VocabParallelEmbeddingWithLoRA(BaseLayerWithLoRA): class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): - def __init__( self, base_layer: ColumnParallelLinear, @@ -400,7 +451,6 @@ class ColumnParallelLinearWithLoRA(BaseLayerWithLoRA): class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): - def __init__( self, base_layer: MergedColumnParallelLinear, @@ -422,7 +472,6 @@ class MergedColumnParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): - def __init__( self, base_layer: QKVParallelLinear, @@ -455,7 +504,6 @@ class QKVParallelLinearWithLoRA(ColumnParallelLinearWithLoRA): class RowParallelLinearWithLoRA(BaseLayerWithLoRA): - def __init__( self, base_layer: RowParallelLinear, 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 c1cc5c02e..ad6f842a6 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines/ltx_2_pipeline.py @@ -494,19 +494,25 @@ class LTX2TwoStageDeviceManager: if "stage2" in self._phase_ready_events: return if self._snapshot_low_vram_mode: - stage1_module = self.pipeline.get_module("transformer") - stage1_param = ( - next(stage1_module.parameters(), None) - if stage1_module is not None - else None - ) - if stage1_param is not None and stage1_param.device.type == "cuda": - self._release_module_to_cpu_snapshot("transformer") + self._release_stage1_for_low_vram() self._schedule_phase_prefetch( "stage2", self.pipeline.get_module("transformer_2") ) + def prepare_upsample_after_stage1(self) -> bool: + if ( + not self.should_use_premerged + or self.mode != "snapshot" + or not self.server_args.dit_cpu_offload + or not self._snapshot_low_vram_mode + ): + return False + if "stage2" in self._phase_ready_events: + return False + self._release_stage1_for_low_vram() + return True + def ensure_phase_ready(self, phase: str | None) -> None: if not self.should_use_premerged or phase not in ("stage1", "stage2"): return @@ -616,6 +622,16 @@ class LTX2TwoStageDeviceManager: phase = "stage2" if module_name == "transformer_2" else "stage1" self._phase_ready_events.pop(phase, None) + def _release_stage1_for_low_vram(self) -> None: + stage1_module = self.pipeline.get_module("transformer") + stage1_param = ( + next(stage1_module.parameters(), None) + if stage1_module is not None + else None + ) + if stage1_param is not None and stage1_param.device.type == "cuda": + self._release_module_to_cpu_snapshot("transformer") + def _ensure_on_gpu(self, module_name: str) -> None: module = self.pipeline.get_module(module_name) if module is None: @@ -807,6 +823,9 @@ class LTX2TwoStagePipeline(_BaseLTX2Pipeline): def prefetch_ltx2_stage2_after_stage1(self) -> None: self._device_manager.prefetch_stage2_after_stage1() + def prepare_ltx2_upsample_after_stage1(self) -> bool: + return self._device_manager.prepare_upsample_after_stage1() + def should_skip_ltx2_lora_switch_stage(self) -> bool: return self._use_premerged_stage2_transformer and self._device_manager.mode in ( "snapshot", diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py index ca4cafe1b..fb93a3f01 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/lora_pipeline.py @@ -99,7 +99,9 @@ class LoRAPipeline(ComposedPipelineBase): if self.lora_path is not None: self.convert_to_lora_layers() self.set_lora( - self.lora_nickname, self.lora_path, strength=self.server_args.lora_scale # type: ignore + self.lora_nickname, + self.lora_path, + strength=self.server_args.lora_scale, # type: ignore ) # type: ignore def is_target_layer(self, module_name: str) -> bool: @@ -426,6 +428,8 @@ class LoRAPipeline(ComposedPipelineBase): ) adapted_count = 0 + missing_layers_by_adapter = [[] for _ in lora_nicknames] + applied_count_by_adapter = [0 for _ in lora_nicknames] for name, layer in lora_layers.items(): # Apply all LoRA adapters in order for idx, (nickname, path, lora_strength) in enumerate( @@ -465,13 +469,9 @@ class LoRAPipeline(ComposedPipelineBase): ), # Only clear on first LoRA ) adapted_count += 1 + applied_count_by_adapter[idx] += 1 else: - if rank == 0 and idx == 0: # Only warn for first missing LoRA - logger.warning( - "LoRA adapter %s does not contain the weights for layer '%s'. LoRA will not be applied to it.", - path, - name, - ) + missing_layers_by_adapter[idx].append(name) # Only disable if no LoRA was applied at all if idx == len(lora_nicknames) - 1: has_any_lora = any( @@ -481,6 +481,37 @@ class LoRAPipeline(ComposedPipelineBase): ) if not has_any_lora: layer.disable_lora = True + + if rank == 0: + total_layers = len(lora_layers) + example_limit = 8 + for idx, path in enumerate(lora_paths): + missing_layers = missing_layers_by_adapter[idx] + if not missing_layers: + continue + missing_count = len(missing_layers) + applied_count = applied_count_by_adapter[idx] + examples = ", ".join(missing_layers[:example_limit]) + if missing_count > example_limit: + examples += ", ..." + if applied_count == 0: + logger.warning( + "LoRA adapter %s did not match any LoRA layer. " + "Checked %d layers; examples: %s", + path, + total_layers, + examples, + ) + else: + logger.info( + "LoRA adapter %s covers %d/%d LoRA layers; " + "%d layers use base weights. Examples: %s", + path, + applied_count, + total_layers, + missing_count, + examples, + ) return adapted_count def is_lora_effective(self, target: str = "all") -> bool: 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 59d95021e..c8c488c3f 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 @@ -217,6 +217,14 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): 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" + pipeline = self.pipeline() if self.pipeline else None + ensure_phase_ready = ( + getattr(pipeline, "ensure_ltx2_phase_ready", None) + if pipeline is not None + else None + ) + if callable(ensure_phase_ready): + ensure_phase_ready("stage2") original_clean_latent_background = getattr( batch, "ltx2_ti2v_clean_latent_background", None ) @@ -250,6 +258,7 @@ class LTX2RefinementStage(LTX2AVDenoisingStage): renoise_generator = None if is_native_ti2v: prepared_latents, denoise_mask, _ = self._prepare_ltx2_ti2v_clean_state( + batch=batch, latents=batch.latents, image_latent=batch.image_latent, num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), diff --git a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py index 13a9e779f..cb232ff33 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/image_encoding.py @@ -417,25 +417,48 @@ class LTX2ImageEncodingStage(PipelineStage): ): return self._condition_image_encoder(video_condition) + @staticmethod + def _normalize_ltx2_image_paths(image_path: str | list[str]) -> list[str]: + image_paths = image_path if isinstance(image_path, list) else [image_path] + if len(image_paths) > 2: + raise ValueError( + "LTX-2 TI2V currently supports at most two conditioning images " + "([first_frame, last_frame])." + ) + return image_paths + + @staticmethod + def _normalize_ltx2_image_latents( + image_latent: torch.Tensor | list[torch.Tensor] | None, + ) -> list[torch.Tensor]: + if image_latent is None: + return [] + return image_latent if isinstance(image_latent, list) else [image_latent] + # -- forward --------------------------------------------------------- @torch.no_grad() def forward(self, batch: Req, server_args: ServerArgs) -> Req: if batch.image_path is None: return batch + image_paths = self._normalize_ltx2_image_paths(batch.image_path) + + vae_sf = int(server_args.pipeline_config.vae_scale_factor) + patch = int(server_args.pipeline_config.patch_size) + expected_tokens = (int(batch.height) // vae_sf // patch) * ( + int(batch.width) // vae_sf // patch + ) if ( batch.image_latent is not None and int(getattr(batch, "ltx2_num_image_tokens", 0)) > 0 ): # Re-encode if resolution changed (e.g. two-stage upsample between stages) - vae_sf = int(server_args.pipeline_config.vae_scale_factor) - patch = int(server_args.pipeline_config.patch_size) - expected = (int(batch.height) // vae_sf // patch) * ( - int(batch.width) // vae_sf // patch - ) - if int(batch.image_latent.shape[1]) == expected: + existing_latents = self._normalize_ltx2_image_latents(batch.image_latent) + if len(existing_latents) == len(image_paths) and all( + int(latent.shape[1]) == expected_tokens for latent in existing_latents + ): return batch - # Resolution mismatch — clear and re-encode below + # Resolution or reference-count mismatch — clear and re-encode below batch.image_latent = None batch.ltx2_num_image_tokens = 0 @@ -447,20 +470,23 @@ class LTX2ImageEncodingStage(PipelineStage): from sglang.multimodal_gen.runtime.models.vision_utils import load_image - # 1. Load image, apply codec compression, resize for condition_image - image_path = ( - batch.image_path[0] - if isinstance(batch.image_path, list) - else batch.image_path - ) - img = load_image(image_path) - arr = np.array(img).astype(np.uint8)[..., :3] - arr = self._apply_video_codec_compression(arr, crf=33) - conditioned_img = PIL.Image.fromarray(arr) - batch.condition_image = conditioned_img.resize( - (int(batch.width), int(batch.height)), - resample=PIL.Image.Resampling.BILINEAR, - ) + # 1. Load images, apply codec compression, resize for condition_image + conditioned_imgs = [] + for image_path in image_paths: + img = load_image(image_path) + arr = np.array(img).astype(np.uint8)[..., :3] + arr = self._apply_video_codec_compression(arr, crf=33) + conditioned_img = PIL.Image.fromarray(arr) + conditioned_imgs.append(conditioned_img) + batch.condition_image = [ + img.resize( + (int(batch.width), int(batch.height)), + resample=PIL.Image.Resampling.BILINEAR, + ) + for img in conditioned_imgs + ] + if len(batch.condition_image) == 1: + batch.condition_image = batch.condition_image[0] # 2. Load encoder(s) to device, cast to encode_dtype use_condition_encoder = self._ensure_condition_image_encoder(server_args) @@ -478,21 +504,35 @@ class LTX2ImageEncodingStage(PipelineStage): else: self.vae = self.vae.to(dtype=encode_dtype) - video_condition = self._pil_to_video_tensor( - conditioned_img, - width=int(batch.width), - height=int(batch.height), - device=device, - dtype=encode_dtype, - ) - - # 3. Encode - if use_condition_encoder: - latent = self._condition_encode(video_condition, server_args).to( - dtype=encode_dtype + packed_latents = [] + for conditioned_img in conditioned_imgs: + video_condition = self._pil_to_video_tensor( + conditioned_img, + width=int(batch.width), + height=int(batch.height), + device=device, + dtype=encode_dtype, ) - else: - latent = self._vae_encode(video_condition, server_args, batch.generator) + + # 3. Encode + if use_condition_encoder: + latent = self._condition_encode(video_condition, server_args).to( + dtype=encode_dtype + ) + else: + latent = self._vae_encode(video_condition, server_args, batch.generator) + + packed = server_args.pipeline_config.maybe_pack_latents( + latent, latent.shape[0], batch + ) + if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): + raise ValueError("Expected packed image latents [B, S0, D].") + if int(packed.shape[1]) != expected_tokens: + raise ValueError( + f"LTX-2 conditioning token count mismatch: " + f"{packed.shape[1]=} {expected_tokens=}." + ) + packed_latents.append(packed) # Restore VAE to its config dtype (shared with decoding stage) if not use_condition_encoder: @@ -501,32 +541,16 @@ class LTX2ImageEncodingStage(PipelineStage): ] self.vae = self.vae.to(dtype=original_dtype) - # 4. Pack into token latents and validate - packed = server_args.pipeline_config.maybe_pack_latents( - latent, latent.shape[0], batch + batch.image_latent = ( + packed_latents[0] if len(packed_latents) == 1 else packed_latents ) - if not (isinstance(packed, torch.Tensor) and packed.ndim == 3): - raise ValueError("Expected packed image latents [B, S0, D].") - - vae_sf = int(server_args.pipeline_config.vae_scale_factor) - patch = int(server_args.pipeline_config.patch_size) - expected_tokens = (int(batch.height) // vae_sf // patch) * ( - int(batch.width) // vae_sf // patch - ) - if int(packed.shape[1]) != expected_tokens: - raise ValueError( - f"LTX-2 conditioning token count mismatch: " - f"{packed.shape[1]=} {expected_tokens=}." - ) - - batch.image_latent = packed - batch.ltx2_num_image_tokens = int(packed.shape[1]) + batch.ltx2_num_image_tokens = int(packed_latents[0].shape[1]) if batch.debug: logger.info( "LTX2 TI2V: %d tokens (shape=%s) for %sx%s", batch.ltx2_num_image_tokens, - tuple(batch.image_latent.shape), + tuple(packed_latents[0].shape), batch.width, batch.height, ) @@ -703,7 +727,6 @@ class ImageVAEEncodingStage(PipelineStage): self, image: torch.Tensor | PIL.Image.Image, ) -> torch.Tensor: - if isinstance(image, PIL.Image.Image): image = pil_to_numpy(image) # to np image = numpy_to_pt(image) # to pt 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 492ba154f..9f1e68265 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 @@ -516,24 +516,111 @@ class LTX2DenoisingStage(DenoisingStage): return next_video, next_audio @staticmethod - def _prepare_ltx2_ti2v_clean_state( + def _normalize_ltx2_condition_latents( + image_latent: torch.Tensor | list[torch.Tensor] | None, + ) -> list[torch.Tensor]: + if image_latent is None: + return [] + return image_latent if isinstance(image_latent, list) else [image_latent] + + @classmethod + def _get_ltx2_condition_spans( + cls, + batch: Req, latents: torch.Tensor, - image_latent: torch.Tensor, + image_latent: torch.Tensor | list[torch.Tensor] | None, + num_img_tokens: int, + ) -> list[tuple[int, torch.Tensor]]: + if num_img_tokens <= 0: + return [] + if not (isinstance(latents, torch.Tensor) and latents.ndim == 3): + raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") + + condition_latents = cls._normalize_ltx2_condition_latents(image_latent) + if not condition_latents: + return [] + if len(condition_latents) > 2: + raise ValueError( + "LTX-2 TI2V currently supports at most two conditioning images." + ) + + for cond in condition_latents: + if not (isinstance(cond, torch.Tensor) and cond.ndim == 3): + raise ValueError( + "Expected LTX-2 conditioning latents to be packed tensors [B, S, D]." + ) + if int(cond.shape[1]) < int(num_img_tokens): + raise ValueError( + "LTX-2 conditioning latent is shorter than one frame token span." + ) + + did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) + if not did_sp_shard: + if int(latents.shape[1]) < int(num_img_tokens): + raise ValueError( + "LTX-2 latent sequence is shorter than one conditioning frame." + ) + if len(condition_latents) == 1: + return [(0, condition_latents[0])] + return [ + (0, condition_latents[0]), + (int(latents.shape[1]) - int(num_img_tokens), condition_latents[1]), + ] + + tokens_per_frame = int(getattr(batch, "sp_video_tokens_per_frame", 0)) + if tokens_per_frame <= 0: + raise ValueError( + "SP-sharded LTX-2 TI2V requires batch.sp_video_tokens_per_frame." + ) + if int(num_img_tokens) != int(tokens_per_frame): + raise ValueError( + "LTX-2 conditioning token count must match one latent frame when using SP." + ) + + raw_shape = getattr(batch, "raw_latent_shape", None) + if raw_shape is None: + raise ValueError("SP-sharded LTX-2 TI2V requires batch.raw_latent_shape.") + global_seq_len = int(raw_shape[1]) + if global_seq_len % tokens_per_frame != 0: + raise ValueError( + "SP-sharded LTX-2 TI2V expected raw seq_len divisible by tokens_per_frame." + ) + + global_num_frames = global_seq_len // tokens_per_frame + local_start_frame = int(getattr(batch, "sp_video_start_frame", 0)) + local_num_frames = int(getattr(batch, "sp_video_latent_num_frames", 0)) + local_end_frame = local_start_frame + local_num_frames + + spans: list[tuple[int, torch.Tensor]] = [] + if local_start_frame == 0: + spans.append((0, condition_latents[0])) + + if len(condition_latents) == 2: + last_global_frame = global_num_frames - 1 + if local_start_frame <= last_global_frame < local_end_frame: + local_last_frame = last_global_frame - local_start_frame + spans.append( + (local_last_frame * tokens_per_frame, condition_latents[1]) + ) + + return spans + + @classmethod + def _prepare_ltx2_ti2v_clean_state( + cls, + batch: Req, + latents: torch.Tensor, + image_latent: torch.Tensor | list[torch.Tensor] | None, num_img_tokens: int, zero_clean_latent: bool, clean_latent_background: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: latents = latents.clone() - conditioned = image_latent[:, :num_img_tokens, :].to( - device=latents.device, dtype=latents.dtype - ) - latents[:, :num_img_tokens, :] = conditioned denoise_mask = torch.ones( (latents.shape[0], latents.shape[1], 1), device=latents.device, dtype=torch.float32, ) - denoise_mask[:, :num_img_tokens, :] = 0.0 if clean_latent_background is not None: clean_latent = ( clean_latent_background.detach() @@ -544,7 +631,24 @@ class LTX2DenoisingStage(DenoisingStage): clean_latent = torch.zeros_like(latents) else: clean_latent = latents.detach().clone() - clean_latent[:, :num_img_tokens, :] = conditioned + + spans = cls._get_ltx2_condition_spans( + batch=batch, + latents=latents, + image_latent=image_latent, + num_img_tokens=num_img_tokens, + ) + for start, cond in spans: + stop = int(start) + int(num_img_tokens) + conditioned = cls._repeat_batch_dim( + cond[:, :num_img_tokens, :].to( + device=latents.device, dtype=latents.dtype + ), + int(latents.shape[0]), + ) + latents[:, start:stop, :] = conditioned + denoise_mask[:, start:stop, :] = 0.0 + clean_latent[:, start:stop, :] = conditioned return latents, denoise_mask, clean_latent @staticmethod @@ -915,23 +1019,6 @@ class LTX2DenoisingStage(DenoisingStage): super()._preprocess_sp_latents(batch, server_args) batch.image_latent = saved - @staticmethod - def _should_apply_ltx2_ti2v(batch: Req) -> bool: - """True if we have an image-latent token prefix to condition with. - - SP note: when token latents are time-sharded, only the rank that owns the - *global* first latent frame should apply TI2V conditioning (rank with start_frame==0). - """ - if ( - batch.image_latent is None - or int(getattr(batch, "ltx2_num_image_tokens", 0)) <= 0 - ): - return False - did_sp_shard = bool(getattr(batch, "did_sp_shard_latents", False)) - if not did_sp_shard: - 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" @@ -998,8 +1085,6 @@ class LTX2DenoisingStage(DenoisingStage): # Video and audio keep separate scheduler state throughout the denoising loop. ctx.audio_scheduler = copy.deepcopy(self.scheduler) - do_ti2v = self._should_apply_ltx2_ti2v(batch) - if ctx.use_ltx23_legacy_one_stage: batch.ltx23_audio_replicated_for_sp = False batch.did_sp_shard_audio_latents = False @@ -1038,6 +1123,13 @@ class LTX2DenoisingStage(DenoisingStage): batch.width // server_args.pipeline_config.vae_config.arch_config.spatial_compression_ratio ) + ti2v_spans = self._get_ltx2_condition_spans( + batch=batch, + latents=ctx.latents, + image_latent=batch.image_latent, + num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), + ) + do_ti2v = bool(ti2v_spans) if do_ti2v: if not (isinstance(ctx.latents, torch.Tensor) and ctx.latents.ndim == 3): raise ValueError("LTX-2 TI2V expects packed token latents [B, S, D].") @@ -1052,6 +1144,7 @@ class LTX2DenoisingStage(DenoisingStage): # Keep conditioned tokens clean and reuse the mask during every step update. ctx.latents, ctx.denoise_mask, ctx.clean_latent = ( self._prepare_ltx2_ti2v_clean_state( + batch=batch, latents=ctx.latents, image_latent=batch.image_latent, num_img_tokens=int(getattr(batch, "ltx2_num_image_tokens", 0)), 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 b2d7d20f3..d8f5fb24c 100644 --- a/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py +++ b/python/sglang/multimodal_gen/runtime/pipelines_core/stages/upsampling.py @@ -84,8 +84,8 @@ class LTX2UpsampleStage(PipelineStage): device=device, dtype=latents.dtype ) latents = self.spatial_upsampler(latents) - if server_args.vae_cpu_offload: - self.spatial_upsampler = self.spatial_upsampler.to("cpu") + # Keep the small spatial upsampler resident after warmup; moving it + # every request dominates the measured two-stage upsample latency. latents = (latents - vae_mean) / vae_std return latents @@ -118,18 +118,31 @@ class LTX2UpsampleStage(PipelineStage): ) def forward(self, batch: Req, server_args: ServerArgs) -> Req: + delay_stage2_prefetch = False + if self.pipeline is not None: + prepare_upsample = getattr( + self.pipeline, "prepare_ltx2_upsample_after_stage1", None + ) + if callable(prepare_upsample): + delay_stage2_prefetch = prepare_upsample() prefetch_stage2 = ( getattr(self.pipeline, "prefetch_ltx2_stage2_after_stage1", None) if self.pipeline is not None else None ) - if callable(prefetch_stage2): + if callable(prefetch_stage2) and not delay_stage2_prefetch: prefetch_stage2() device = get_local_torch_device() latents = self._upsample_video_latents(batch.latents, server_args, device) + if callable(prefetch_stage2) and delay_stage2_prefetch: + prefetch_stage2() logger.info("Upsampled video latents: %s", list(latents.shape)) self._restore_full_resolution(batch) + batch.image_latent = None + batch.ltx2_num_image_tokens = 0 + batch.did_sp_shard_latents = False + batch.did_sp_shard_audio_latents = False self._pack_video_latents(batch, latents, server_args) logger.info( "Packed video latents for Stage 2: %s (resolution %dx%d)", diff --git a/python/sglang/multimodal_gen/test/server/accuracy_config.py b/python/sglang/multimodal_gen/test/server/accuracy_config.py index 469f58644..955c84ad1 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_config.py +++ b/python/sglang/multimodal_gen/test/server/accuracy_config.py @@ -142,17 +142,6 @@ SKIP_COMPONENTS: Dict[str, Dict[ComponentType, ComponentSkip]] = { "Representative text encoder accuracy is already covered by flux_2_image_t2i for the same source component and topology" ), }, - "wan2_1_t2v_1.3b_text_encoder_cpu_offload": { - ComponentType.VAE: ComponentSkip( - "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology" - ), - ComponentType.TRANSFORMER: ComponentSkip( - "Representative transformer accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology" - ), - ComponentType.TEXT_ENCODER: ComponentSkip( - "Representative text encoder accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology" - ), - }, "wan2_1_t2v_1.3b_teacache_enabled": { ComponentType.VAE: ComponentSkip( "Representative VAE accuracy is already covered by wan2_1_t2v_1.3b for the same source component and topology" diff --git a/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py b/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py index fedec989d..83268b989 100644 --- a/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py +++ b/python/sglang/multimodal_gen/test/server/accuracy_testcase_configs.py @@ -31,7 +31,6 @@ ACCURACY_ONE_GPU_CASE_IDS = ( "flux_2_image_t2i_upscaling_4x", "mova_360p_1gpu", "wan2_1_t2v_1.3b", - "wan2_1_t2v_1.3b_text_encoder_cpu_offload", "wan2_1_t2v_1.3b_teacache_enabled", "wan2_1_t2v_1.3b_frame_interp_2x", "wan2_1_t2v_1.3b_upscaling_4x", diff --git a/python/sglang/multimodal_gen/test/server/gpu_cases.py b/python/sglang/multimodal_gen/test/server/gpu_cases.py index 637800b60..31edfa1ed 100644 --- a/python/sglang/multimodal_gen/test/server/gpu_cases.py +++ b/python/sglang/multimodal_gen/test/server/gpu_cases.py @@ -175,14 +175,6 @@ ONE_GPU_CASES: list[DiffusionTestCase] = [ ), T2V_sampling_params, ), - DiffusionTestCase( - "wan2_1_t2v_1.3b_text_encoder_cpu_offload", - DiffusionServerArgs( - model_path=DEFAULT_WAN_2_1_T2V_1_3B_MODEL_NAME_FOR_TEST, - text_encoder_cpu_offload=True, - ), - T2V_sampling_params, - ), # TeaCache acceleration test for Wan video model DiffusionTestCase( "wan2_1_t2v_1.3b_teacache_enabled", diff --git a/python/sglang/multimodal_gen/test/server/perf_baselines.json b/python/sglang/multimodal_gen/test/server/perf_baselines.json index 51c5a4b4b..12c74a092 100644 --- a/python/sglang/multimodal_gen/test/server/perf_baselines.json +++ b/python/sglang/multimodal_gen/test/server/perf_baselines.json @@ -951,72 +951,6 @@ "expected_median_denoise_ms": 145.57, "estimated_full_test_time_s": 129.3 }, - "wan2_1_t2v_1.3b_text_encoder_cpu_offload": { - "stages_ms": { - "DecodingStage": 675.91, - "TextEncodingStage": 1072.93, - "TimestepPreparationStage": 2.43, - "LatentPreparationStage": 0.14, - "InputValidationStage": 0.07, - "DenoisingStage": 7221.86 - }, - "denoise_step_ms": { - "0": 128.64, - "1": 108.46, - "2": 140.93, - "3": 140.8, - "4": 141.21, - "5": 141.84, - "6": 141.56, - "7": 142.35, - "8": 142.08, - "9": 141.68, - "10": 141.69, - "11": 141.46, - "12": 141.39, - "13": 141.36, - "14": 141.65, - "15": 141.55, - "16": 142.02, - "17": 141.53, - "18": 140.98, - "19": 142.4, - "20": 141.84, - "21": 141.3, - "22": 141.41, - "23": 141.47, - "24": 141.78, - "25": 141.6, - "26": 142.19, - "27": 141.09, - "28": 141.2, - "29": 141.22, - "30": 141.2, - "31": 141.23, - "32": 141.41, - "33": 141.5, - "34": 141.56, - "35": 141.51, - "36": 141.25, - "37": 141.49, - "38": 141.56, - "39": 141.52, - "40": 141.17, - "41": 141.83, - "42": 141.72, - "43": 142.31, - "44": 141.56, - "45": 141.91, - "46": 141.93, - "47": 141.9, - "48": 141.52, - "49": 141.34 - }, - "expected_e2e_ms": 9296.17, - "expected_avg_denoise_ms": 144.33, - "expected_median_denoise_ms": 144.84, - "estimated_full_test_time_s": 129.3 - }, "wan2_1_t2v_1.3b_cfg_parallel": { "stages_ms": { "LatentPreparationStage": 0.29,